TQDM不会更新时期

Question

    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_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()

这是我的代码。

这是打印到控制台的示例：

重点是，顶部计数器（与时期1平行）仅从1/10000升级到10/10000始终停止，即使时代大于10。

Answer 1

您可以在每个时期都可以做类似的事情，

for epoch in range(10):
     trainBar = tqdm.tqdm(trainData)
     valBar = tqdm.tqdm(valData)
     for batch, data in enumerate(trainBar):
       do smthing
     for batch, data in enumerate(valBar):
       do smthing

每个时期都会创建一个新的TQDM栏，而您不必担心重置它。