machrisaa / tensorflow-vgg

VGG19 and VGG16 on Tensorflow
2.21k stars 1.08k forks source link

Memory Usage #28

Closed jhoh10 closed 7 years ago

jhoh10 commented 7 years ago

I am looking to fine tune the VGG19 model for an image classification problem, and am running into memory issues on my GPU.

I was wondering how much memory is expected to be used when training this network? When running it on a K80 with a 11 gb available, it was using 6.4 gb while training on just the two demo images.

I have heard of people running vgg19 on graphics cards with less than 4 gb of memory, so I was wondering if there room for improvement reducing the memory usage of this implementation?

machrisaa commented 7 years ago

Hi @jhoh10, my GPU only has 2GB and at most 1.8GB is available for training. I am afraid that the memory size cannot be reduced. But you may move some calculations to be calculated by the CPU and use the general memory instead.

The bottleneck is in the fully connected layers. In most of my project that using the VGG, when I found the GPU memory issue, I will add with tf.device('/cpu:0'): to the largest layer or even the second largest layer.

Hope that is it is helpful.