为啥我的 Colab 会话内存不足？

Posted 2023-02-16

【中文标题】为啥我的 Colab 会话内存不足？

【英文标题】：Why does my Colab session run out of RAM?为什么我的 Colab 会话内存不足？

【发布时间】：2021-12-13 06:49:07

【问题描述】：

我正在使用 Keras 基于此 paper 中描述的模型构建图像去模糊模型。我使用以下训练代码在 Colab 上训练模型：

x_train, y_train = load_h5_dataset()

def train(batch_size=16, epoch_num=5, critic_updates=5, log_dir='drive/MyDrive/train_logs'):
    g = make_resnet_generator_model()
    d = make_discriminator_model()
    gan = make_gan(g, d)

    d_opt = Adam(learning_rate=1e-4, beta_1=0.9, beta_2=0.999, epsilon=1e-8)
    gan_opt = Adam(learning_rate=1e-4, beta_1=0.9, beta_2=0.999, epsilon=1e-8)

    d.trainable = True
    d.compile(optimizer=d_opt, loss=wasserstein_loss)
    d.trainable = False
    loss = [perceptual_loss, wasserstein_loss]
    loss_weights = [100, 1]
    gan.compile(optimizer=gan_opt, loss=loss, loss_weights=loss_weights)
    d.trainable = True

    output_true_batch, output_false_batch = np.ones((batch_size, 1)), -np.ones((batch_size, 1))
    writer = tf.summary.create_file_writer(log_dir)

    for epoch in tqdm(range(epoch_num)):
        print(f"Epoch epoch + 1/epoch_num...")
    
        permuted_indexes = np.random.permutation(x_train.shape[0])

        d_losses = []
        gan_losses = []
        x_train = dataset['sharp_img']
        for index in range(int(x_train.shape[0] / batch_size)):
            batch_indexes = permuted_indexes[index * batch_size:(index + 1) * batch_size]
            image_blur_batch = x_train[batch_indexes]
            image_full_batch = y_train[batch_indexes]

            generated_images = g.predict(x=image_blur_batch, batch_size=batch_size)

            for _ in range(critic_updates):
                d_loss_real = d.train_on_batch(image_full_batch, output_true_batch)
                d_loss_fake = d.train_on_batch(generated_images, output_false_batch)
                d_loss = 0.5 * np.add(d_loss_fake, d_loss_real)
                d_losses.append(d_loss)

            d.trainable = False

            gan_loss = gan.train_on_batch(image_blur_batch, [image_full_batch, output_true_batch])
            gan_losses.append(gan_loss)

            d.trainable = True

    write_logs(writer, ['d_loss', 'gan_loss'], [np.mean(d_losses), np.mean(gan_losses)], epoch)
    save_weights(d, g, epoch, int(np.mean(gan_losses)))

在上面的训练代码中，感知损失是使用在 ImageNet 上预训练的 VGG16 网络计算的。函数load_h5_dataset() 用于加载保存为.hdf5 文件的数据集。执行此代码时遇到两个问题：

当我在 Colab 上运行它时，它一直在 Colab 上耗尽 RAM 并停止执行。但是，数据集的大小为 6GB，远低于 Colab 的可用 RAM 大小。 当我在本地机器（有 16GB 内存和 6GB 容量的 NVIDIA GeForce GTX 1660 Ti）上运行此代码时，我遇到了这个错误：tensorflow.python.framework.errors_impl.ResourceExhaustedError: OOM when allocating tensor with shape[16,256,128,128] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc [Op:Conv2D]

有人可以看看我的代码，看看这里出了什么问题吗？非常感谢。

【参考方案1】：

你能检查一下这个问题吗https://github.com/tensorflow/models/issues/1993

你也可以

del whatevervariable

然后 RAM 将是空闲的

【讨论】：

我已经尝试过减小批量大小和删除冗余数组，但问题仍然存在。

尝试降低数据。

并尝试降低模型的架构