一、模型结构
import torch
from torch import nn
class Cifar10Model(nn.Module):
def __init__(self):
super(Cifar10Model, self).__init__()
self.model=nn.Sequential(
nn.Conv2d(3,32,5,1,2),
nn.MaxPool2d(2),
nn.Conv2d(32,32,5,1,2),
nn.MaxPool2d(2),
nn.Conv2d(32,64,5,1,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
if __name__=='__main__':
tudui=Cifar10Model()
input=torch.ones(64,3,32,32)
output=tudui(input)
print(output.shape)
训练和测试:
import torch.optim
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from Model.Cifar10Model import Cifar10Model
train_data= torchvision.datasets.CIFAR10(root="../dataset",train=True,transform=torchvision.transforms.ToTensor(),download=True)
test_data= torchvision.datasets.CIFAR10(root="../dataset",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)
#创建网络模型
cifar10Model=Cifar10Model()
#损失函数
loss_fn=nn.CrossEntropyLoss()
#优化器
#1e-2=1*(10)^(-2)=1/100=0.01
learning_rate=1e-2
optimzer=torch.optim.SGD(cifar10Model.parameters(),lr=learning_rate)
#设置训练网络的一些参数
total_train_step=0
total_test_step=0
epoch=10
#添加 tensorboard
write= SummaryWriter("../logs_train")
for i in range(epoch):
print("----第{}轮训练开始-----".format(i+1))
#训练步骤开始
for data in train_dataloader:
imgs,targets=data
outputs=cifar10Model(imgs)
loss =loss_fn(outputs,targets)
#优化器优化模型
optimzer.zero_grad()
loss.backward()
optimzer.step()
total_train_step=total_train_step+1
if total_train_step%100==0:
print("训练次数:{},Loss:{}".format(total_train_step,loss.item()))
write.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:
ims,targets=data
output=cifar10Model(ims)
loss=loss_fn(output,targets)
total_test_loss=total_test_loss+loss.item()
accuray=(output.argmax(1)==targets.sum())
total_accuracy=total_accuracy+accuray
print("整体测试集上的loss:{}".format(total_test_loss))
print("整体测试集上的正确率:{}".format(total_accuracy/test_data_size))
write.add_scalar("tess_loss",total_test_loss,total_test_step)
write.add_scalar("test_accuracy",total_accuracy/test_data_size,total_test_step)
total_test_step=total_test_step+1
torch.save(cifar10Model,"cifar10Model_{}.path".format(i))
print("模型已保存")
write.close()