pytorch内存泄漏

news2024/9/25 11:20:55

问题描述：

内存泄漏积累过多最终会导致内存溢出，当内存占用过大，进程会被killed掉。

解决过程：

在代码的运行阶段输出内存占用量，观察在哪一块存在内存剧烈增加或者显存异常变化的情况。但是在这个过程中要分级确认问题点，也即如果存在三个文件main.py、train.py、model.py。
在此种思路下，应该先在main.py中确定问题点，然后，从main.py中进入到train.py中，再次输出显存占用量，确定问题点在哪。随后，再从train.py中的问题点，进入到model.py中，再次确认。如果还有更深层次的调用，可以继续追溯下去。

import psutil
process = psutil.Process()
current_memory = process.memory_info().rss
print(f"0--------------Current memory usage: {current_memory / (1024 ** 3):.4f} GB")

具体使用的代码

for epoch in range(start_epoch+1, args.epochs+1):
        process = psutil.Process()
        current_memory = process.memory_info().rss
        print(f"0--------------Current memory usage: {current_memory / (1024 ** 3):.4f} GB")
        count_step = (epoch-1)*len(train_loader)  
        print(f"1--------------Current memory usage: {current_memory / (1024 ** 3):.4f} GB")
        mean_loss, lr = train_one_epoch(model, optimizer, train_loader, device, epoch, count_step,writer,lr_scheduler,
                                         print_freq=args.print_freq)
        print(f"2--------------Current memory usage: {current_memory / (1024 ** 3):.4f} GB")
        val_info = evaluate_vgg(model, epoch, val_loader, device, writer, num_classes=num_classes)
        print(f"3--------------Current memory usage: {current_memory / (1024 ** 3):.4f} GB")
        with open(results_file, "a") as f:
            # 记录每个epoch对应的train_loss、lr以及验证集各指标       
            train_info = f"[epoch: {epoch}]\n" \
                         f"train_loss: {mean_loss:.4f}\n" \
                         f"lr: {lr:.6f}\n"
            f.write(train_info + val_info + "\n\n")

        save_vgg_file = {"model": model.state_dict(),
                     "optimizer": optimizer.state_dict(),
                    #  "lr_scheduler": lr_scheduler.state_dict(),
                     "epoch": epoch,
                     "args": args}
        
        torch.save(save_vgg_file, 'checkpoints/fcn_model_Adam-StepLR_1e-2.pth')
        print(f"update checkpoints/fcn_model_Adam-StepLR_1e-2.pth")
        print(f"4--------------Current memory usage: {current_memory / (1024 ** 3):.4f} GB")

在这里插入图片描述
每个epoch训练完之后所占内存会不断增加，也就是说，每轮跑完之后会有冗余的数据一直在消耗内存。于是criterion、train_one_epoch、evaluate三个部分

criterion部分
Mem usage：5310 MiB 在这里插入图片描述 train_one_epoch部分
Mem usage：4439 MiB

evaluate部分
Mem usage：10644

在evaluate部分可以看到，所占用内存突然增大，并且之后的代码也占用了大量内存，继续监控得知在下一个epoch中criterion部分占用内存也是16064MiB，由此推测出内存消耗在evaluate部分。