使用 Pytorch 进行训练：由于 CUDA 内存问题导致错误

Question

我正在尝试在 Cityscapes 数据集上训练模型以进行分割。我使用 torchvision deeplabv3_resnet50 模型及其 Cityscapes 数据集 class 和转换。万一重要，我是运行 Jupyter notebook 中的代码。

数据集和数据加载器都在工作。当我尝试训练时，我总是在第一批试图通过网络（one_epoch 函数中的 y_ = net(xb)）时出现此错误。

运行时错误：CUDA 内存不足。尝试分配 128.00 MiB（GPU 0；6.00 GiB 总容量；4.20 GiB 已分配；6.87 MiB 空闲；PyTorch 总共保留 4.20 GiB）

奇怪的是，无论批量大小（bs）是多少，根据错误的空闲内存量都是一个比尝试分配的内存量少一点的值，例如对于 bs=16 我得到：

运行时错误：CUDA 内存不足。尝试分配 2.00 GiB（GPU 0；6.00 GiB 总容量；2.90 GiB 已分配；1.70 GiB 可用；PyTorch 总共保留 2.92 GiB）

我有一个更复杂的模型运行，它适用于 bs=16。该模型从头开始构建所有内容。但我真的希望能够使用 torchvision 的模型动物园和数据集似乎具有的简单性。

我的代码在下面，仅是最基本的代码，足以显示它是否运行在 GPU 上正常。

def one_epoch(net, loss, dl, opt=None, metric=None):

if opt:
    net.train()  # only affects some layers
else:
    net.eval()
    rq_stored = []
    for p in net.parameters():
        rq_stored.append(p.requires_grad)
        p.requires_grad = False

L, M = [], []
dl_it = iter(dl)
for xb, yb in tqdm(dl_it, leave=False):
    xb, yb = xb.cuda(), yb.cuda()
    y_ = net(xb)
    l = loss(y_, yb)
    if opt:
        opt.zero_grad()
        l.backward()
        opt.step()
    L.append(l.detach().cpu().numpy())
    if metric: M.append(metric(y_, yb).cpu().numpy())

if not opt:
    for p,rq in zip(net.parameters(), rq_stored): p.requires_grad = rq

return L, M

accuracy = lambda y_,yb: (y_.max(dim=1)[1] == yb).float().mean()

def fit(net, tr_dl, val_dl, loss=nn.CrossEntropyLoss(), epochs=3, lr=3e-3, wd=1e-3):   

opt = optim.Adam(net.parameters(), lr=lr, weight_decay=wd)

Ltr_hist, Lval_hist = [], []
for epoch in trange(epochs):
    Ltr,  _    = one_epoch(net, loss, tr_dl,  opt)
    Lval, Aval = one_epoch(net, loss, val_dl, None, accuracy)
    Ltr_hist.append(np.mean(Ltr))
    Lval_hist.append(np.mean(Lval))
    print(f'epoch: {epoch+1}\ttraining loss: {np.mean(Ltr):0.4f}\tvalidation loss: {np.mean(Lval):0.4f}\tvalidation accuracy: {np.mean(Aval):0.2f}')

return Ltr_hist, Lval_hist

class To3ch(object):
def __call__(self, pic):
    if pic.shape[0]==1: pic = pic.repeat(3,1,1)
    return pic

bs = 1
imagenet_stats = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])

transf = transforms.Compose([
    transforms.ToTensor(),
    To3ch(),
    transforms.Normalize(*imagenet_stats)
])

train_ds = datasets.Cityscapes('C:/cityscapes_ds', split='train', target_type='semantic', transform=transf, target_transform=transf)
val_ds = datasets.Cityscapes('C:/cityscapes_ds', split='val', target_type='semantic', transform=transf, target_transform=transf)

train_dl  = DataLoader(train_ds,  batch_size=bs,   shuffle=True,  num_workers=0)
val_dl = DataLoader(val_ds, batch_size=2*bs, shuffle=False, num_workers=0)

net = models.segmentation.deeplabv3_resnet50(num_classes=20)
fit(net.cuda(), train_dl, val_dl, loss=nn.CrossEntropyLoss(), epochs=1, lr=1e-4, wd=1e-4, plot=True)

Answer 1

您没有指定，但如果您使用的是原始 Cityscapes，则完全可以预料到此 OOM。

原始的 Cityscapes 数据集包含大图像（类似于 1024x2048，IIRC），看起来您有 6GB GPU。仅供参考，我无法将 batch_size=2 放入具有这种大小输入的 12GB GPU 中。

训练 DeepLab 模型时，对输入应用变换（例如，随机裁剪、调整大小、缩放等）很常见，但您似乎没有应用任何变换。

当你说：

I have a much more complicated model running, that will work with bs=16.

也许您正在研究另一种复杂性，它对内存需求的影响没有您想象的那么大。

使用 Pytorch 进行训练：由于 CUDA 内存问题导致错误

Training with Pytorch: error due to CUDA memory issue

pytorch

out-of-memory

torchvision