#1_Pytorch下的CNN图片分类器

导入所需框架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_Pytorch下的CNN图片分类器_第1张图片

 训练准备

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数据集包含的类别及对应                          网上找一张小狗的照片

#1_Pytorch下的CNN图片分类器_第2张图片                                        #1_Pytorch下的CNN图片分类器_第3张图片

                                                                                         dog.png

运行结果

#1_Pytorch下的CNN图片分类器_第4张图片

 结果即tensor[5]='dog',算是识别出来了...

你可能感兴趣的:(机器学习)