pytorch8-基于FashionMNIST搭建卷积神经网络进行识别

import torch
import torch.nn as nn
import torch.utils.data as Data
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.metrics import accuracy_score, confusion_matrix, classification_report
import torch
import seaborn as sns
import copy
import time
from torchvision import transforms
from torchvision.datasets import FashionMNIST

train_data = FashionMNIST(
    root="./data/FashionMNIST",  # 数据路径
    train=True, # 只使用训练数据集
    transform=transforms.ToTensor(),
    download=True
)
class_label = train_data.classes
class_label[0] = "T-shirt"

# 定义一个训练数据加载器
train_loader = Data.DataLoader(
    dataset=train_data,
    batch_size=64,
    shuffle=False,
    num_workers=2,
)

# 计算train_loader有多少个batch
print("train_loader 的batch数目为：", len(train_loader))

# for step, (b_x, b_y) in enumerate(train_loader):
#     if step > 0:
#         break
#
#     # 可视化一个batch的图像
#     batch_x = b_x.squeeze().numpy()
#     # print("batch_x", batch_x)
#     batch_y = b_y.numpy()
#     # print("batch_y", batch_y)
#
#     class_label = train_data.classes
#     class_label[0] = "T-shirt"
#     plt.figure(figsize=(12, 5))
#     for ii in np.arange(len(batch_y)):
#         plt.subplot(4, 16, ii + 1)
#         plt.imshow(batch_x[ii, :, :], cmap=plt.cm.gray)
#         plt.title(class_label[batch_y[ii]], size=9)
#         plt.axis("off")
#         plt.subplots_adjust(wspace=0.05)
#     # plt.show()

# 定义一个测试数据加载器
test_data = FashionMNIST(
    root="./data/FashionMNIST",  # 数据路径
    train=True,
    download=True
)

# 为数据添加一个通道纬度，并且取值范围缩小到0-1之间
test_data_x = test_data.data.type(torch.FloatTensor) / 255.0
test_data_x = torch.unsqueeze(test_data_x, dim=1)
test_data_y = test_data.targets  # 测试集的标签
print("test_data_x.shape", test_data_x.shape)
print("test_data_y.shape", test_data_y.shape)


class MyConvNet(nn.Module):
    def __init__(self):
        super(MyConvNet, self).__init__()  # 继承MyConvNet父类里的全部东西，即nn.Module里的东西

        # 定义一个卷积层
        self.conv1 = nn.Sequential(
            nn.Conv2d(
                in_channels=1,  # 输入
                out_channels=16,  # 输出
                kernel_size=3,  # 卷积核
                stride=1,  # 步长
                padding=1,  # 填充
            ),  # 卷积后（1*28*28）-（16*28*28）

            nn.ReLU(),  # 激活函数
            nn.AvgPool2d(
                kernel_size=2,
                stride=2,
            ),  # 平均值池化后：（16*28*28）-（16*14*14）
        )

        # 定义第二个卷积层
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, 3, 1, 0),  # 卷积操作（16*14*14）-（32*12*12）
            nn.ReLU(),
            nn.AvgPool2d(2, 2)  # 最大池化操作 （32*12*12）-（32*6*6）
        )

        # 定义分类器
        self.classifiler = nn.Sequential(
            nn.Linear(32 * 6 * 6, 256),  # 定义第一个全连接层
            nn.ReLU(),
            nn.Linear(256, 128),  # 定义第二个全连接层
            nn.ReLU(),
            nn.Linear(128, 10)  # 定义输出层
        )

    # 定义网络的前向传播路径
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)
        output = self.classifiler(x)
        return output


def train_model(model, traindataloader, train_rate, criterion, optimizer, num_epochs=25):
    # model:网络模型  traindataloader：数据加载器  train_rate：训练数据集
    # criterion：损失函数  optimizer：优化方法  num_epochs：训练的轮数

    # 计算训练使用的batch数量
    batch_num = len(traindataloader)
    train_batch_num = round(batch_num * train_rate)  # 保留四舍五入值，训练阶段batch的数目

    # 复制模型的参数
    best_model_wts = copy.deepcopy(model.state_dict())

    best_acc = 0.0
    train_loss_all = []
    train_acc_all = []
    val_loss_all = []
    val_acc_all = []

    since = time.time()

    for epoch in range(num_epochs):
        print('Epoch{}/{}'.format(epoch, num_epochs - 1))  # print方式------
        print('-' * 50)

        # 每个epoch有两个阶段
        train_loss = 0.0
        train_corrects = 0
        train_num = 0
        val_loss = 0.0
        val_corrects = 0
        val_num = 0

        for step, (b_x, b_y) in enumerate(traindataloader):
            if step < train_batch_num:
                model.train()  # 将模型设置为训练模式
                output = model(b_x)
                pre_lab = torch.argmax(output, 1)
                loss = criterion(output, b_y)
                optimizer.zero_grad()  # 反向传播前将梯度数据清空
                loss.backward()  # 反向传播，计算出梯度
                optimizer.step()  # 调用step函数进行网络参数更新

                train_loss += loss.item() * b_x.size(0)
                train_corrects += torch.sum(pre_lab == b_y.data)
                train_num += b_x.size(0)
            else:
                model.eval()  # 设置模型为评估模式
                output = model(b_x)

                pre_lab = torch.argmax(output, 1)
                loss = criterion(output, b_y)
                val_loss += loss.item() * b_x.size(0)
                val_corrects += torch.sum(pre_lab == b_y.data)
                val_num += b_x.size(0)

        # 计算一个epoch在训练集和验证集上的损失和精度
        train_loss_all.append(train_loss / train_num)
        train_acc_all.append(train_corrects.double().item() / train_num)
        val_loss_all.append(val_loss / val_num)
        val_acc_all.append(val_corrects.double().item() / val_num)
        print('{}Train Loss: {:.4f}  Train Acc: {:.4f}'.format(epoch, train_loss_all[-1], train_acc_all[-1]))
        print('{}Val Loss: {:.4f}  val Acc: {:.4f}'.format(epoch, val_loss_all[-1], val_acc_all[-1]))

        # 拷贝模型下最高精度的参数
        if val_acc_all[-1] > best_acc:
            best_acc = val_acc_all[-1]
            best_model_wts = copy.deepcopy(model.state_dict())
        time_use = time.time() - since
        print("Train and val complete in {:.0f}m {:.0f}s".format(time_use // 60, time_use % 60))

    print("train_loss_all", train_loss_all)
    print("val_loss_all", val_loss_all)
    print("val_acc_all", val_acc_all)
    print("epoch", range(num_epochs))

    # 使用最好模型的参数
    model.load_state_dict(best_model_wts)
    train_process = pd.DataFrame(
        data={"epoch": range(num_epochs),
              "train_loss_all": train_loss_all,
              "val_loss_all": val_loss_all,
              "train_acc_all": train_acc_all,
              "val_acc_all": val_acc_all})
    return model, train_process


if __name__ == "__main__":
    myconvenet = MyConvNet()
    # print(myconvenet)
    optimizer = torch.optim.Adam(myconvenet.parameters(), lr=0.003)  # 定义优化函数----优化器
    criterion = nn.CrossEntropyLoss()  # 定义损失函数
    myconvenet, train_process = train_model(myconvenet, train_loader, 0.8, criterion, optimizer, num_epochs=25)

    plt.figure(figsize=(12, 4))  # 定义画面大小
    # 一行两列的第一列
    plt.subplot(1, 2, 1)
    plt.plot(train_process.epoch, train_process.train_loss_all, "ro-", label="Train loss")
    plt.plot(train_process.epoch, train_process.val_loss_all, "bs-", label="Val loss")
    plt.legend()  # 给图像加上图例
    plt.xlabel("epoch")
    plt.ylabel("Loss")

    # 一行两列的第二列
    plt.subplot(1, 2, 2)
    plt.plot(train_process.epoch, train_process.train_acc_all, "ro-", label="Train acc")
    plt.plot(train_process.epoch, train_process.val_acc_all, "bs-", label="Val acc")
    plt.xlabel("epoch")
    plt.ylabel("acc")
    plt.legend()  # 给图像加上图例

    # 对测试集进行预测，并进行可视化预测结果
    myconvenet.eval()
    output = myconvenet(test_data_x)
    pre_lab = torch.argmax(output, 1)
    acc = accuracy_score(test_data_y, pre_lab)
    print('在测试集上的预测精度为：{}'.format(acc))

    conf_mat = confusion_matrix(test_data_y, pre_lab)
    df_cm = pd.DataFrame(conf_mat, index=class_label, columns=class_label)

    heatmap = sns.heatmap(df_cm, annot=True, fmt="d", cmap="YlGnBu")
    heatmap.yaxis.set_ticklabels(heatmap.yaxis.get_ticklabels(), rotation=0, ha='right')
    heatmap.xaxis.set_ticklabels(heatmap.xaxis.get_ticklabels(), rotation=45, ha='right')

    plt.ylabel('True label')
    plt.xlabel('Predicted label')
    plt.show()