import torch
import torchvision
from torch import nn
from torch.nn import Conv2d, MaxPool2d, Flatten, Linear
from torch.utils.data import DataLoader
# 1.准备数据集并加载
from torch.utils.tensorboard import SummaryWriter
train = torchvision.datasets.CIFAR10('./torchvision_dataset',train=True,download=False,transform=torchvision.transforms.ToTensor())
test = torchvision.datasets.CIFAR10('./torchvision_dataset',train=False,download=False,transform=torchvision.transforms.ToTensor())
train_loader = DataLoader(train,batch_size=64,drop_last=True)
test_loader = DataLoader(test,batch_size=64,drop_last=True)
#获取两个数据集的长度
train_len = len(train)
test_len = len(test)
#定义网络结构
class Zkl(nn.Module):
def __init__(self):
super(Zkl, self).__init__()
# self.conv1 = Conv2d(3,32,5,padding=2)
# self.maxpool1 = MaxPool2d(2)
# self.conv2 = Conv2d(32,32,5,padding=2)
# self.maxpool2 = MaxPool2d(2)
# self.conv3 = Conv2d(32,64,5,padding=2)
# self.maxpool3 = MaxPool2d(2)
# self.flatten = Flatten()
# self.linear1 = Linear(1024,64)
# self.linear2 = Linear(64,10)
self.modle1 = nn.Sequential(
Conv2d(3, 32, 5, padding=2),
MaxPool2d(2),
Conv2d(32, 32, 5, padding=2),
MaxPool2d(2),
Conv2d(32, 64, 5, padding=2),
MaxPool2d(2),
Flatten(),
Linear(1024, 64),
Linear(64,10)
)
def forward(self,x):
# x = self.conv1(x)
# x = self.maxpool1(x)
# x = self.conv2(x)
# x = self.maxpool2(x)
# x = self.conv3(x)
# x = self.maxpool3(x)
# x = self.flatten(x)
# x = self.linear1(x)
# x = self.linear2(x)
x = self.modle1(x)
return x
#定义模型对象
zkl = Zkl()
#定义损失函数
loss = nn.CrossEntropyLoss()
#定义优化器
learn_rate = 1e-2
optim = torch.optim.SGD(zkl.parameters(),lr=learn_rate)
#设置训练网络的一些参数
#记录训练的次数
total_train_step = 0
#记录测试的次数
total_test_step = 0
#定义训练轮数
epoch=10
#添加tensorboard
writer = SummaryWriter('train_log')
for i in range(epoch):
print(f"-------第 {i} 轮训练开始----------")
zkl.train() # 设置这个是使模型进入一个特定状态,只对dropout层或者batchnorml层有作用,并且有这些层就必须调用
for data in train_loader:
imgs,targets = data
output = zkl(imgs)
train_loss=loss(output,targets)
#优化器调优模型
optim.zero_grad()
train_loss.backward()
optim.step()
total_train_step += 1
if total_train_step%100 == 0: # 减少打印次数,防止看不清楚结果
print(f"训练次数 {total_train_step}:loss={train_loss.item()}") # 加上.item()是为了使输出的是一个值,而不是tensor(1)这样的表达形式
writer.add_scalar('model_train',train_loss,total_train_step)
#每一轮训练完后开始测试
zkl.eval() # 设置这个是使模型进入一个特定状态,只对dropout层或者batchnorml层有作用,并且有这些层就必须调用
total_test_loss = 0
total_accuracy = 0
with torch.no_grad(): # 不进行优化器在测试过程中,只需要使用训练好的参数
for data in test_loader:
imgs,targets = data
output = zkl(imgs)
loss_test = loss(output,targets)
total_test_loss += loss_test.item()
accuracy = (output.argmax(1) == targets).sum()
total_accuracy += accuracy
print(f"第 {i} 轮训练后测试的loss: {total_test_loss}")
print(f"第 {i} 轮训练后测试的accuracy: {total_accuracy/test_len}")
writer.add_scalar('model_test', loss_test, total_test_step)
writer.add_scalar('model_accuracy',total_accuracy/test_len,total_test_step)
total_test_step += 1
#保存每一轮训练的结果
torch.save(zkl,f'modle_train_{i}.pth')
print('模型已保存')
writer.close()
pytorch入门19:模型的完整训练套路
最新推荐文章于 2024-03-15 22:27:23 发布