pytorch - 验证阶段完成后gpu内存仍然被占用，pytorch

Question

据我所知，在使用 GPU 训练和验证模型时，GPU 内存主要用于加载数据，向前和向后。据我所知，我认为 GPU 内存使用应该相同 1) 训练前，2) 训练后，3) 验证前，4) 验证后。但在我的例子中，验证阶段使用的 GPU 内存在训练阶段仍然被占用，反之亦然。它不会在每个时期都增加，所以我确信这不是像 loss.item() 这样的常见错误。

这是我的问题的总结

• 是否应该先清理一个阶段使用的 GPU 内存(模型权重除外)？
• 如果应该的话，我在这里是否犯了任何新手错误......？

感谢您的帮助。

这是训练循环的代码

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)

我读过关于将验证阶段设为函数有何帮助，因为 python 是函数作用域语言。
所以 evaluate() 是

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)

这是 colab 中的 GPU 使用情况。 1 是第一次调用 evaluate() 的点，2 是火车开始的时间。

Answer 1

分配 GPU 内存很慢。 PyTorch 保留它分配的 GPU 内存，即使没有更多张量引用该内存。您可以调用 torch.cuda.empty_cache() 来释放任何不可访问的 GPU 内存。