用于递归连接的 TensorFlow 高效共享内存分配

Posted 2023-02-23

【中文标题】用于递归连接的 TensorFlow 高效共享内存分配

【英文标题】：TensorFlow efficient shared memory allocation for recursive concatenation

【发布时间】：2017-09-08 21:50:19

【问题描述】：

DenseNets 在 TensorFlow 中往往会占用大量内存，因为每个 concat 操作都存储在单独的分配中。最近的一篇论文Memory-Efficient Implementation of DenseNets 表明，通过共享分配可以显着降低内存利用率。这张来自论文 + pytorch 实现的图片说明了共享内存方法：

这如何用 TensorFlow 实现？如果不能通过python实现，那么如何在支持CPU和GPU的Op中正确实现呢？

Pytorch efficient DenseNet implementation Keras DenseNet Implementation 具有“幼稚”分配，适用于 TensorFlow 后端。

我创建了一个TensorFlow Feature Request for necessary allocation functionality。

【问题讨论】：

这是一个强大的想法，如果它不能在 Python 中完成，我不知道为什么，因为 Python 中的所有其他东西绝对容易受到这种技术的影响。

但是，如果 Tensorflow 所需的部分是用另一种语言编写的，共享内存的处理方式不同或特殊安排，或者出现其他问题，你最终将不得不进行大手术，并且替换它时可能会出现很多级联问题和错误。

【参考方案1】：

现在可以在以下位置获得内存高效实现：

https://github.com/joeyearsley/efficient_densenet_tensorflow

上面链接的相关函数是：

# Gradient checkpoint the layer
_x = tf.contrib.layers.recompute_grad(_x)