tf.identity是否在指定设备上不传输或保持张量？

Question

我正在尝试使用Tensorflow构建vDNN模型。这里的主要思想是将特征映射换出到主机内存，并在计算梯度时将其取回。

我使用私有函数Operation._update_input将生成的计算图修改为：

def remap_gradients(self):
    for t in self._transfer_ts:
        t_name = t.name.rsplit('/', 1)[0]
        with tf.device('/cpu:0'):
            identity = tf.identity(t, name='%s/transfer_to_cpu' % t_name)

        for op in t.consumers():
            for i, tensor in enumerate(op.inputs):
                if tensor == t and op != identity.op:
                    op._update_input(i, identity, False)

其中self._transfer_ts是我要转移到主机内存的张量。重新映射后，我得到了新的结构：

所有梯度计算都将重定向到identity操作，该操作应在主机内存上运行。

但是在GPU上运行此模型时，我仍然会用光内存。似乎identity操作无法正常运行，或者它会将功能图传回并再次将其缓存在GPU上。那么，这里出了什么问题？

我试图分配给在主机内存上分配的变量。但是有一个问题：变量在分配形状时需要指定形状。但是，某些小批量中的批量大小发生了变化，这将在转移时导致形状不匹配。

EDIT1：

我试图将tf.identity与下一层断开连接，这意味着标识值仅用于梯度计算。结果似乎是传输了数据，但比正向计算要慢。它仍然没有足够的内存，但是我知道了：

Current usage from device: /job:localhost/replica:0/task:0/device:GPU:0, allocator: GPU_0_bfc
  196.00MiB from block2_conv1/Conv2D

但是如果我禁用传递功能，则会得到：

Current usage from device: /job:localhost/replica:0/task:0/device:GPU:0, allocator: GPU_0_bfc
  392.00MiB from block1_conv1/Conv2D
  392.00MiB from block1_conv2/Conv2D
  196.00MiB from block2_conv1/Conv2D
  98.00MiB from block1_pool/MaxPool

似乎释放了一些张量，但是内存没有被重用。