PyTorch验证阶段结束后GPU内存仍被占用的问题排查

据我所知，用GPU训练和验证模型时，GPU内存主要用于加载数据、前向传播与反向传播。我认为以下节点的GPU内存占用应该一致：

• 训练前
• 训练后
• 验证前
• 验证后

但实际场景中，验证阶段占用的GPU内存会保留到训练阶段，反之亦然，且内存占用不会随epoch增加，已经排除了loss.item()这类常见错误。

我的问题：

1. 除模型权重外，一个阶段的GPU内存难道不该在进入另一阶段前被清理吗？
2. 若应该清理，我是不是犯了新手常犯的错误？

训练循环代码 #

eval_result = evaluate(model,val_loader,True,True)
print(eval_result)

print('start training')
for epoch in range(num_epoch):
    model.train()
    time_ = datetime.datetime.now()
    for iter_, data in enumerate(tr_loader):
        x, y = data
        x = x.to(device).view(x.shape[0],1,*(x.shape[1:]))
        y = y.to(device).long()
        
        pred = model.forward(x)
        loss = loss_fn(pred,y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        # print
        print_iter = 16
        if (iter_+1) % print_iter == 0:
            elapsed = datetime.datetime.now() - time_
            expected = elapsed * (num_batches / print_iter)
            _epoch = epoch + ((iter_ + 1) / num_batches)
            print('\rTRAIN [{:.3f}/{}] loss({}) '
                'elapsed {} expected per epoch {}'.format(
                    _epoch,num_epoch, loss.item(), elapsed, expected)
                 ,end="\t\t\t")
            time_ = datetime.datetime.now()
        
        
        
    print()
    eval_result = evaluate(model,val_loader,True,True)
    print(eval_result)
    scheduler.step(eval_result[0])
    
    if (epoch+1) %1 == 0:
        save_model(model, optimizer, scheduler)

验证函数代码 #

def evaluate(model, val_loader, get_acc = True, get_IOU = True):
    """
    pred: Tensor of shape B C D H W
    label Tensor of shape B D H W
    """
    val_loss = 0
    val_acc = 0
    val_IOU = 0
    with torch.no_grad():
        model.eval()
        for data in tqdm(val_loader):
            x, y = data
            x = x.to(device).view(x.shape[0],1,*(x.shape[1:]))
            y = y.to(device).long()

            pred = model.forward(x)
                    
            loss = loss_fn(pred,y)
            val_loss += loss.item()
        
            pred = torch.argmax(pred, dim=1)
            
            if get_acc:
                total = np.prod(y.shape)
                total = total if total != 0 else 1
                val_acc += torch.sum((pred == y)).cpu().item()/total
            if get_IOU:
                iou = 0
                for class_num in range(1,8):
                    iou += torch.sum((pred==class_num)&(y==class_num)).cpu().item()\
                        / torch.sum((pred==class_num)|(y==class_num)).cpu().item()
                val_IOU += iou/7
                
        val_loss /= len(val_loader)
        val_acc /= len(val_loader)
        val_IOU /= len(val_loader)
    return (val_loss, val_acc, val_IOU)

GPU使用情况说明 #

Colab的GPU使用曲线显示：

• 节点1：首次调用evaluate()的时刻，GPU内存占用上升
• 节点2：训练开始的时刻，GPU内存没有下降，而是在之前的基础上继续保持占用

内容的提问来源于stack exchange，提问作者Luana Kwon