尝试在 Google Colab 上训练 ResNet 时出现 ResourceExhaustedError

Question

我正在尝试在自定义数据集上在 Google Colab 上训练 ResNet56，其中每个图像的尺寸为 299x299x1。这是我得到的错误：

ResourceExhaustedError:  OOM when allocating tensor with shape[32,16,299,299] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
     [[node resnet/conv2d_21/Conv2D (defined at <ipython-input-15-3b824ba8fe2a>:3) ]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.
 [Op:__inference_train_function_21542]

Function call stack:
train_function

这是我的模型配置：

TRAINING_SIZE = 9287
VALIDATION_SIZE = 1194

AUTO = tf.data.experimental.AUTOTUNE # used in tf.data.Dataset API
BATCH_SIZE = 32

model_checkpoint_path = "/content/drive/My Drive/Patch Classifier/Data/patch_classifier_checkpoint"
if not os.path.exists(model_checkpoint_path):
    os.mkdir(model_checkpoint_path)

CALLBACKS = [
              EpochCheckpoint(model_checkpoint_path, every=2, startAt=0),
              TrainingMonitor("/content/drive/My Drive/Patch Classifier/Training/resnet56.png",
                              jsonPath="/content/drive/My Drive/Patch Classifier/Training/resnet56",
                              startAt=0)
              ]

compute_steps_per_epoch = lambda x: int(math.ceil(1. * x / BATCH_SIZE))
steps_per_epoch = compute_steps_per_epoch(TRAINING_SIZE)
val_steps = compute_steps_per_epoch(VALIDATION_SIZE)

opt = SGD(lr=1e-1)
model = ResNet.build(299, 299, 1, 5, (9, 9, 9), (64, 64, 128, 256), reg=0.005)
model.compile(loss="categorical_crossentropy", optimizer=opt, metrics=["accuracy"])

history = model.fit(get_batched_dataset("/content/drive/My Drive/Patch Classifier/Data/patch_classifier_train_0.tfrecords"), steps_per_epoch=steps_per_epoch, epochs=10,
                    validation_data=get_batched_dataset("/content/drive/My Drive/Patch Classifier/Data/patch_classifier_val_0.tfrecords"), validation_steps=val_steps,
                    callbacks=CALLBACKS)

有什么想法吗？

Answer 1

如果内存不足，您可以做的事情并不多。

我能想到的要么是

1. 减少BATCH_SIZE
2. 减小图像输入大小。

如果您选择减小批量大小，那么您可能还需要降低学习率，如果您觉得它没有收敛。

P.S：如果你把动力放在那里，SGD 会做得更好，比如SGD(lr=1e-1, momentum=0.9)

Answer 2

我也遇到了同样的错误，这是因为我使用了 512*512 的图像大小和 10 的批处理大小。 我将 batch size 减少到 2，它开始为我工作。