Tensorflow Float16 for VGG19 model parameters

Question

I'm doing neural style transfer using TensorFlow(with the 19.01 Nvidia TensorFlow Docker image) in Python. I have an Nvidia 2070 graphics card, and I've been getting Out of Memory errors when I try to run a newer version of the TensorFlow docker image(19.08 for example). So I decided that perhaps it is time to consider using 16 bit precision instead of 32 bit for storing the parameters of the VGG19 CNN.

My initial research when I built my machine had led me to believe that switching from 32 to 16 was a cakewalk, but that hasn't been my experience now that I'm actively trying to make the transition.

This is what I have done:

1. I set tf.keras.backend.set_floatx('float16')
2. Set tf.keras.backend.set_epsilon(1e-4)
3. Change my image input to the VGG19 network to a float16, and any other miscellaneous parts of my code that use the float32 datatype in conjunction with the float16.

When I run the code, nvidia-smi still reports that essentially 100% of my GPU is being used. Has anyone had any success with reducing their model memory footprint by switching to float16 in TensorFlow?

Answer 1

TensorFlow has various ways of managing mixed precision. The most suitable mixed precision approach depends on which optimizer you plan to use. Keras optimizers, for example, have an API designed to easily port code one way or the other . It is called mixed_precision .

If you are using a TFv1 optimizer, or one of the other non keras optimizers offered by TensorFlow, you can use their graph rewrite function to convert various pieces of the graph to float 16.