Pytorch完整训练模型套路
-
以CIFAR_10模型为例
-
使用tensorboard将训练过程可视化
-
模型文件:
model.py
import torch from torch import nn # 搭建神经网络 class CIFAR_10_Model(nn.Module): def __init__(self): super().__init__() self.model=nn.Sequential( nn.Conv2d(3,32,5,1,2), nn.MaxPool2d(2), nn.Conv2d(32,16,5,1,2), nn.MaxPool2d(2), nn.Flatten(), nn.Linear(64*4*4,64), nn.Linear(64,10) ) def forward(self,input): return self.model(input) if __name__=="__main__": cifar10=CIFAR_10_Model() input=torch.ones((64,3,32,32)) output=cifar10(input) print(output.shape)
-
训练文件
train.py
import torchvision from torch.utils.tensorboard import SummaryWriter from model import * from torch import nn from torch.utils.data import DataLoader # 准备数据集 train_data=torchvision.datasets.CIFAR10(root="../data/dataset",train=True, transform=torchvision.transforms.ToTensor(), download=True) test_data=torchvision.datasets.CIFAR10(root="../data/dataset",train=False, transform=torchvision.transforms.ToTensor(), download=True) # 数据集的长度 train_data_size=len(train_data) test_data_size=len(test_data) print(f"train_set长度:{train_data_size}") print(f"test_set长度:{test_data_size}") # 利用DataLoader加载数据 train_dataloader=DataLoader(train_data,batch_size=64) test_dataloader=DataLoader(test_data,batch_size=64) # 创建CIFAR_10模型实例 Cifar10=CIFAR_10_Model() #损失函数 loss_fn=nn.CrossEntropyLoss() #学习率 learning_rate=1e-2 #优化器 optimizer=torch.optim.SGD(Cifar10.parameters(),lr=learning_rate) #训练次数 total_train_step=0 #测试次数 total_test_step=0 #训练轮数 epoch=10 # 添加tensorboard writer=SummaryWriter("./dataset_logs") for i in range(epoch): print(f"----------第{i+1}轮训练开始----------") #开始训练 Cifar10.train() for data in train_dataloader: imgs,targets=data outputs=Cifar10(imgs) loss=loss_fn(outputs,targets) #反向传播,梯度下降 optimizer.zero_grad() loss.backward() optimizer.step() total_train_step=total_test_step+1 if (total_train_step+1)%100==0: print(f"训练次数:{total_train_step},Loss:{loss}") writer.add_scalar("train_loss",loss.item(),total_train_step) #开始测试 Cifar10.eval() total_test_loss=0 total_accuray=0 with torch.no_grad(): for data in test_dataloader: imgs,targets=data outputs=Cifar10(imgs) loss=loss_fn(outputs,targets) total_test_loss=total_test_loss+loss.item() accuracy=(outputs.argmax(1)==targets).sum() total_accuray=total_accuray+accuracy print(f"整体测试集的Loss:{total_test_loss}") print(f"整体测试集的正确率:{total_accuray/test_data_size}") writer.add_scalar("test_loss",total_test_loss,total_test_step) writer.add_scalar("test_accuracy",total_accuray/test_data_size,total_test_step) total_test_step=total_test_step+1 writer.close()