2

我正在研究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）。

tf.GradientTape() as tape我为基本的mnist数据尝试了相同的格式，然后它就可以正常工作了。
那么尺寸重要吗？但是当我降低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)

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

4

1 回答 1

1

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

于 2022-01-07T02:11:54.420 回答