環境

• 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_datapred_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

與50位技術專家面對面20年技術見證，附贈技術全景圖

總結

以上是生活随笔為你收集整理的Debug Tensorflow: 随着训练进行，内存消耗越来越大的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。