使用pytorch时 报错：OSError: [Errno 12] Cannot allocate memory

问题描述

在使用Pytorch跑模型时，虽然用了GPU，但随着训练的进行，后台占用内存（不是显存）越来越多，直到内存不足，程序崩溃：

OSError: [Errno 12] Cannot allocate memory

解决办法

造成这个问题的原因是我在使用循环计算一个epoch内所有batch的loss之和时，直接累加了tensor，导致梯度信息一直保持，内存占用不断增加：

   for epoch in range(1, n_epoch+1):
        loss_epoch = 0
        model.train()
        for batch in loader:
            ......
            loss = criterion(......)
            loss.backward()
            optimizer.step()
            loss_epoch += loss # All loss in a batch

修改方法是，loss先转成浮点数，再累加：

   for epoch in range(1, n_epoch+1):
        loss_epoch = 0
        model.train()
        for batch in loader:
            ......
            loss = criterion(......)
            loss.backward()
            optimizer.step()
            loss_epoch += loss.item() # All loss in a batch