Tensorflow 错误：分配具有奇怪形状的张量时出现 OOM

Question

当我收到此错误时，我正在尝试使用一些 512x512 图像执行 VGG16 的小变化：

tensorflow.python.framework.errors_impl.ResourceExhaustedError: OOM when allocating tensor with shape[131072,4096]

这是我的model.summary()

Layer (type)                 Output Shape              Param #   
=================================================================
conv2d (Conv2D)              (None, 512, 512, 64)      1792      
_________________________________________________________________
conv2d_1 (Conv2D)            (None, 512, 512, 64)      36928     
_________________________________________________________________
max_pooling2d (MaxPooling2D) (None, 256, 256, 64)      0         
_________________________________________________________________
conv2d_2 (Conv2D)            (None, 256, 256, 128)     73856     
_________________________________________________________________
conv2d_3 (Conv2D)            (None, 256, 256, 128)     147584    
_________________________________________________________________
max_pooling2d_1 (MaxPooling2 (None, 128, 128, 128)     0         
_________________________________________________________________
conv2d_4 (Conv2D)            (None, 128, 128, 256)     295168    
_________________________________________________________________
conv2d_5 (Conv2D)            (None, 128, 128, 256)     590080    
_________________________________________________________________
conv2d_6 (Conv2D)            (None, 128, 128, 256)     590080    
_________________________________________________________________
max_pooling2d_2 (MaxPooling2 (None, 64, 64, 256)       0         
_________________________________________________________________
conv2d_7 (Conv2D)            (None, 64, 64, 512)       1180160   
_________________________________________________________________
conv2d_8 (Conv2D)            (None, 64, 64, 512)       2359808   
_________________________________________________________________
conv2d_9 (Conv2D)            (None, 64, 64, 32)        147488    
_________________________________________________________________
max_pooling2d_3 (MaxPooling2 (None, 32, 32, 32)        0         
_________________________________________________________________
flatten (Flatten)            (None, 32768)             0         
_________________________________________________________________
dense (Dense)                (None, 4096)              134221824 
_________________________________________________________________
dense_1 (Dense)              (None, 4096)              16781312  
_________________________________________________________________
dense_2 (Dense)              (None, 3)                 12291     
=================================================================
Total params: 156,438,371
Trainable params: 156,438,371
Non-trainable params: 0

这里奇怪的是错误信息中的shape[131072, 4096]。

在我的模型中，正如我们在model.summary() 中看到的那样，唯一的4096 位于最后，在Dense 层中，在flatten 层之后。但是 flatten 层会产生32768 神经元的输出。所以，这种情况下的错误应该是shape[32768, 4096]。我尝试写 2048 或 1024 而不是 4096。但错误总是出现在 shape[131072, 4096] 上，而且这个形状永远不会改变。

问题是：

1. shape[131072, 4096] 来自哪里？不应该是shape[32768, 4096]吗？
2. 如何解决？

我读了一些有同样问题的问题，比如我们：

Shape of image after MaxPooling2D with padding ='same' --calculating layer-by-layer shape in convolution autoencoder

或

Error: OOM when allocating tensor with shape

但他们都没有解释为什么形状不正确。

Answer 1

由于您使用的是 VGG16，因此您需要确保您的输入是 224x224，而您的情况并非如此。另外，你能告诉你正在使用的批量大小吗，错误是因为你没有足够的 GPU 内存。因此，请考虑减少批量大小。鉴于您使用的是 512x512，GPU 没有足够的内存是可以理解的。此外，作为最佳实践，请记住大多数 SOTA 模型的输入图像永远不要超过 300x300。建议不要使用这么大的尺寸。

至于您要问的形状，让我问您是否使用批量大小 32。因为 32 x 256 x 256 x 64 = 131072 x 4096。因此张量实际上是正在创建的初始张量，但是由于内存较少，无法创建。至于为什么会出现这样的形状可能是由于 TensorFlow 的内部工作，我无法对此发表评论，因为我对此的了解有限，对此我深表歉意。但是，我希望我能帮助您了解形状的来源。