在Nvidia RTX2080 Ti上使用Tensorflow2.0的“简单” ResNet50进行OOM

Question

我很惊讶在Nvidia RTX2080Ti（具有11Gb内存！）上使用tf.keras.applications.ResNet50实现遇到内存不足错误。

题：

我使用的工作流程有问题吗？

笔记：

• 我正在使用tensorflow-gpu==2.0.0b1和CUDA v10.1
• 我上分割任务工作，从而大output_shape
• 我自己构建批次，因此使用train_on_batch()
• 设置即使memory_growth至True ，内存被填满时从700MB到10850Mb在几分之一秒。

码：

import tensorflow as tf
import tensorflow.keras as ke
import numpy as np

ke.backend.clear_session()

inputs = ke.layers.Input(shape=(512,1024,3), dtype="float32")
outputs = ke.applications.ResNet50(include_top=False, weights="imagenet")(inputs)
outputs = ke.layers.Lambda(lambda x: tf.compat.v1.image.resize_bilinear(x, size=(512,1024)))(outputs)
outputs = ke.layers.Conv2D(2, 1, activation="softmax")(outputs)
model = ke.Model(inputs=inputs, outputs=outputs)

model.compile(optimizer=ke.optimizers.RMSprop(lr=0.001), loss=ke.losses.CategoricalCrossentropy())

images = np.zeros((1,512,1024,3), dtype=np.float32)
targets = np.zeros((1,512,1024,2), dtype=np.float32)

model.train_on_batch(images, targets)

Answer 1

Resnet是复杂的复杂模型，输入的尺寸可能是OOM错误的原因。 尝试减小尺寸和相应的批处理大小（在内存可以容纳的范围内）并尝试。

如评论中所述，它适用于批处理大小为1，尺寸为700 * 512。