211024-Pytorch模型训练中显示进度条

# 🌟 1. 定义进度条
# 🌟 2. 设置迭代器
# 🌟 3. 设置开头
# 🌟 4. 设置结尾

# ⚠️：tqdm会在一定程度上降低训练的速度
# ⚠️：如果需要观察变化，需要使用time.sleep

from tqdm import tqdm
import time
    # ! with tqdm
    def fit(self, args, model, device, train_loader, optimizer, epoch, vis=True):
        # print('Epochs: ', epoch)
        correct = 0
        L = len(train_loader.dataset)
        model.train()
        with tqdm(train_loader, unit="batch") as tepoch: # 🌟 1. 定义进度条
            for data, target in tepoch:              # 🌟 2. 设置迭代器
                tepoch.set_description(f"Epoch {epoch}") # 🌟 3. 设置开头
                data, target = data.to(device), target.to(device)#.squeeze(1)       # Data to device
                optimizer.zero_grad()                                              # Zero gradient 
                output = model(data)                                               # Forward propagation
                losstr = F.nll_loss(output,target)                                 # Calculate loss
                losstr.backward()                                                  # Back propagation
                optimizer.step()                                                   # Optimize parameters
        
                predict = output.argmax(dim=1, keepdim=True)                           # Get the index 
                correct = predict.eq(target.view_as(predict)).sum().item() 
            
                model.global_step += 1
                # model.vis_moni_train(losstr.item(),correct,model.global_step) if vis==True else None # Visulize
                accuracy = correct/len(data)
                tepoch.set_postfix(loss=losstr.item(), accuracy='{:.3f}'.format(accuracy)) # 🌟 4. 设置结尾
                sleep(0.0001)

    # ! withour tqdm
    def fit(self, args, model, device, train_loader, optimizer, epoch, vis=True):
        print('Epochs: ', epoch)
        correct = 0
        L = len(train_loader.dataset)
        model.train()
        for data, target in tqdm(train_loader):              # Load from cpu
            data, target = data.to(device), target.to(device)#.squeeze(1)       # Data to device
            optimizer.zero_grad()                                              # Zero gradient 
            output = model(data)                                               # Forward propagation
            losstr = F.nll_loss(output,target)                                 # Calculate loss
            losstr.backward()                                                  # Back propagation
            optimizer.step()                                                   # Optimize parameters
    
            predict = output.argmax(dim=1, keepdim=True)                           # Get the index 
            correct = predict.eq(target.view_as(predict)).sum().item() 
        
            model.global_step += 1
            model.vis_moni_train(losstr.item(),correct,model.global_step) if vis==True else None # Visulize
            accuracy = correct/len(data)

• https://towardsdatascience.com/training-models-with-a-progress-a-bar-2b664de3e13e
• https://blog.csdn.net/dreaming_coder/article/details/113486645