情况：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在进行深度学习训练的时候，有4大部分的显存开销，分别是
①模型参数(parameters)
②模型参数的梯度(gradients)
③优化器状态(optimizer states)
④中间激活值(intermediate activations) 或者叫中间结果(intermediate results)。
可以根据这几部分考虑排查下。
排查过程：
-
确认使用 with torch.no_grad()
-
排查模型结构中每一步的显存占用
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)
Original: https://blog.csdn.net/xxzhix/article/details/124365824
Author: xxzhix
Title: 模型预测时GPU占用不断增长
原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/711179/
转载文章受原作者版权保护。转载请注明原作者出处！