导入所需框架package
import torch
import torchvision
from torch import nn
from torch.nn import Sequential, Conv2d, MaxPool2d, Flatten, Linear,ReLU
from torch.utils.data import DataLoader
取得数据集
所使用的数据集为torchvision自带的CIFAR10数据集。
该数据集共有60000张彩色图像,这些图像是32*32,分为10个类,每类6000张图。这里面有50000张用于训练,构成了5个训练批,每一批10000张图;另外10000用于测试,单独构成一批。测试批的数据里,取自10类中的每一类,每一类随机取1000张。抽剩下的就随机排列组成了训练批
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)
数据集放在项目路径dataset文件下,将其转化为tensor数据类型,并下载
数据集分为训练数据集和测试数据集: train=Ture表示使用CIFAR10的训练数据集
train=False表示使用CIFAR10的测试数据集
数据集装载
test_loader=DataLoader(dataset=test_data,batch_size=64,shuffle=True)
train_loader=DataLoader(dataset=test_data,batch_size=64,shuffle=True)
batch_size:即一次训练所抓取的数据样本数量,其大小影响训练速度和模型优化,一般在16~256
shuffle:每一轮(epoch)训练是否将训练模型的数据集进行洗牌
准备工作已经完成,进入正题
最简神经网络搭建
CNN:卷积神经网络主要由这几类层构成:输入层、卷积层,ReLU层、池化(Pooling)层和全连接层(全连接层和常规神经网络中的一样)。通过将这些层叠加起来,就可以构建一个完整的卷积神经网络。
先看代码
class Demon(nn.Module):
def __init__(self):
super(Demon, self).__init__()
self.model=Sequential(
Conv2d(in_channels=3,out_channels=32,kernel_size=5,padding=2), #卷积层
ReLU(),
MaxPool2d(kernel_size=2),#最大池化层
Conv2d(32, 32, 5, padding = 2),
ReLU(),
MaxPool2d(2),
Conv2d(32, 64, 5, padding = 2),
ReLU(),
MaxPool2d(2),
Flatten(),#铺平为一维
Linear(1024,64),#线性层 输入:1024/输出:64
Linear(64,10),#线性层 输入:64/输出:10
)
def forward(self,x):
x=self.model(x)
return x
参数解释如下:
训练准备
1.实例化网络模型
2.设定学习速率,与梯度下降的效果密切相关
3.给定优化器
4.设定训练轮数
#创建网络模型
demon=Demon()
#损失函数用以反向传播
loss_fn=nn.CrossEntropyLoss()
#给定优化器
learning_rate=0.01
optimizer=torch.optim.SGD(demon.parameters(),lr=learning_rate)#引入随机梯度下降的优化器 用于update module中的参数
#训练+测试
train_step=0#记录训练次数
test_step=0#测试次数
#训练轮数=100
epoch=200
训练&测试
训练步前设置为训练模式:demon.train()
测试步前设置为测试模式:demon.eval()
for i in range(epoch):
print("---------------------------第{}轮训练---------------------------".format(i+1))
#训练步骤
demon.train()#设置为训练模式
for data in train_loader:
imgs, label = data
outputs = demon(imgs)
result_loss = loss_fn(outputs, label)
#优化器优化
optimizer.zero_grad() # 每一次都要把上一次的梯度清零
result_loss.backward() # 反向传播
optimizer.step()#update
train_step+=1
if(train_step%50==0):
print("{}次训练的loss为:{}".format(train_step,result_loss))
#测试步骤
demon.eval()
total_loss=0
total_accuracy=0
with torch.no_grad():#不用算梯度 加快运行速度
for data in test_loader:
imgs,label=data
outputs=demon(imgs)
loss=loss_fn(outputs,label)
total_loss+=loss
total_accuracy+=(outputs.argmax(1)==label).sum()
print("测试集上的正确率:{}".format(total_accuracy/test_data_size))
print("测试集loss:{}".format(total_loss))
torch.save(demon,"demon_{}.pth".format(i))
print("模型已保存")
每50次输出一次train_loss,test_loss以及测试集上的正确率
torch.save(model,"moedel_name")保存已训练好的模型
开始训练......等(因为是在CPU上训练所以比较慢)
训练的效果如何?
测试
从目录下读取一张图片,转为RGB三通道,通过transform()转为32×32
torch.load("model_name")装载训练时的模型
取有最大概率的类别为识别得到类别
import torch
from PIL import Image
import torchvision
image_path="dog.png"
image=Image.open(image_path)
image=image.convert('RGB')
transform=torchvision.transforms.Compose([torchvision.transforms.Resize((32,32)),torchvision.transforms.ToTensor()])
image_tensor=transform(image)
model=torch.load("demon.pth")
print(model)
image_tensor=torch.reshape(image_tensor,(1,3,32,32))
model.eval()
with torch.no_grad():
output=model(image_tensor)
print(output.argmax(1))
测试结果
CFAR10数据集包含的类别及对应 网上找一张小狗的照片
dog.png
运行结果
结果即tensor[5]='dog',算是识别出来了...