[TF 2.0] TF 2.0 consumes twice as much memory as TF 1.x or CNTK.

[KichangKim]

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow):
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Windows 10 x64 1809
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device:
• TensorFlow installed from (source or binary): pip install tensorflow-gpu==2.0.0-alpha0
• TensorFlow version (use command below): v1.12.0-9492-g2c319fb415 2.0.0-alpha0
• Python version: 3.6.7
• Bazel version (if compiling from source):
• GCC/Compiler version (if compiling from source):
• CUDA/cuDNN version: CUDA 10.0, cudnn-10.0-windows10-x64-v7.5.0.56
• GPU model and memory: GeForce GTX 1070 8GB

Describe the current behavior
Evaluating TF 2.0 keras model allocates twice as much memory as TF 1.x or CNTK.

Describe the expected behavior
Memory usage of TF 2.0 should be same or similar to other libraries, not double.

Code to reproduce the issue
For TF 2.0 or 1.x

import numpy as np
import tensorflow as tf

# tf.config.gpu.set_per_process_memory_growth(True)

size = 28000

inputs = tf.keras.Input((size,), dtype='float32')
outputs = tf.keras.layers.Dense(size)(inputs)
model = tf.keras.models.Model(inputs=inputs, outputs=outputs)

model.predict(np.ones((1, size,), dtype=np.float32))

print('complete')

while True:
    pass

For TF 1.x or CNTK with keras

import keras
import numpy as np

size = 28000

inputs = keras.Input((size,), dtype='float32')
outputs = keras.layers.Dense(size)(inputs)
model = keras.models.Model(inputs=inputs, outputs=outputs)

model.predict(np.ones((1, size,), dtype=np.float32))

print('complete')

while True:
    pass

With 8GB VRAM GPU, TF 1.x and CNTK works successfully, and TF 2.0 code are failed due to Resource exhausted exception.

[robieta]

TL;DR It's actually a somewhat niche case (a single weight is O(gpu memory size)), but it's also a pretty easy fix.

In your example, the model is failing when it builds the kernel for the dense layer; or more specifically, when it calls the initializer. If you consider a much simpler example:

a = tf.ones((size, size))
b = a + 1
c = b + 1

In graph mode when you evaluate c (which is a symbolic tensor), TensorFlow's memory allocator will reuse blocks of memory when it detects that they are no longer being used, so by the time c is computed the buffer that backed a is long gone. By contrast, in eager a handle to all three tensors co-exist in the body, so the runtime is not allowed to free them until they go out of scope on the python side. (Since a user could presumably add other computations later, unlike the graph case where all computations are known.)

Fortunately, wrapping the initializer call in a tf.function fixes the issue. In general I suspect more and more of these utility functions are going to get function'd for precisely this sort of reason.

[KichangKim]

@robieta Thanks for investigation. Can you provide sample code of tf.function usage? What is the initializer?

Also, my actual code is big CNN network with TF 2.0 keras API and has manual training loop by using tf.data.Dataset and GradientTape, samely suffer from this memory and performance issue. If you provide sample code of the usage of tf.function for keras model and manual training loop, it will be very helpful.

[robieta]

For the case I described:

size = 40000  # Slightly larger b/c I was testing on a GPU w/ 16 GB
init_fn = tf.keras.initializers.glorot_uniform()
init_fn = tf.function(init_fn, autograph=False) # <== This is what prevents the OOM (Comment it out to test)
layer = tf.keras.layers.Dense(size, kernel_initializer=init_fn)
layer.build((size,))

And my plan is to make builtin initializers do the tf.function wrap automatically.

In general functions are documented in https://www.tensorflow.org/alpha/tutorials/eager/tf_function, https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/function