1
import os
2
import cv2
3
import PIL
4
import torch
5
import torchvision
6
import numpy as np
7
from torch import nn
8
from torch.utils.data import DataLoader
9
from torchvision.models import VGG16_Weights
10
from torch.utils.tensorboard import SummaryWriter

1
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
2
print("使用设备：{}".format(device))

1
# 准备数据集
2
train_dataset = torchvision.datasets.CIFAR10(
3
    root        = "./dataset/CIFAR10",
4
    train       = True,
5
    download    = True,
6
    transform   = torchvision.transforms.ToTensor()
7
)
8
train_dataset_size = len(train_dataset)
9

10
test_dataset = torchvision.datasets.CIFAR10(
11
    root        = "./dataset/CIFAR10",
12
    train       = False,
13
    download    = True,
14
    transform   = torchvision.transforms.ToTensor()
15
)
16
test_dataset_size  = len(test_dataset)
17

18
print("----数据集导入成功----")
19
print("训练集大小：{}".format(train_dataset_size))
20
print("测试集大小：{}".format(test_dataset_size))
21

22

23
# 加载数据集
24
train_dataloader = DataLoader(
25
    dataset     = train_dataset,
26
    batch_size  = 64,
27
    shuffle     = True,
28
    num_workers = 0,
29
    drop_last   = False
30
)
31

32
test_dataloader = DataLoader(
33
    dataset     = test_dataset,
34
    batch_size  = 64,
35
    shuffle     = True,
36
    num_workers = 0,
37
    drop_last   = False
38
)

1
class My_Model(nn.Module):
2
    def __init__(self):
3
        super().__init__()
4
        self.conv1      = nn.Conv2d(
5
            in_channels     = 3,
6
            out_channels    = 32,
7
            kernel_size     = 5,
8
            padding         = 2
9
        )
10
        self.maxpool1   = nn.MaxPool2d(
11
            kernel_size = 2,
12
        )
13
        self.conv2      = nn.Conv2d(
14
            in_channels     = 32,
15
            out_channels    = 32,
16
            kernel_size     = 5,
17
            padding         = 2
18
        )
19
        self.maxpool2   = nn.MaxPool2d(
20
            kernel_size = 2,
21
        )
22
        self.conv3      = nn.Conv2d(
23
            in_channels     = 32,
24
            out_channels    = 64,
25
            kernel_size     = 5,
26
            padding         = 2
27
        )
28
        self.maxpool3   = nn.MaxPool2d(
29
            kernel_size = 2,
30
        )
31
        self.flatten    = nn.Flatten()
32
        self.linear1    = nn.Linear(
33
            in_features    = 64 * 4 * 4,
34
            out_features   = 64
35
        )
36
        self.linear2    = nn.Linear(
37
            in_features    = 64,
38
            out_features   = 10
39
        )
40
        self.seq        = nn.Sequential(
41
            self.conv1,
42
            self.maxpool1,
43
            self.conv2,
44
            self.maxpool2,
45
            self.conv3,
46
            self.maxpool3,
47
            self.flatten,
48
            self.linear1,
49
            self.linear2
50
        )
51

52
    def forward(self, input):
53
        output = self.seq(input)
54
        return output
55

56
my_model = My_Model()
57
my_model = my_model.to(device)

1
loss_fn = nn.CrossEntropyLoss()
2
loss_fn = loss_fn.to(device)

1
Learnig_rate = 0.002
2
optim = torch.optim.Adam(
3
    params      = my_model.parameters(),
4
    lr          = Learnig_rate
5
)

1
total_train_step = 0 # 记录训练的次数
2
total_test_step  = 0 # 记录测试的次数
3
epoch = 10          # 训练轮数
4
writer = SummaryWriter("./logs")  # TensorBoard 日志保存路径
5

6
for i in range(epoch):
7
    print("--------第 {} 轮训练开始--------".format(i + 1))
8

9
    # 训练步骤
10
    my_model.train()
11
    for data in train_dataloader:
12
        imgs, targets = data
13
        imgs    = imgs.to(device)
14
        targets = targets.to(device)
15

16
        outputs = my_model(imgs)
17
        loss = loss_fn(outputs, targets)
18

19
        optim.zero_grad()   # 清空上一步的残余更新参数值
20
        loss.backward()     # 反向传播，计算参数更新值
21
        optim.step()       # 更新参数
22

23
        total_train_step += 1
24
        if total_train_step % 100 == 0:
25
            print("训练次数：{}，Loss：{}".format(total_train_step, loss.item()))
26
            writer.add_scalar("train_loss", loss.item(), total_train_step)
27

28
    torch.save(my_model.state_dict(), "./My_Model/model_{}.pth".format(i))
29
    print("模型已保存")
30

31
    # 测试步骤
32
    my_model.eval()
33
    total_test_loss = 0
34
    total_accuracy  = 0
35
    with torch.no_grad():
36
        for data in test_dataloader:
37
            imgs, targets = data
38
            imgs    = imgs.to(device)
39
            targets = targets.to(device)
40

41
            outputs = my_model(imgs)
42
            loss = loss_fn(outputs, targets)
43
            total_test_loss += loss.item()
44
            accuracy = (outputs.argmax(1) == targets).sum()
45
            total_accuracy += accuracy
46

47
    print("整体测试集上的Loss：{}".format(total_test_loss))
48
    print("整体测试集上的正确率：{}".format(total_accuracy / test_dataset_size))
49
    writer.add_scalar("test_loss", total_test_loss, total_test_step)
50
    writer.add_scalar("test_accuracy", total_accuracy / test_dataset_size, total_test_step)
51
    total_test_step += 1