TQDM训练进度条epoch计数器更新至10/10000后停止问题

问题表现 #

• 扩散模型训练代码运行时，tqdm生成的、与Epoch标识同栏的顶部进度计数器，仅从1/10000更新到10/10000后就完全停滞，实际训练epoch轮次超过10后计数器也不会继续推进。
• 控制台输出效果参考：
  

原训练代码如下：

for epoch in range(args.num_epochs):
    model.train()
    # print(f"Epoch {epoch}")
    with tqdm(total=len(input_tensor_catted), unit="ba") as pbar:
        pbar.set_description(f"Epoch {epoch}")
        pbar.update(1)
        # for step, batch in enumerate(train_dataloader):
        for step in range(len(input_tensor_catted) // args.batch_size):

            indices = torch.multinomial(torch.ones(len(input_tensor_catted)) / len(input_tensor_catted), args.batch_size, replacement=True)

            clean_inputs = input_tensor_catted[indices, :]
            clean_conditioning = original_cost_tensor_catted[indices, :].to(clean_inputs.device)

            # clean clean_inputs = batch["input"]
            noise_samples = torch.randn(clean_inputs.shape).to(clean_inputs.device)
            bsz = clean_inputs.shape[0]

            timesteps = torch.randint(0, noise_scheduler.timesteps, (bsz,), device=clean_inputs.device).long()

            # add noise onto the clean images according to the noise magnitude at each timestep
            # (this is the forward diffusion process)
            noisy_images = noise_scheduler.training_step(clean_inputs, noise_samples, timesteps)

            if step % args.gradient_accumulation_steps != 0:
                with accelerator.no_sync(model):
                    # from noisy images, predict epsilon
                    output = model(noisy_images, timesteps, clean_conditioning)
                    # predict the noise residual
                    loss = F.mse_loss(output, noise_samples)
                    loss = loss / args.gradient_accumulation_steps
                    accelerator.backward(loss)
            else:
                output = model(noisy_images, timesteps, clean_conditioning)

                # predict the noise residual
                loss = F.mse_loss(output, noise_samples)
                loss = loss / args.gradient_accumulation_steps
                accelerator.backward(loss)
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()
            pbar.update(1)
            pbar.set_postfix(loss=loss.detach().item(), lr=optimizer.param_groups[0]["lr"])

            optimizer.step()

故障根因 #

代码存在3处逻辑错误，直接导致进度条异常，同时会破坏训练正确性：

1. 无效初始进度更新：进入tqdm上下文管理器后、未执行任何训练步时就调用pbar.update(1)，平白多计数1个单位进度。
2. 进度条总长度与实际更新次数不匹配：tqdm初始化时设置total=len(input_tensor_catted)（即总样本数，对应你看到的分母10000），但内层训练步的循环次数为len(input_tensor_catted) // args.batch_size，单轮epoch内仅会调用对应次数的pbar.update(1)。若batch_size为1000，单轮epoch仅会更新10次进度，刚好对应计数到10就停止的现象——内层循环结束后当前epoch的pbar生命周期就已结束，进入下一轮epoch会创建新的pbar实例，旧进度条自然不会再更新。
3. 重复优化器步进：训练循环末尾多余的optimizer.step()会和梯度累积分支的更新逻辑冲突，导致梯度更新规则完全错乱，训练无法正常收敛。

修复代码 #

针对以上问题修改后可正常运行：

for epoch in range(args.num_epochs):
    model.train()
    # 对齐tqdm总长度与单epoch实际训练步数
    total_steps_per_epoch = len(input_tensor_catted) // args.batch_size
    with tqdm(total=total_steps_per_epoch, unit="ba") as pbar:
        pbar.set_description(f"Epoch {epoch}")
        # 删除无效的初始update调用
        for step in range(total_steps_per_epoch):
            indices = torch.multinomial(torch.ones(len(input_tensor_catted)) / len(input_tensor_catted), args.batch_size, replacement=True)
            clean_inputs = input_tensor_catted[indices, :]
            clean_conditioning = original_cost_tensor_catted[indices, :].to(clean_inputs.device)

            noise_samples = torch.randn(clean_inputs.shape).to(clean_inputs.device)
            bsz = clean_inputs.shape[0]
            timesteps = torch.randint(0, noise_scheduler.timesteps, (bsz,), device=clean_inputs.device).long()
            noisy_images = noise_scheduler.training_step(clean_inputs, noise_samples, timesteps)

            if step % args.gradient_accumulation_steps != 0:
                with accelerator.no_sync(model):
                    output = model(noisy_images, timesteps, clean_conditioning)
                    loss = F.mse_loss(output, noise_samples)
                    loss = loss / args.gradient_accumulation_steps
                    accelerator.backward(loss)
            else:
                output = model(noisy_images, timesteps, clean_conditioning)
                loss = F.mse_loss(output, noise_samples)
                loss = loss / args.gradient_accumulation_steps
                accelerator.backward(loss)
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()
            
            pbar.update(1)
            pbar.set_postfix(loss=loss.detach().item(), lr=optimizer.param_groups[0]["lr"])
        # 删除循环末尾多余的optimizer.step()调用

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