pytorch系列(七)：CNN实现数字图像分类

浅浅的花香味﹌ 2023-07-10 14:04 1阅读 0赞

import  torch
    import torchvision#包含一些数据集，像如mnist
    import torch.utils.data as Data
    import torch.nn as nn
    import matplotlib.pyplot as plt
    
    #超参数
    EPOCH=1
    BATCH_SIZE=50
    LR=0.01
    DOWNLOAD_MNIST=True#是否下载数据集
    
    
    #mnist手写数字数据集
    train_data=torchvision.datasets.MNIST(
        root="mnist",
        train=True,#是否是训练
        transform=torchvision.transforms.ToTensor(),#将图像或numpy数组数据形式转化成tensor的floattensor(C*H*M)形式
        download=DOWNLOAD_MNIST #初始下载数据集用，后面就不用了
    )
    
    #测试集
    test_data=torchvision.datasets.MNIST("mnist",train=False)
    #批训练和dataloader数据封装
    train_loader=Data.DataLoader(
        dataset=train_data,
        batch_size=BATCH_SIZE,
        shuffle=True
        
    )
    
    #测试数据  只取前2000个
    test_x = torch.unsqueeze(test_data.test_data, dim=1).type(torch.FloatTensor)[:2000]/255.   # shape from (2000, 28, 28) to (2000, 1, 28, 28), value in range(0,1)
    test_y = test_data.test_labels[:2000]
    
    #实现卷积层模型
    class CNN(torch.nn.Module):
        def __init__(self):
            super(CNN,self).__init__()
            #第一层卷积
            self.conv1=torch.nn.Sequential(
                torch.nn.Conv2d(
                    in_channels=1,
                    out_channels=16,
                    kernel_size=5,
                    stride=1,
                    padding=2#padding=(kernel_size-1)/2时，卷积出来的图片尺寸不变
                ),#此时输入是(,28*28)输出是(16*28*28)
                torch.nn.ReLU(),
                torch.nn.MaxPool2d(kernel_size=2)#此时图片尺寸减小一半  16*28*28--》16*14*14
            )
            
            #第二层卷积
            self.conv2=torch.nn.Sequential(
                torch.nn.Conv2d(16,32,5,1,2),
                torch.nn.ReLU(),
                torch.nn.MaxPool2d(kernel_size=2)
            )#输入是16*14*14  输出是32*7*7
            
            #添加全连接层
            self.output=torch.nn.Linear(32*7*7,10)
            
            
        #前向传播  卷积
        def forward(self,x):
            x=self.conv1(x)
            x=self.conv2(x)
            #将此时的32*7*7的图片铺平  拉成一维32*7*7的数组形式
            x=x.view(x.size(0),-1)# 展平多维的卷积图成 (batch_size, 32 * 7 * 7)
            output=self.output(x)
            return output
            
            
            
    
    #实例化卷积神经网络
    cnn=CNN()
    print(cnn)#打印网络结构
    
    #训练
    
    #优化器
    optimizer=torch.optim.Adam(cnn.parameters(),lr=LR)
    #损失函数
    loss_func=torch.nn.CrossEntropyLoss()#用交叉熵损失函数  标签是one-hot编码
    
    #训练和测试
    for epoch in range(EPOCH):
        #批训练
        for step,(batch_x,batch_y) in enumerate(train_loader):
            output=cnn(batch_x)
            loss=loss_func(output,batch_y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
    
    
    #预测
    test_output = cnn(test_x[:10])
    pred_y = torch.max(test_output,1)[1].data.numpy().squeeze()
    print(pred_y, 'prediction number')
    print(test_y[:10].numpy(), 'real number')