模型预测时GPU占用不断增长

情况：pytorch框架，深度模型单GPU预测时，显存占用会不断增长，由有1G增长到5G，甚至更多。

我这边排查下来的原因是卷积的输入尺寸变化，有很大尺寸的输入导致GPU占用不断增长。如果你的模型中没有使用卷积层，可以不往下阅读。

如何看出显存占用不断增长？

登录代码运行服务器，使用 watch -n 1 nvidia-smi ，观察显存变化

预测部分代码如下：

with torch.no_grad():
    # 模型预测
    output = self.model(**input_data_item)
    
    #获得输出概率和输出类别
    _, predict = torch.max(prob, dim=-1)

    # 将输出类别转换成中文类别
    predict_labels = [labels_id2b[p] for p in predict]
 

对pytorch的显存占用说明可以参考(81条消息) Pytorch显存机制与显存占用(一) —— 理论分析(memory)_hxxjxw的博客-CSDN博客_pytorch 显存

pytorch在进行深度学习训练的时候，有4大部分的显存开销，分别是

      ①模型参数(parameters)

      ②模型参数的梯度(gradients)

      ③优化器状态(optimizer states)

      ④中间激活值(intermediate activations) 或者叫中间结果(intermediate results)。

可以根据这几部分考虑排查下。

排查过程：

1. 确认使用 with torch.no_grad()

 2. 排查模型结构中每一步的显存占用

    def forward(self, x1, x2, x3, x4, x5, x6, x7):
        print("init", torch.cuda.max_memory_allocated()//1024**2)
        embed_1 = self.encoder1(x4, x5)
        print("encoder1", torch.cuda.max_memory_allocated()/1024**2)
        embed_2 = self.encoder2(x1, x6, x7, x2, x3, )
        print("encoder2", torch.cuda.max_memory_allocated()/1024**2)
        _, hidden_embed = self.graph(embed_1, embed_2)
        print("graph", torch.cuda.max_memory_allocated()/1024**2)

打印结果为

# 这是模型的第一个batch输入之后的打印结果
init 69.58447265625  # 69M为加载的模型的大小
encoder1 69.58740234375
encoder2 294.62255859375
graph 294.62255859375

 可以看到encoder2的显存占用变大了，但这不是GPU占用不稳定增长的原因，这只是表示encoder2需要使用更多显存来计算，是正常的。

当第二个，第三个，..., 第k个batch的显存占用输出可能是稳定的，类似

init 294.62255859375 
encoder1 294.62255859375
encoder2 294.62255859375
graph 294.62255859375

 当第K+1个batch的打印如下

init 294.62255859375 
encoder1 294.62255859375
encoder2 698.58624792378
graph 698.58624792378

可以看到在这个batch内，encoder2这一层的占用突然增长，这是可以用相同的方式不断打印encoder2层内的每一步计算，定位导致GPU增长的那行代码。

我这边定位到

# images_segments是在计算过程中动态生成的，
# 前k个batch的image_segments的形状是 torch.Size([625, 512, 7, 7])
# 第k+1个batch的image_segments的形状是 torch.Size([3024, 512, 7, 7])
# 导致了显存占用增长，就算后面batch的形状变小了，占用的显存也不释放了，
# 于是遇到更大形状的输入就不断增长
image_segments_out = self.conv(image_segments)

     

参考链接：(81条消息) Pytorch显存机制与显存占用(一) —— 理论分析(memory)_hxxjxw的博客-CSDN博客_pytorch 显存