Memory leak in Conv2D/Activation on GPU

[jan-x-marek]

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 20.04
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device:
• TensorFlow installed from (source or binary): Binary, the standard docker distribution
• TensorFlow version (use command below): v2.4.0-rc4-71-g582c8d236cb 2.4.0
• Python version: 3.6.9
• Bazel version (if compiling from source):
• GCC/Compiler version (if compiling from source):
• CUDA/cuDNN version: 11.0
• GPU model and memory: GeForce RTX 2070, 8GB

Describe the current behavior
I upgraded to TF 2.4.0 from TF 2.1.2, and training a very simple convolutional network, which worked fine in 2.1.2, started running out of memory during training. I distilled a simple reproducible example that demonstrates the issue. Each training epoch consumes about 50MB of additional memory and, given enough epochs, it grows to infinity (or 32 GB in my case). It only occurs on GPU, the same thing runs fine on CPU.

Describe the expected behavior
Memory not growing, or growing only very little

Standalone code to reproduce the issue

import gc
import os
import psutil
import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Dense, Conv2D, Flatten, BatchNormalization, Activation

physical_devices = tf.config.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(physical_devices[0], True)


input_tensor = tf.keras.layers.Input(shape=(512,64,1))

x = Conv2D(filters=32, kernel_size=(5,5), strides=(2,2), padding='same')(input_tensor)
# Commented out on purpose - see Note 1 below
# x = BatchNormalization()(x)
x = Activation('relu')(x)

x = Conv2D(filters=64, kernel_size=(4,4), strides=(2,2), padding='same')(x)
# Commented out on purpose - see Note 1 below
# x = BatchNormalization()(x)
x = Activation('relu')(x)

x = Conv2D(filters=128, kernel_size=(4,4), strides=(2,1), padding='same')(x)
# Commented out on purpose - see Note 1 below
# x = BatchNormalization()(x)
x = Activation('relu')(x)

x = Conv2D(filters=128, kernel_size=(4,4), strides=(2,1), padding='same')(x)
# Commented out on purpose - see Note 1 below
# x = BatchNormalization()(x)
x = Activation('relu')(x)

x = Flatten()(x)

x = Dense(5, activation='sigmoid')(x)

model = tf.keras.Model(inputs=input_tensor, outputs=x)


train_x = np.random.random((2048, 512, 64, 1))
train_y = np.random.random((2048, 5))

model.compile(loss='binary_crossentropy', optimizer=tf.keras.optimizers.Adam())

process = psutil.Process(os.getpid())

for i in range(50):
    model.fit(train_x, train_y, epochs=1, batch_size=32, verbose=0)
    gc.collect()
    print(i, process.memory_info().rss // 1000000)

Note 1
Now, if you uncomment the BatchNormalization() layers creation, the memory problem disappears. So, it is somehow caused by the Activation layer following immediately the Conv2D

Note 2
The memory problem also occurs if I train multiple epochs in a single fit() call, such as

model.fit(train_x, train_y, epochs=50, batch_size=32)

I used the for loop only to be able to call garbage collection and print the memory.

Note 3
A Conv2D layer with activation embedded in it, such as

Conv2D(filters=128, kernel_size=(4,4), strides=(2,1), padding='same', activation='relu')

also causes the memory issue

[zlee406]

I think I have the same issue. I've narrowed it a bit further. Memory usage continually grows until no more memory is available. Here is my minimal working example:

import tensorflow as tf
import numpy as np
import sys
physical_devices = tf.config.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(physical_devices[0], True)

def create_model():
    inputs = tf.keras.Input(shape=(80, 80, 4))
    x = tf.keras.layers.Conv2D(32, (8, 8), strides=(4, 4), activation='relu')(inputs)
    x = tf.keras.layers.Conv2D(64, (4, 4), strides=(2, 2), activation='relu')(x)
    x = tf.keras.layers.Flatten()(x)
    outputs = tf.keras.layers.Dense(5, activation='linear')(x)

    model = tf.keras.Model(inputs=inputs, outputs=outputs)
    return model

model = create_model()

@tf.function
def get_action(state):
    action_values = model(state)
    return tf.math.argmax(action_values, axis=1)

for i in range(150000):
    states = np.random.randint(0, 255, (32, 80, 80, 4))
    s = get_action(states)
    if i % 50 == 0:
        sys.stdout.write(f'\rIter: {i}')

A few things I've noticed:

1. It functions normally without the convolution layers (both of them). Similar to the original post's Note 1.
2. It functions normally without tf.function or on CPU

Update:

1. The memory leak only occurs with ReLu activation function. LeakyRelu does not cause the memory leak unless setting alpha=0.
2. Tanh activation causes a crash with exit code: -1073741571 (0xC00000FD)

[Saduf2019]

@jan-x-marek
i ran the code shared and do not see any memory leak on colab, please find the gist here.

[zlee406]

@Saduf2019 I am not getting the issue in the colab. @jan-x-marek and I are both using Cuda 11.0, but the colab is on 10.1. From what I could tell here, TF 2.4 uses Cuda 11.0, correct?

[jan-x-marek]

@Saduf2019

As @zlee406 says, your Colab notebook is running on CUDA V10.1.243. Me and @zlee406 are running on CUDA 11 (V11.0.221 in my case), as required by the install guide: https://www.tensorflow.org/install/gpu

This is how my memory grows when I run the sample code:
0 4194
1 4280
2 4376
3 4439
4 4530
5 4583
6 4614
7 4670
8 4731
9 4773
10 4802
11 4871
12 4920
13 4960
14 5016
15 5100
16 5159
17 5228
18 5289
19 5367
20 5422
21 5469
22 5526
23 5561
24 5598
25 5659
26 5726
27 5786
28 5846
29 5873
30 5959
31 6008
32 6065
33 6155
34 6198
35 6252
36 6287
37 6338
38 6412
39 6456
40 6546
41 6598
42 6653
43 6730
44 6783
45 6800
46 6874
47 6926
48 6967
49 7064

[Uture]

I have the same issue. It seems that the problem only occurs when using Conv2D with stride != 1.

[zlee406]

A potential substitute until the issue is resolved is just to replace the ReLu activation with tf.where(x <= 0., 0., x)