tqdm库的主要作用是可视化当前网络训练的进程,下载方式为在终端中输入:
pip install tqdm
考虑下面一个常见的训练过程:
for epoch in range(num_epochs):
## 对训练数据的迭代器进行迭代计算
for step, (b_x1, b_y) in enumerate(train_loader):
b_x1 = b_x1.cuda()
b_y = b_y.cuda()
output = mlp1(b_x1).flatten() # MLP在训练batch上的输出
output = output.cuda()
train_loss = loss_func(output, b_y) # 平方根误差
optimizer.zero_grad() # 每个迭代步的梯度初始化为0
train_loss.backward() # 损失的后向传播,计算梯度
optimizer.step() # 使用梯度进行优化
train_loss_one += train_loss.item() * b_x1.size(0)
train_num += b_x1.size(0)
train_loss_all.append(train_loss.item())
train_loss_all.append(train_loss_one / train_num)
我们希望在训练过程能看到训练的进度,例如当前的epoch是多少,loss是多少等等,可以按照下方的方式进行改进:
for epoch in range(num_epochs):
loop = tqdm(enumerate(train_loader), total=len(train_loader))
## 对训练数据的迭代器进行迭代计算
for step, (b_x1, b_y) in loop:
b_x1 = b_x1.cuda()
b_y = b_y.cuda()
output = mlp1(b_x1).flatten() # MLP在训练batch上的输出
output = output.cuda()
train_loss = loss_func(output, b_y) # 平方根误差
optimizer.zero_grad() # 每个迭代步的梯度初始化为0
train_loss.backward() # 损失的后向传播,计算梯度
optimizer.step() # 使用梯度进行优化
train_loss_one += train_loss.item() * b_x1.size(0)
train_num += b_x1.size(0)
train_loss_all.append(train_loss.item())
loop.set_description(f'Epoch [{epoch}/{num_epochs}]')
loop.set_postfix(loss=train_loss_one / train_num)
首先是创建一个loop,目的是在后面for循环时能通过tqdm去调用,再通过set_description与set_postfix设置其他信息。
效果:
?
|