情况: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)。
可以根据这几部分考虑排查下。
排查过程:
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)