pytorch使用TensorBoard可视化损失函数曲线、精度信息

79 篇文章 10 订阅

PyTorch已经内置了TensorBoard的相关接口，用户在安装后便可调用相关接口进行数据可视化

加载数据

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms as T
from tqdm import tqdm

data_transform = T.Compose([
T.RandomResizedCrop(32),
T.ToTensor(),
])

train_dataset = torchvision.datasets.MNIST('./', train=True, transform=data_transform)
val_dataset = torchvision.datasets.MNIST('./', train=False, transform=data_transform)



定义网络

class LeNet5(nn.Module):
def __init__(self,num_class=10):
super(LeNet5,self).__init__()
self.conv1 = nn.Conv2d(1, 6, 5)
self.pool1 = nn.AvgPool2d((2, 2))

self.conv2 = nn.Conv2d(6, 16, 5)
self.pool2 = nn.AvgPool2d((2, 2))

self.conv3 = nn.Conv2d(16, 120, 5)

self.relu = nn.ReLU()

self.fc1 = nn.Linear(120, 84)
self.fc2 = nn.Linear(84, num_class)

def forward(self, x):
# x: torch.Size([1, 3, 32, 32])

x = self.conv1(x) # torch.Size([1, 6, 28, 28])
x = self.relu(x)
x = self.pool1(x) # torch.Size([1, 6, 14, 14])

x = self.conv2(x) # torch.Size([1, 16, 10, 10])
x = self.relu(x)
x = self.pool2(x) # torch.Size([1, 16, 5, 5])

x = self.conv3(x) # torch.Size([1, 120, 1, 1])
x = self.relu(x)

x = x.flatten(start_dim=1) # torch.Size([1, 120])

x = self.fc1(x) # torch.Size([1, 84])
x = self.relu(x)
x = self.fc2(x) # torch.Size([1, 5])

return x


定义模型等组件

model = LeNet5(10)
loss_function = nn.CrossEntropyLoss()

# 日志信息保存路径
logger = SummaryWriter(log_dir='./log')


模型训练

for epoch in range(50):
train_accuracy = 0
train_loss = 0
val_accuracy = 0
val_loss = 0
epoch_acc_count = 0
count = 0
running_loss = 0

model.train()
images, labels = data
output = model(images)
loss = loss_function(output, labels)
loss.backward()
optimizer.step()

running_loss += loss.item()

# 计算每个epoch正确的个数
epoch_acc_count += (output.argmax(axis=1) == labels.view(-1)).sum()
count += len(images)

# 写入日志信息
logger.add_scalar('train_accuracy', epoch_acc_count / count, epoch + 1)

running_loss = 0
epoch_acc_count = 0
count = 0

model.eval()
images, labels = data
output = model(images)
loss = loss_function(output, labels)
loss.backward()
optimizer.step()

running_loss += loss.item()

# 计算每个epoch正确的个数
epoch_acc_count += (output.argmax(axis=1) == labels.view(-1)).sum()
count += len(images)

logger.add_scalar('val_accuracy', epoch_acc_count / count, epoch + 1)


命令启动TensorBoard

tensorboard --logdir=log_path

• 2
点赞
• 35
收藏
觉得还不错? 一键收藏
• 打赏
• 0
评论
09-16
11-30 1万+
04-07 208
03-30 927
02-23 1078
06-14 4301
04-17 4469
06-15 888
10-18 3576

“相关推荐”对你有帮助么？

• 非常没帮助
• 没帮助
• 一般
• 有帮助
• 非常有帮助

¥2 ¥4 ¥6 ¥10 ¥20

1.余额是钱包充值的虚拟货币，按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载，可以购买VIP、C币套餐、付费专栏及课程。