PyTorch学习记录十二——tqdm

tqdm

[tqdm三方库官方文档]https://tqdm.github.io/

• tqdm是一个第三方库，可以在python中实现进度条功能。
• 可以在长循环中添加一个进度提示信息，用户只需要封装任意的迭代器tqdm(iterator)。
• 可以帮助我们检测程序运行的进度，估计运行的时长。

一、基本使用方法

1.传入可迭代对象

import time
from tqdm import *  # *代表导入所有可导入对象到当前命名空间（包括类、函数、变量等）
for i in tqdm(range(100)):
    time.sleep(0.1)

输出结果如下(共用时10s)

使用trange

trange(i)是tqdm(range(i))的简单写法

import time
from tqdm import *  
for i in trange(100):
    time.sleep(0.1)

2.为进度条设置描述

在for循环外部初始化tqdm，可以打印其他相关信息。
通过set_description和set_postfix方法设置进度条显示信息。

import time
from tqdm import *
from random import random,  randint
pbar = tqdm(["a", "b", "c", "d"])
for char in pbar:
    # 设置进度条左边显示的信息
    pbar.set_description("Processing %s" % char)
    # 设置进度条右边显示的信息
    pbar.set_postfix(loss=random(), gen=randint(1, 999), str="h", lst=[1, 2])
    time.sleep(1)

输出如下

3.write放方法

import time
from tqdm import *
bar = trange(10)
for i in bar:
    time.sleep(1)
    if not (i % 3):
        tqdm.write("Done task %i" % i)

输出如下

4.手动控制进度条

通过update方法可以控制每次进度条更新的进度：

from tqdm import tqdm
import time
# total参数设置进度条的总长度
pbar = tqdm(total=100)
for i in range(30):
    time.sleep(0.5)
    # 每次更新进度条的长度
    pbar.update(2)
# 关闭占用的资源
pbar.close()

输出结果如下

二、在DL中的使用

[教程]https://blog.youkuaiyun.com/wxd1233/article/details/118371404

• 以手写数字识别代码作为例子。
• 对于class tqdm(Comparable)，装饰的是一个可迭代对象，返回一个迭代器，该迭代器的行为与原始可迭代对象完全相同，但每次请求值时都会打印一个动态更新的进度条。
• 对于在模型训练中来说，train_loader是PyTorch中的DataLoader对象（torch.utils.data.DataLoader类的一个实例），用于批量加载训练数据。我们将train_loader作为参数传递给tqdm函数，tqdm函数会将train_loader包装成一个新的可迭代对象loop，并且在迭代过程中显示进度条。
• 至于train_loader本身其实也是一个可迭代对象，所以我们可以用for循环直接遍历，只不过被包装成了一个新的。

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from tqdm import tqdm


class CNN(nn.Module):
    def __init__(self, in_channels=1, num_classes=10):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.pool = nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2))
        self.conv2 = nn.Conv2d(in_channels=8, out_channels=16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.fc = nn.Linear(16*7*7, num_classes)

    def forward(self, x):
        x = F.relu(self.conv1(x))  # 激活函数也可以在forward中来实现
        x = self.pool(x)
        x = F.relu(self.conv2(x))
        x = self.pool(x)
        # 将图像数据四维张量展平为二维张量（batch_size, channels*height*weight）。确保展平后每个样本仍然独立。
        x = x.reshape(x.shape[0], -1)
        # 将特征图输入全连接层之前，需要将其展平为二维张量。
        # 全连接层（nn.Linear）虽然通常被描述为接收一维输入，但实际上它可以处理二维输入，并且在实际应用中，处理二维输入是非常常见的。
        # 全连接层会对每个样本独立地应用相同的线性变换，即对每一行（代表一个样本）进行操作，而不会影响批量维度，可以处理批量数据。
        x = self.fc(x)
        return x


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

in_channels = 1
num_classes = 10
learning_rate = 0.001
batch_size = 64
num_epochs = 5

train_dataset = datasets.MNIST(root="dataset/", train=True, transform=transforms.ToTensor(), download=True)
test_dataset = datasets.MNIST(root="dataset/", train=False, transform=transforms.ToTensor(), download=True)
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=True)

model = CNN().to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

for epoch in range(num_epochs):
    losses = []
    accuracy = []
    loop = tqdm(train_loader, total=len(train_loader))
    for data, targets in loop:
        data = data.to(device=device)
        targets = targets.to(device=device)
        scores = model(data)  # 形状为（batch_size, num_classes）,对每一个样本是每一个类别都有一个得分。
        loss = criterion(scores, targets)
        losses.append(loss)
        optimizer.zero_grad()  # 在反向传播开始之前，要将上一批次反向传播保存的梯度信息（.grad）清除，否则会导致两批次的梯度累加。
        loss.backward()  # 本次反向传播用的梯度信息都是在backward过程中计算的，与上一批次的无关，也无关前相传播（只提供计算图，即函数关系）
        _, predictions = scores.max(1)  # 忽略最大值，只需要最大值索引。
        num_correct = (predictions == targets).sum()
        running_train_acc = float(num_correct) / float(data.shape[0])
        accuracy.append(running_train_acc)
        optimizer.step()  # 使用优化器，根据反向传播得到的梯度更新模型的参数
        # 当遇到{}时，Python会把其中的内容当作一个表达式进行求值，然后将求值结果转换为字符串，并替换掉{}及其内部的内容。
        # 直接将变量括起来，就实现了一个占位符的作用，会被当前循环中变量的真实值替换掉。
        loop.set_description(f'Epoch [{epoch + 1}/{num_epochs}]')  # 添加tqdm左侧信息
        # criterion返回的是一个标量形式的张量（只包含一个元素的张量，0阶张量），使用.item()将其转换为python标量
        loop.set_postfix(loss=loss.item(), acc=running_train_acc)  # 添加tqdm右侧信息

输出结果如下

三、ViT中的使用