#训练过程函数 def fit(epoch,model,trainloader,testloader): correct = 0 total = 0 running_loss =0 model.train() #指明这是train模式需要bn和drop for x,y in trainloader: if torch.cuda.is_available(): x,y =x.to('cuda'),y.to('cuda') y_pred =model(x) loss = loss_fn(y_pred,y) optim.zero_grad() loss.backward() optim.step() with torch.no_grad(): y_pred = torch.argmax(y_pred,dim=1) correct +=(y_pred==y).sum().item() total += y.size(0) running_loss += loss.item() epoch_loss = running_loss/len(trainloader) epoch_acc =correct/total test_correct = 0 test_total = 0 test_running_loss =0 model.eval() with torch.no_grad(): for x,y in testloader: if torch.cuda.is_available(): x,y = x.to('cuda'),y.to('cuda') y_pred =model(x) loss = loss_fn(y_pred,y) y_pred = torch.argmax(y_pred,dim=1) test_correct +=(y_pred==y).sum().item() test_total +=y.size(0) test_running_loss +=loss.item() epoch_tst_loss =test_running_loss/len(testloader) epoch_tst_acc = test_correct/test_total print('epoch',epoch,'loss',round(epoch_loss,3), 'acc:',round(epoch_acc,3), 'test_loss:',round(epoch_tst_loss,3), 'test_acc:',round(epoch_tst_acc,3)) return epoch_loss ,epoch_acc,epoch_tst_loss,epoch_tst_acc #打印训练epoch epochs =30 train_loss =[] train_acc =[] test_loss =[] test_acc=[] for epoch in range(epochs): epoch_loss,epoch_acc,epoch_tst_loss,epoch_tst_acc =fit(epoch,model,train_dl,test_dl) train_loss.append(epoch_loss) train_acc.append(epoch_acc) test_loss.append(epoch_tst_loss) test_acc.append(epoch_tst_acc) #绘制图像: plt.plot(range(1,epochs+1),train_loss,label='train_loss') plt.plot(range(1,epochs+1),test_loss,label = 'test_loss' ) plt.legend() plt.show()
写完模型之后的训练函数
猜你喜欢
转载自blog.csdn.net/qq_45675231/article/details/129899983
今日推荐
周排行