Debug Tensorflow: 随着训练进行，内存消耗越来越大

环境

• ubuntu 18.04
• Python 3.8
• TensorFlow-gpu 2.3.1
• CUDA 11.1
• Tensorflow-yolov5

问题

内存泄漏。训练时消耗内存越来越多，直到内存完全被占用，服务器连接被迫断开。无报错。正常情况是训练时，内存不应该增加

解决方案

1. 监控linux内存情况
   

2. 安装memory-profiler
   

3. 定位问题所在

from memory_profiler import profile
fp=open('memory_profiler.log','w+')

@profile(stream=fp)
def train_step2(self, image_data, target):
    with tf.GradientTape() as tape:
        # print(tf.reduce_mean(image_data).numpy(), 'matched', tf.reduce_sum(target[0][:,:,:,:,4]).numpy(), tf.reduce_sum(target[1][:,:,:,:,4]).numpy() , tf.reduce_sum(target[2][:,:,:,:,4]).numpy())
        image = image_data
        pred_result = self.model(image, training=True)      

4. 修改代码
   我解决这个问题比较奇怪，删掉shuffle之后就好了。或者说解决了绝大部分，后面好像还是会内存泄漏，但只有一点点了，几十个epoch内存大概上升1%，可以接受。

# if shuffle:
        #     dataset = dataset.shuffle(buffer_size=1111, seed=1949)

其他情况

keras自带有问题：https://github.com/tensorflow/tensorflow/issues/32052