环境
1.系统环境:ubuntu 18.0.4
2.MegEngine版本:0.3.1
3.python版本:3.7.6
使用gpu训练mnist数据集 Lenet时,占用gpu内存不断增加,直至gpu内存溢出(4g内存,batch 128)
请提供关键的代码片段便于追查问题
示例训练代码 gpu训练mnist数据集 Lenet时 加入两层dropout
去除drop 不会出问题
class LeNet(M.Module):
def __init__(self):
super(LeNet, self).__init__()
# 单信道图片, 两层 5x5 卷积 + ReLU + 池化
self.conv1 = M.Conv2d(1, 6, 5)
self.relu1 = M.ReLU()
self.pool1 = M.MaxPool2d(2, 2)
self.conv2 = M.Conv2d(6, 16, 5)
self.relu2 = M.ReLU()
self.pool2 = M.MaxPool2d(2, 2)
# 两层全连接 + ReLU
self.fc1 = M.Linear(16 * 5 * 5, 120)
self.relu3 = M.ReLU()
self.drop1 = M.Dropout(0.5)
self.fc2 = M.Linear(120, 84)
self.relu4 = M.ReLU()
self.drop2 = M.Dropout(0.5)
# 分类器
self.classifer = M.Linear(84, 10)
def forward(self, x):
x = self.pool1(self.relu1(self.conv1(x)))
x = self.pool2(self.relu2(self.conv2(x)))
# F.flatten 将原本形状为 (N, C, H, W) 的张量x从第一个维度(即C)开始拉平成一个维度,
# 得到的新张量形状为 (N, C*H*W) 。 等价于 reshape 操作: x = x.reshape(x.shape[0], -1)
x = F.flatten(x, 1)
x = self.relu3(self.fc1(x))
x = self.drop1(x)
x = self.relu4(self.fc2(x))
x = self.drop2(x)
x = self.classifer(x)
return x
请提供完整的日志及报错信息
egBrainError: MegBrain core throws exception: mgb::MegDNNError
curand call failed: status=102(CURAND_STATUS_ALLOCATION_FAILED) call=curandGenerateUniform(m_curand_handle.gen(), dst.ptr<dt_float32>(), dst.layout.total_nr_elems())
解答
动态图模式下存在以下问题,我们将持续优化:
- 当前实现导致动态创建的 Tensor 显存不能自动回收,目前需要使用 set_value() 方法手动复用 Tensor 来避免显存持续增加
- 动态模式下的显存占用较高
- PyTorch 子图,megengine.Funtion 等算子占用显存会持续增加
- random operator 以及依赖它的 operator 暂时无法去重,导致显存使用持续增加,例如 dropout