Tensorflow：将张量从 GPU 移动到 CPU 以释放内存

Question

我正在尝试进行一些手动内存管理，但不确定操作完成后如何将张量移至 CPU 上。对于我的用例，我想在 GPU 上运行 STFT，然后立即将张量移到 CPU 上，因为它占用大量内存。就像是。

with tf.device("GPU:0"):
    foo = tf.signal.stft(bar, ...)

with tf.device("CPU:0"):
    baz = foo
del foo

Answer 1

一般情况下，作用域会调整操作的行为，但纯赋值不会生成 TF 操作。像这样的事情会起作用吗：

with tf.device("CPU:0"):
  baz = tf.identity(foo)