Pytorch关于CIFAR-10测试完整代码

#_*_ coding:utf-8 _*_
# pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117
import torch
from torch import nn
from torch.nn import Conv2d,MaxPool2d,Linear,Sequential,Flatten
from torch.utils.data import DataLoader
import torch.optim.optimizer
import torchvision
from torch.utils.tensorboard import SummaryWriter
print(SummaryWriter)
class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.covn1=Conv2d(in_channels=3,out_channels=32,kernel_size=5,padding=2)
        self.maxpool1=MaxPool2d(kernel_size=2)
        self.covn2=Conv2d(in_channels=32,out_channels=32,kernel_size=5,padding=2)
        self.maxpool2=MaxPool2d(kernel_size=2)
        self.covn3=Conv2d(in_channels=32,out_channels=64,kernel_size=5,padding=2)
        self.maxpool3=MaxPool2d(kernel_size=2)
        self.flaten=Flatten()
        self.linear1=Linear(in_features=1024,out_features=64)
        self.linear2=Linear(in_features=64,out_features=10)
 
    def forward(self,x):
        x= self.covn1(x)
        x=self.maxpool1(x)
        x=self.covn2(x)
        x=self.maxpool2(x)
        x=self.covn3(x)
        x=self.maxpool3(x)
        x=self.flaten(x)
        x=self.linear1(x)
        x=self.linear2(x)
 
        return x
 
 
if __name__ == "__main__":
    model = Model()
    print(model)
 
    #下载数据集
    train_data = torchvision.datasets.CIFAR10('../CIFAR-10/',train=True,transform=torchvision.transforms.ToTensor(),download=False)
    test_data = torchvision.datasets.CIFAR10('../CIFAR-10/',train=False,transform=torchvision.transforms.ToTensor(),download=False)
    #数据加载器
    train_data_loader = DataLoader(train_data,batch_size=64)
    test_data_loader = DataLoader(test_data,batch_size=64)
 
    #损失函数
    loss = nn.CrossEntropyLoss()
    #优化器
    optim = torch.optim.SGD(model.parameters(),lr=0.01)
 
    #创建可视化
    write = SummaryWriter('CIFAR_logs')
    #开始训练数据
 
    eplo =10
    train_step = 0
    for i in range(10):
        print("开始训练:".format(i))
        for img,laber in train_data_loader:
            input_d = model(img)
            loss_fn =loss(input_d,laber)
 
            optim.zero_grad() #梯度清零
            loss_fn.backward() #反向计算
            optim.step()#更新梯度
 
            train_step += 1
            if train_step%100==0:
                print("训练次数：{}".format(train_step),"损失：{}".format(loss_fn.item()))
                write.add_scalar('train_loss',loss_fn.item(),train_step)
 
        total_accuracy = 0
        total_test_loss = 0
        total_test_step = 0
        with torch.no_grad():
            for imgs,labler in test_data_loader:
                output = model(imgs)
                loss_re=loss(output,labler)
 
                accuracy = (output.argmax(1) == labler).sum()
                total_accuracy += accuracy
                total_test_loss += loss_re
                total_test_step += 1
 
        print("测试集上的正确率{}".format(total_accuracy/len(test_data)))
        print("测试集上的损失{}".format(total_test_loss))
        write.add_scalar("test_loss",total_test_loss.item(),total_test_step)
        torch.save(model,"model10")
    write.close()
 
 
 
# tensorboard --logdir=CIFAR_logs --port=2017

预测代码：

from torchvision import datasets, transforms
import numpy as np
from PIL import Image
import torch
import torch.nn.functional as F
from cov01 import Model
 
classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
 
if __name__ == '__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = torch.load('model10')  # 加载模型
    model = model.to(device)
    model.eval()  # 把模型转为test模式
 
    img = Image.open("../bird.jpg")
 
    trans = transforms.Compose(
        [
            transforms.CenterCrop(32),
            transforms.ToTensor(),
            # transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
        ])
 
    img = trans(img)
    img = img.to(device)
    img = img.unsqueeze(0)  # 图片扩展多一维,因为输入到保存的模型中是4维的[batch_size,通道,长，宽]，而普通图片只有三维，[通道,长，宽]
    # 扩展后，为[1，1，28，28]
    output = model(img)
    prob = F.softmax(output, dim=1)  # prob是10个分类的概率
    print(prob)
    value, predicted = torch.max(output.data, 1) #按照维度返回最大概率dim = 0 表示按列求最大值，并返回最大值的索引，dim = 1 表示按行求最大值，并返回最大值的索引
    print(predicted.item())
    print(value)
    pred_class = classes[predicted.item()]
    print(pred_class)
 
 
'''
记住：
torch.max()[0]， 只返回最大值的每个数
troch.max()[1]， 只返回最大值的每个索引
torch.max()[1].data 只返回variable中的数据部分（去掉Variable containing:）
torch.max()[1].data.numpy() 把数据转化成numpy ndarry
torch.max()[1].data.numpy().squeeze() 把数据条目中维度为1 的删除掉
torch.max(tensor1,tensor2) element-wise 比较tensor1 和tensor2 中的元素，返回较大的那个值
'''

E:\开发工具\pythonProject\studyLL\venv\Scripts\python.exe E:/开发工具/pythonProject/studyLL/pytorch01/predict.py
<class 'torch.utils.tensorboard.writer.SummaryWriter'>
tensor([[0.1517, 0.0265, 0.3715, 0.1244, 0.1860, 0.0556, 0.0084, 0.0167, 0.0457,
         0.0134]], device='cuda:0', grad_fn=<SoftmaxBackward0>)
2
tensor([1.9028], device='cuda:0')
bird

Process finished with exit code 0

 

成功预测为bird