PyTorch训练循环加入评估模块后停滞，进度条不更新

问题 #

用PyTorch构建了一个简单的神经网络，运行下述训练循环时，程序卡在首次迭代，tqdm_notebook进度条无更新。移除整个评估模块后，循环与进度条运行正常；在granular_MLP.eval()前添加print语句，循环也能正常执行。请问该现象的原因是什么？

train_loss, train_auc, train_auprc, test_loss, test_auc, test_auprc = [], [], [], [], [], []

for epoch in tqdm_notebook(range(100)): 
    batch_train_loss, batch_train_auc, batch_train_auprc = [], [], []
    granular_MLP.train()
    
    for i, (X,y) in enumerate(granular_train_data_loader):
        X = X.cuda('cuda:0')
        y = y.cuda('cuda:0')
        output = granular_MLP(X.view(len(X), granular_n_input))
        loss = granular_criterion(output, y.view(len(y), 1).float())
        
        batch_train_loss.append(loss.item())
        batch_train_auc.append(roc_auc_score(y.detach().cpu(), output.detach().cpu()))
        batch_train_auprc.append(average_precision_score(y.detach().cpu(), output.detach().cpu()))
        
        granular_optimizer.zero_grad()
        loss.backward()
        granular_optimizer.step()
        
    train_loss.append(np.mean(batch_train_loss))
    train_auc.append(np.mean(batch_train_auc))
    train_auprc.append(np.mean(batch_train_auprc))
    
    # model evaluation 
    granular_MLP.eval()
    batch_test_loss, batch_test_auc, batch_test_auprc = [], [], []
    for i, (X,y) in enumerate(granular_test_data_loader):
        X = X.cuda('cuda:0')
        y = y.cuda('cuda:0')
        output = granular_MLP(X.view(len(X), granular_n_input))
        loss = granular_criterion(output, y.view(len(y), 1).float())

        batch_test_loss.append(loss.item())
        batch_test_auc.append(roc_auc_score(y.detach().cpu(), output.detach().cpu()))
        batch_test_auprc.append(average_precision_score(y.detach().cpu(), output.detach().cpu()))
        
    test_loss.append(np.mean(batch_test_loss))
    test_auc.append(np.mean(batch_test_auc))
    test_auprc.append(np.mean(batch_test_auprc))

原因分析与解决办法 #

核心原因 #

问题出在评估阶段未使用torch.no_grad()上下文管理器，导致评估时仍在计算并存储梯度，显存占用急剧上升，最终引发CUDA操作阻塞，程序停滞。添加print语句时，IO操作会触发CUDA同步，间接释放部分临时显存，让程序能继续执行，但这只是临时 workaround；移除评估模块则直接避免了额外的显存消耗，所以循环能正常运行。

关键细节 #

• PyTorch的eval()模式仅关闭批量归一化（BatchNorm）、Dropout等训练特有的行为，不会自动禁用梯度计算，前向传播依然会构建计算图并保留梯度信息。
• 当测试集数据量大时，持续的梯度存储会快速耗尽显存，导致CUDA kernel无法继续执行，表现为进度条无更新、程序卡住。

修复后的代码 #

在评估循环外层添加torch.no_grad()，彻底关闭梯度计算：

train_loss, train_auc, train_auprc, test_loss, test_auc, test_auprc = [], [], [], [], [], []

for epoch in tqdm_notebook(range(100)): 
    batch_train_loss, batch_train_auc, batch_train_auprc = [], [], []
    granular_MLP.train()
    
    for i, (X,y) in enumerate(granular_train_data_loader):
        X = X.cuda('cuda:0')
        y = y.cuda('cuda:0')
        output = granular_MLP(X.view(len(X), granular_n_input))
        loss = granular_criterion(output, y.view(len(y), 1).float())
        
        batch_train_loss.append(loss.item())
        batch_train_auc.append(roc_auc_score(y.detach().cpu(), output.detach().cpu()))
        batch_train_auprc.append(average_precision_score(y.detach().cpu(), output.detach().cpu()))
        
        granular_optimizer.zero_grad()
        loss.backward()
        granular_optimizer.step()
        
    train_loss.append(np.mean(batch_train_loss))
    train_auc.append(np.mean(batch_train_auc))
    train_auprc.append(np.mean(batch_train_auprc))
    
    # model evaluation 
    granular_MLP.eval()
    batch_test_loss, batch_test_auc, batch_test_auprc = [], [], []
    # 禁用梯度计算，避免显存占用过高
    with torch.no_grad():
        for i, (X,y) in enumerate(granular_test_data_loader):
            X = X.cuda('cuda:0')
            y = y.cuda('cuda:0')
            output = granular_MLP(X.view(len(X), granular_n_input))
            loss = granular_criterion(output, y.view(len(y), 1).float())

            batch_test_loss.append(loss.item())
            batch_test_auc.append(roc_auc_score(y.detach().cpu(), output.detach().cpu()))
            batch_test_auprc.append(average_precision_score(y.detach().cpu(), output.detach().cpu()))
        
    test_loss.append(np.mean(batch_test_loss))
    test_auc.append(np.mean(batch_test_auc))
    test_auprc.append(np.mean(batch_test_auprc))

内容的提问来源于stack exchange，提问作者Kaihua Hou