python - 在深度学习训练时配置进度条

Question

我从教程中升级了这个小型训练函数。

def train(epoch, tokenizer, model, device, loader, optimizer):
model.train()
with tqdm.tqdm(loader, unit="batch") as tepoch:
  for _,data in enumerate(loader, 0):
      y = data['target_ids'].to(device, dtype = torch.long)
      y_ids = y[:, :-1].contiguous()
      lm_labels = y[:, 1:].clone().detach()
      lm_labels[y[:, 1:] == tokenizer.pad_token_id] = -100
      ids = data['source_ids'].to(device, dtype = torch.long)
      mask = data['source_mask'].to(device, dtype = torch.long)

      outputs = model(input_ids = ids, attention_mask = mask, decoder_input_ids=y_ids, labels=lm_labels)
      loss = outputs[0]

      tepoch.set_description(f"Epoch {epoch}")
      tepoch.set_postfix(loss=loss.item())
      
      if _%10 == 0:
          wandb.log({"Training Loss": loss.item()})

      if _%1000==0:
          print(f'Epoch: {epoch}, Loss:  {loss.item()}')
  
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()
      # xm.optimizer_step(optimizer)
      # xm.mark_step()

该函数训练得很好，问题是我似乎无法使进度条正常工作。我尝试过它，但还没有找到可以正确更新损失并告诉我还剩多少时间的配置。 有人对我可能做错了什么有任何指示吗？ 提前致谢!

Answer 1

如果其他人遇到了我同样的问题，感谢之前的回复，我只需对之前所做的事情进行一些调整即可根据需要配置进度条:

def train(epoch, tokenizer, model, device, loader, optimizer):
  model.train()    
  for _,data in tqdm(enumerate(loader, 0), unit="batch", total=len(loader)):

一切都保持不变，现在我有一个进度条显示百分比和损失。我更喜欢这个解决方案，因为它允许我保留我拥有的其他日志记录功能而无需进一步更改。