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

Question

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

数据集正在工作，数据加载器也是如此。当我尝试训练时，我总是会收到此错误，此时第一批尝试通过网络（在 one_epoch 函数中的 y_ = net(xb)）。

RuntimeError: CUDA 内存不足。尝试分配 128.00 MiB（GPU 0；6.00 GiB 总容量；4.20 GiB 已分配；6.87 MiB 空闲；PyTorch 总共保留 4.20 GiB）

奇怪的是，无论批处理大小（bs）是多少，根据错误释放的内存量都是比尝试分配的内存量略小的值，例如对于 bs=16 我得到：

RuntimeError: 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 模型时，通常会在输入上应用转换（例如，随机裁剪、调整大小、缩放等），看起来您没有应用任何转换。

当你说：

我正在运行一个更复杂的模型，它适用于 bs=16。

也许您正在研究一种不同类型的复杂性，它对内存需求的影响比您想象的要小。