Tensorflow Error: OOM when allocating tensor with strange shape

Question

I'm trying to execute a small variation of VGG16 with some 512x512 images when I get this error:

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

Here is my 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

The strange thing here is the shape[131072, 4096] in the error message.

In my model, as we can see in the model.summary(), the only 4096 is at the end, in the Dense layer, after the flatten layer. But the flatten layer produces an output of 32768 neurons. So, the error in this case should be shape[32768, 4096]. I tried to write 2048 or 1024 instead of 4096. but the error is always with shape[131072, 4096], and this shape never changes.

The questions are:

1. Where does this shape[131072, 4096] come from? Shouldn't it be shape[32768, 4096]?
2. How can it be fixed?

I read some questions with the same problem, such us:

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

or

Error: OOM when allocating tensor with shape

but none of them explain why the shape is incorrect.

Answer 1

Since you are using VGG16, you need to make sure that your input is 224x224, which is not the case with you. Also, can you tell the batch size you are using, the error is because you don't have enough GPU memory. So, consider reducing the batch size. Given that you are using 512x512, it is understandable for GPU to not have enough memory. Also, as a best practice remember that most of the SOTA models, never go above 300x300 for input images. It would be advisable to refrain from using such big sizes.

As for the shape you are asking, let me ask you if you are using batch size 32. Because 32 x 256 x 256 x 64 = 131072 x 4096. And thus the tensor is actually the initial tensor that is being created but can't be created because of less memory. As to why the shape is like this may be due to the internal working of TensorFlow and I can't comment on that, because of my limited knowledge on it for which I apologize. But, I hope I have helped you understand where the shape came from.