使用 tf.GradientTape() 会耗尽所有的 gpu 内存，没有它也没关系

Question

我正在研究 Convolution Tasnet，我制作的模型大小约为 505 万个变量。

我想使用自定义训练循环来训练它，问题是，

for i, (input_batch, target_batch) in enumerate(train_ds): # each shape is (64, 32000, 1)
    with tf.GradientTape() as tape:
        predicted_batch = cv_tasnet(input_batch, training=True) # model name
        loss = calculate_sisnr(predicted_batch, target_batch) # some custom loss
    trainable_vars = cv_tasnet.trainable_variables
    gradients = tape.gradient(loss, trainable_vars)
    cv_tasnet.optimizer.apply_gradients(zip(gradients, trainable_vars))

这部分耗尽了所有的GPU内存（可用24GB）..
当我尝试不使用 tf.GradientTape() as tape、

for i, (input_batch, target_batch) in enumerate(train_ds):
        predicted_batch = cv_tasnet(input_batch, training=True)
        loss = calculate_sisnr(predicted_batch, target_batch)

这使用了合理数量的 gpu 内存（大约 5~6GB）。

我对基本的 mnist 数据尝试了 tf.GradientTape() as tape 的相同格式，然后它可以正常工作。
那么大小有关系吗？但是当我将 BATCH_SIZE 降低到 32 或更小时出现同样的错误。

为什么第一个代码块耗尽了所有的 gpu 内存？

当然，我把

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        # Currently, memory growth needs to be the same across GPUs
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        # Memory growth must be set before GPUs have been initialized
        print(e)

此代码位于第一个单元格。

Answer 1

梯度磁带触发自动微分，这需要跟踪所有权重和激活的梯度。 Autodiff 需要更多内存。这个是正常的。您必须手动调整批量大小，直到找到一个有效的批量大小，然后再调整您的 LR。通常，曲调仅表示猜测和检查或网格搜索。（我正在开发一款产品来为您完成所有这些工作，但我不是来这里插入它的）。

使用 tf.GradientTape() 会耗尽所有的 gpu 内存，没有它也没关系

Use of tf.GradientTape() exhausts all the gpu memory, without it it doesn't matter

gpu

tensorflow

tensorflow2.0

gradienttape