# 🌟 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)
