当前位置:   article > 正文

tqdm在pytorch中的使用_tqdm(enumerate(train_loader), total=train_loader._

tqdm(enumerate(train_loader), total=train_loader.__len__())

tqdm库的主要作用是可视化当前网络训练的进程,下载方式为在终端中输入:

pip install tqdm

考虑下面一个常见的训练过程:

  1. for epoch in range(num_epochs):
  2. ## 对训练数据的迭代器进行迭代计算
  3. for step, (b_x1, b_y) in enumerate(train_loader):
  4. b_x1 = b_x1.cuda()
  5. b_y = b_y.cuda()
  6. output = mlp1(b_x1).flatten() # MLP在训练batch上的输出
  7. output = output.cuda()
  8. train_loss = loss_func(output, b_y) # 平方根误差
  9. optimizer.zero_grad() # 每个迭代步的梯度初始化为0
  10. train_loss.backward() # 损失的后向传播,计算梯度
  11. optimizer.step() # 使用梯度进行优化
  12. train_loss_one += train_loss.item() * b_x1.size(0)
  13. train_num += b_x1.size(0)
  14. train_loss_all.append(train_loss.item())
  15. train_loss_all.append(train_loss_one / train_num)

我们希望在训练过程能看到训练的进度,例如当前的epoch是多少,loss是多少等等,可以按照下方的方式进行改进:

  1. for epoch in range(num_epochs):
  2. loop = tqdm(enumerate(train_loader), total=len(train_loader))
  3. ## 对训练数据的迭代器进行迭代计算
  4. for step, (b_x1, b_y) in loop:
  5. b_x1 = b_x1.cuda()
  6. b_y = b_y.cuda()
  7. output = mlp1(b_x1).flatten() # MLP在训练batch上的输出
  8. output = output.cuda()
  9. train_loss = loss_func(output, b_y) # 平方根误差
  10. optimizer.zero_grad() # 每个迭代步的梯度初始化为0
  11. train_loss.backward() # 损失的后向传播,计算梯度
  12. optimizer.step() # 使用梯度进行优化
  13. train_loss_one += train_loss.item() * b_x1.size(0)
  14. train_num += b_x1.size(0)
  15. train_loss_all.append(train_loss.item())
  16. loop.set_description(f'Epoch [{epoch}/{num_epochs}]')
  17. loop.set_postfix(loss=train_loss_one / train_num)

首先是创建一个loop,目的是在后面for循环时能通过tqdm去调用,再通过set_description与set_postfix设置其他信息。

效果:

 

声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/盐析白兔/article/detail/101783
推荐阅读
相关标签
  

闽ICP备14008679号