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"])
显示实时动图,看着蛮炫酷的。