使用Gluon训练模型时,第一个epoch后,占用显存会提升,之后平稳不再上涨。

如上图所示,在模型训练第一个epoch后,GPU显存占用会明显提升,之后平稳。
在每一个训练batch之后增加mx.nd.waitall(),依旧无法缓解该问题。

该问题影响训练时可以调整的最大网络计算量,不知是否有人会遇到相同的问题。

训练pipeline如下:

        batch_data, batch_label = self.batch_fn(batch)

        with mx.autograd.record():
            batch_output = [self.network(*data) for data in batch_data]

            batch_loss = [self.loss(output, label)
                            for output, label in zip(batch_output, batch_label)]

            losses = []
            for i_loss in batch_loss:  # for each device
                for j_loss in i_loss:  # for each output
                    losses.append(j_loss)
            mx.autograd.backward(losses)

        self.trainer.step(1)

因为第一个epoch结束后会进行第一次val测试,所以在第一个epoch后显存会增加,之后保持平稳。

with torch.zero_grad():