毕业设计记录-Pytorch学习-HidderLayer训练过程可视化

Pytorch学习-HidderLayer训练过程可视化

import torch
import torch.nn as nn
import torchvision
import torchvision.utils as vutils
import torch.utils.data as Data
from sklearn.metrics import accuracy_score
import matplotlib.pyplot as plt
import hiddenlayer as hl
from torchviz import make_dot
from tensorboardX import SummaryWriter
import os
import time

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
train_data = torchvision.datasets.MNIST(
    root="./data/FashionMNIST",
    train=True,
    transform=torchvision.transforms.ToTensor(),
    download=False
)
train_loader = Data.DataLoader(
    dataset=train_data,
    batch_size=256,
    shuffle=True,
    num_workers=2,
    # pin_memory=True  # 将内存的tensor转移到gpu
)
test_data = torchvision.datasets.MNIST(
    root="./data/FashionMNIST",
    train=False,
    transform=torchvision.transforms.ToTensor(),
    download=False
)
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
test_data_x = test_data_x.to(device)
test_data_y = test_data_y.to(device)
# print("test_data_x.shape:", test_data_x.shape)
# print("test_data_y.shape:", test_data_y.shape)
# 输出网络结构


class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        # 第一层卷积
        # 图片尺寸 f=28
        self.conv1 = nn.Sequential(
            nn.Conv2d(
                in_channels=1,
                out_channels=16,
                kernel_size=3,
                stride=1,
                padding=1,
            ),
            nn.ReLU(),
            # 卷积之后，图片的尺寸 f=(28+2*1-3)/1 + 1 = 28
            nn.AvgPool2d(
                kernel_size=2,
                stride=2,
            )
            # 经过平均池化后，图片的尺寸 f=(28-2)/2 + 1 = 14
        )
        # 第二层卷积
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, 3, 1, 1),
            # 卷积之后，图片的尺寸 f=(14+2*1-3)/1 + 1 = 14
            nn.ReLU(),
            nn.AvgPool2d(2, 2)
            # 经过平均池化后，图片的尺寸 f=(14-2)/2 + 1 = 7
        )
        # 全连接层
        self.fc = nn.Sequential(
            nn.Linear(
                in_features=32*7*7,
                out_features=128
            ),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU()
        )
        self.out = nn.Linear(64, 10)
        # 定义网络的前向传播路径

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        output = self.out(x)
        return output


if __name__ == "__main__":

    MyConvNet = ConvNet().to(device)
    # 定义优化器
    optimizer = torch.optim.Adam(MyConvNet.parameters(), lr=0.003)
    # 交叉熵损失函数
    loss_func = nn.CrossEntropyLoss()
    history1 = hl.History()
    canvas1 = hl.Canvas()
    print_step = 100  # 每经过100次迭代之后输出损失
    # 对模型进行迭代训练，对所有的数据训练epoch轮
    for epoch in range(5):
        # 对训练数据的加载器进行迭代计算
        for step, (b_x, b_y) in enumerate(train_loader):
            b_x = b_x.to(device)
            b_y = b_y.to(device)
            # 计算每个batch的损失
            output = MyConvNet(b_x)  # 卷积神经网络计算得出的结果
            loss = loss_func(output, b_y)  # 交叉熵损失函数
            loss = loss.to(device)
            optimizer.zero_grad()  # 每个迭代步的梯度初始化为0
            loss.backward()  # 后向传播
            optimizer.step()  # 使用梯度进行优化
            # 计算迭代次数
            # 计算每经过print_step次迭代后的输出
            if step % print_step == 0:
                # 计算在测试集上的精度
                output = MyConvNet(test_data_x)
                _, pre_lab = torch.max(output, 1)
                test_data_y = test_data_y.cpu()
                pre_lab = pre_lab.cpu()
                acc = accuracy_score(test_data_y, pre_lab)
                # 计算每个epoch和step的模型的输出特征
                history1.log(
                    (epoch, step),
                    train_loss=loss,
                    test_acc=acc,
                    # 第二个全连接层权重
                    hidden_weight=MyConvNet.fc[2].weight
                )
            with canvas1:
                canvas1.draw_plot(history1["train_loss"])
                canvas1.draw_plot(history1["test_acc"])
                canvas1.draw_image(history1["hidden_weight"])

显示实时动图，看着蛮炫酷的。