import torch import torchvision.datasets from torch import nn from torch.utils.data import DataLoader from torch.utils.tensorboard import SummaryWriter from Model import * # 准备数据集 train_data=torchvision.datasets.CIFAR10(root="cifar-10-python",train=True,transform=torchvision.transforms.ToTensor(), download=True) test_data=torchvision.datasets.CIFAR10(root="cifar-10-python",train=False,transform=torchvision.transforms.ToTensor(), download=True) # 看看数据集长度 train_data_size=len(train_data) test_data_size=len(test_data) print("训练数据集的长度为:{}".format(train_data_size)) print("测试数据集的长度为:{}".format(test_data_size)) #利用dataloader加载数据集 train_dataloader=DataLoader(train_data,batch_size=64) test_dataloader=DataLoader(test_data,batch_size=64) # 创建网络模型 model=Model() # 创建损失函数 loss_fn=nn.CrossEntropyLoss() #优化器,随机梯度下降 learning_rate=0.01 optimizer=torch.optim.SGD(model.parameters(),lr=learning_rate) # 设置训练网络的一些参数 #记录训练和测试的次数 total_train_step=0 total_test_step=0 #训练轮数 epoch=10 #添加tensorboard writer=SummaryWriter("logs_train") for i in range(epoch): print("----------第{}轮训练开始----------".format(i+1)) #训练步骤开始 for data in train_dataloader: imgs,targets=data outputs=model(imgs) loss=loss_fn(outputs,targets) #优化 optimizer.zero_grad() loss.backward() optimizer.step() total_train_step+=1 if(total_train_step%100==0): print("训练次数:{},loss:{}".format(total_train_step,loss.item())) writer.add_scalar("train_loss",loss.item(),total_train_step) #测试步骤开始 total_test_loss=0 total_accuracy=0 with torch.no_grad(): for data in test_dataloader: imgs,targets=data outputs=model(imgs) loss=loss_fn(outputs,targets) total_test_loss=total_test_loss+loss.item() accuracy=(outputs.argmax(1)==targets).sum() total_accuracy+=accuracy print("整体测试集上的loss:{}".format(total_test_loss)) print("整体测试集上的正确率:{}".format(total_accuracy/test_data_size)) writer.add_scalar("test_loss",total_test_loss,total_test_step) writer.add_scalar("test_accuracy",total_accuracy/test_data_size,total_test_step) total_test_step+=1 torch.save(model,"model_{}.pth".format(i)) #torch,save(model.state_dict(),"model_{}.pth".format(i)) print("模型已保存") writer.close()
正确率:用于分类问题
图源:一只小土堆 bilibili
当从网络上下载模型,如果是cuda的可以用
model=torch.load("文件名"),map_location=torch.device('cpu')
来切换到cpu运行
训练结束后,测试
import torch import torchvision from PIL import Image from torch import nn image_pth="img.png"#这里是输入图片 image=Image.open(image_pth) image=image.convert('RGB')#保留颜色通道 transform=torchvision.transforms.Compose([torchvision.transforms.Resize((32,32)), torchvision.transforms.ToTensor()]) image=transform(image) print(image.shape) class Model(nn.Module): def __init__(self): super(Model,self).__init__() self.model=nn.Sequential( nn.Conv2d(3,32,5,1,padding=2), nn.MaxPool2d(2), nn.Conv2d(32,32,5,1,padding=2), nn.MaxPool2d(2), nn.Conv2d(32,64,5,1,padding=2), nn.MaxPool2d(2), nn.Flatten(), nn.Linear(64*4*4,64), nn.Linear(64,10) ) def forward(self,x): x=self.model(x) return x model=torch.load("model_9.pth") print(model) image=torch.reshape(image,(1,3,32,32)) model.eval() with torch.no_grad(): output=model(image) print(output) print(output.argmax(1))#这里输出tensor数据类型第几类,1表示横向看