环境
- ubuntu 18.04
- Python 3.8
- TensorFlow-gpu 2.3.1
- CUDA 11.1
- Tensorflow-yolov5
问题
内存泄漏。训练时消耗内存越来越多,直到内存完全被占用,服务器连接被迫断开。无报错。正常情况是训练时,内存不应该增加
解决方案
-
监控linux内存情况
-
安装memory-profiler
-
定位问题所在
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)
- 修改代码
我解决这个问题比较奇怪,删掉shuffle之后就好了。或者说解决了绝大部分,后面好像还是会内存泄漏,但只有一点点了,几十个epoch内存大概上升1%,可以接受。
# if shuffle:
# dataset = dataset.shuffle(buffer_size=1111, seed=1949)
其他情况
keras自带有问题:https://github.com/tensorflow/tensorflow/issues/32052