Q: 动态图显存一直增长

环境

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

GPU训练,不使用dropout,动态图训练显存无变化,但是用静态图的时候,显存也一直涨,只是涨的慢,不知道是我用错了么

方便提供一下你的训练代码么?

from lenet5 import LeNet
import megengine as mge
import megengine.optimizer as optim
import megengine.functional as F
from megengine.data import DataLoader
from megengine.data.transform import ToMode, Pad, Normalize, Compose
from megengine.data import RandomSampler
from megengine.data.dataset import MNIST
from megengine.jit import trace
import time

star = time.time()


@trace(symbolic=True)
def train_func(data, label, *, opt, net):
    net.train()
    logits = net(data)
    loss = F.cross_entropy_with_softmax(logits, label)
    opt.backward(loss)
    return logits, loss


mnist_train_dataset = MNIST(root="./dataset/MNIST", train=True, download=True)
dataloader = DataLoader(
    mnist_train_dataset,
    transform=Compose([
        Normalize(mean=0.1307*255, std=0.3081*255),
        Pad(2),
        ToMode('CHW'),
    ]),
    sampler=RandomSampler(dataset=mnist_train_dataset, batch_size=64, drop_last=True),
)

le_net = LeNet()
optimizer = optim.SGD(le_net.parameters(), lr=0.05)

trace.enabled = True

total_epochs = 100

for epoch in range(total_epochs):
    total_loss = 0
    correct = 0
    total = 0
    for step, (data, label) in enumerate(dataloader):
        optimizer.zero_grad()

        label = label.astype('int32')

        logits, loss = train_func(data, label, opt=optimizer, net=le_net)

        optimizer.step()
        total_loss += loss.numpy().item()
        predicted = F.argmax(logits, axis=1)

        correct += (predicted == label).sum().numpy().item()
        total += label.shape[0]

    print("epoch: {}, loss {}, acc {}".format(epoch, total_loss/len(dataloader), correct/total))

print("use time :", time.time() - star)
path = "lenet_j.mge"
mge.save(le_net.state_dict(), path)

原因是 F.argmax 这个仍然跑在了动态图模式下,每次的输入都不一样(logits)所以导致了显存逐步增长。建议这里暂时替换成 numpy.argmax 来避免显存增长。

好了,谢谢 :+1:,文档这么写的,就按照那个来的

收到,我们来更新一下文档。