PyTorch 训练中进度条不更新的解决方法

在 PyTorch 训练过程中，使用 tqdm 库来显示进度条非常方便。但是，有时候我们会遇到进度条不更新的情况，这通常是由于进度条的更新逻辑没有放在正确的循环内造成的。

以下代码展示了一个常见的错误示例：

for epoch in range(Epoch):
    d_epoch_loss = 0
    g_epoch_loss = 0
    count = len(dataloader)
    processBar = tqdm(dataloader,unit="step")
    for step, (img, _) in enumerate(dataloader):
        img = img.to(device)
        size = img.size(0)
        random_noise = torch.randn(size, 100, device=device)
        d_optim.zero_grad()
        # 判别器输入真实图片
        real_output = Dis(img)
        # 判别器在真实图像上的损失
        d_real_loss = loss_function(real_output,torch.ones_like(real_output))
        d_real_loss.backward()
        gen_img = Gen(random_noise)
        fake_output = Dis(gen_img.detach())  # 判别器输入生成图片,fake_output对生成图片的预测
        # gen_img是由生成器得来的，但我们现在只对判别器更新，所以要截断对Gen的更新
        # detach()得到了没有梯度的tensor，求导到这里就停止了，backward的时候就不会求导到Gen了
        d_fake_loss = loss_function(fake_output,torch.zeros_like(fake_output))
        d_fake_loss.backward()
        d_loss = d_real_loss + d_fake_loss
        d_optim.step()
        g_optim.zero_grad()
        fake_output = Dis(gen_img)
        g_loss = loss_function(fake_output,torch.ones_like(fake_output))
        g_loss.backward()
        g_optim.step()

        with torch.no_grad():
            d_epoch_loss += d_loss
            g_epoch_loss += g_loss

    with torch.no_grad():  # 之后的内容不进行梯度的计算（图的构建）
        d_epoch_loss /= count
        g_epoch_loss /= count
        D_loss.append(d_epoch_loss)
        G_loss.append(g_epoch_loss)
        print('Epoch:', epoch+1)
        gen_img_plot(Gen, test_input,epoch==0,(epoch==Epoch))
    processBar.close()

在这个例子中，processBar = tqdm(dataloader,unit="step") 放在了外层循环内，而进度条的更新应该是在内层循环，即每个 step 更新一次。

解决方案：

将进度条的更新放到内层循环中，并在初始化进度条时指定总步数。例如：

processBar = tqdm(total=len(dataloader)*Epoch, unit="step")
for epoch in range(Epoch):
    d_epoch_loss = 0
    g_epoch_loss = 0
    for step, (img, _) in enumerate(dataloader):
        ...
        processBar.update(1)
    ...
processBar.close()

这样就可以在每个 step 更新进度条，确保进度条的及时更新。

总结：

• 进度条的更新应该放在内层循环中，即每个 step 更新一次。
• 在初始化进度条时，需要指定总步数，以便进度条能够正确地显示进度。

通过以上方法，可以轻松解决 PyTorch 训练过程中进度条不更新的问题。