搭建简单的CNN网络模型进行预训练，验证，和预测（准确率不高，单纯记录下模型搭建，训练，验证，预测整套流程）

搭建简单的CNN网络模型进行预训练，验证，和预测（准确率不高，单纯记录下模型搭建，训练，验证，预测整套流程）

以下三个python文件都在同一级目录下
model.py

import torch
from torch import nn


# 搭建神经网络
class Mymodel(nn.Module):
    def __init__(self):
        super(Mymodel, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=32, kernel_size=(5, 5), padding=2, stride=(1, 1)),
            nn.MaxPool2d(2),
            nn.Conv2d(in_channels=32, out_channels=32, kernel_size=(5, 5), padding=2, stride=(1, 1)),
            nn.MaxPool2d(2),
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=(5, 5), padding=2, stride=(1, 1)),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(1024, 64),
            nn.Linear(64, 10)
        )

    def forward(self, x):
        x = self.model(x)
        return x


if __name__ == '__main__':
    # 检查网络的正确性
    mymodel = Mymodel()
    input = torch.ones((64, 3, 32, 32))
    output = mymodel(input)
    print(output.shape)  # torch.Size([64, 10])

train.py

import torch.optim
import torchvision
from torch import nn
from torch.utils.tensorboard import SummaryWriter
import time
from model import Mymodel
# 准备数据集
from torch.utils.data import DataLoader

# 定义训练的设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 如果电脑上有多张显卡
# device=torch.device("cuda:0")   #指定第一张显卡

train_data = torchvision.datasets.CIFAR10(root="./data", train=True, transform=torchvision.transforms.ToTensor(),
                                          download=True)
test_data = torchvision.datasets.CIFAR10(root="./data", train=False, transform=torchvision.transforms.ToTensor(),
                                         download=True)
# length 长度
train_data_size = len(train_data)
test_data_size = len(test_data)
print("训练数据集的长度为:{}".format(train_data_size))
print("测试数据集的长度为:{}".format(test_data_size))

# 利用DataLoader来加载数据集
train_data_loader = DataLoader(train_data, batch_size=64)
test_data_loader = DataLoader(test_data, batch_size=64)

# 创建网络模型
mymodel = Mymodel()
mymodel = mymodel.to(device)
# 创建损失函数
loss_fu = nn.CrossEntropyLoss()
# 调用设备计算损失
loss_fu = loss_fu.to(device)
# 优化器
learning_rate = 0.001
optimizer = torch.optim.SGD(mymodel.parameters(), lr=learning_rate)

# 设置训练网络的一些参数
# 记录训练的次数
total_train_step = 0
# 训练的轮数
epoch = 10

# 添加TensorBoard
writer = SummaryWriter("logs_network")  # 为了画图，可去掉
# 开始时间
start_time = time.time()
for i in range(epoch):
    print("-----------第{}轮训练开始-----------".format(i + 1))
    # 训练步骤开始
    mymodel.train()
    for data in train_data_loader:
        imgs, targets = data
        imgs = imgs.to(device)
        targets = targets.to(device)
        output = mymodel(imgs)
        loss = loss_fu(output, targets)
        # 优化器优化模型
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_train_step += 1
        if total_train_step % 100 == 0:
            end_time = time.time()
            print("训练次数:{},消耗时间:{}s,loss:{}".format(total_train_step, end_time - start_time, loss.item()))

            writer.add_scalar("train_loss", loss.item(), total_train_step)  # 为了画图，可去掉
    # 测试步骤开始
    mymodel.eval()
    # 记录测试集上总的loss
    total_test_loss = 0
    # 记录测试集上总的预测正确的个数
    total_accuracy = 0
    with torch.no_grad():
        for data in test_data_loader:
            imgs, targets = data
            imgs = imgs.to(device)
            targets = targets.to(device)
            outputs = mymodel(imgs)
            loss = loss_fu(outputs, targets)
            total_test_loss += loss.item()
            # 计算正确率
            accuracy = (outputs.argmax(1) == targets).sum()
            total_accuracy += accuracy.item()
        print("第{}轮训练整体测试集上的loss:{}".format(i + 1, total_test_loss))
        print("第{}轮训练整体测试集上的正确率:{}".format(i + 1, total_test_loss / test_data_size))
        writer.add_scalar("test_loss", total_test_loss, i)  # 为了画图，可去掉
        writer.add_scalar("test_accuracy", total_accuracy / test_data_size, i)  # 为了画图，可去掉

        torch.save(mymodel, "./model/mymodel_{}.pth".format(i))
        # torch.save(mymodel.state_dict(),"./model/mymodel_{}.pth".format(i))  #推荐保存方式
        print("模型已保存")
writer.close()

predict.py

import torch
import torchvision
from PIL import Image
from model import Mymodel
img = Image.open("./images/dog.jpg")
print(type(img))
# print(img)
transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((32, 32)),
    torchvision.transforms.ToTensor()
])
img = transform(img)
print(type(img))

mymodel = torch.load("./model/mymodel_29_gpu.pth",map_location=torch.device('cpu'))
# print(mymodel)
# 将输入图片的shape修改成符合输入条件的shape
img = torch.reshape(img, (1, 3, 32, 32))
print(img.shape)
"""
mymodel.eval()
with torch.no_grad():
可以节省预测时间
"""
mymodel.eval()
with torch.no_grad():
    output = mymodel(img)
    print(output.argmax(1))