博客

在Pytorch中简单使用tensorboard

  • Post category:Python

以下是在PyTorch中简单使用TensorBoard的完整攻略，包括两个示例。

在PyTorch中使用TensorBoard的基本步骤

使用TensorBoard的基本步骤如下：

  1. 安装TensorBoard

使用以下命令安装TensorBoard：

pip install tensorboard
  1. 导入TensorBoard

在PyTorch中，可以使用torch.utils.tensorboard模块导入TensorBoard。

from torch.utils.tensorboard import SummaryWriter
  1. 创建SummaryWriter对象

使用SummaryWriter类创建一个SummaryWriter对象，指定TensorBoard的日志目录。

writer = SummaryWriter('logs')
  1. 记录数据

使用add_scalar方法记录标量数据，使用add_histogram方法记录张量数据。

# 记录标量数据
writer.add_scalar('loss', loss, epoch)

# 记录张量数据
writer.add_histogram('weights', model.weights, epoch)
  1. 启TensorBoard

使用以下命令启动TensorBoard：

tensorboard --logdir=logs
  1. 查看TensorBoard

在浏览器中打开http://localhost:6006，即可查看TensorBoard。

以下是两个在PyTorch中使用TensorBoard的示例。

示例一：记录训练损失

以下是记录训练损失的示例代码：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
from torch.utils.tensorboard import SummaryWriter

# 创建SummaryWriter对象
writer = SummaryWriter('logs')

# 加载MNIST数据集
train_dataset = MNIST(root='data', train=True, transform=ToTensor(), download=True)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

# 定义模型
model = nn.Sequential(
    nn.Linear(784, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 训练模型
for epoch in range(10):
    for i, (images, labels) in enumerate(train_loader):
        # 前向传播
        images = images.view(-1, 784)
        outputs = model(images)
        loss = criterion(outputs, labels)

        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # 记录训练损失
        writer.add_scalar('train_loss', loss.item(), epoch * len(train_loader) + i)

# 关闭SummaryWriter对象
writer.close()

上面的代码使用MNIST数据集训练一个简单的神经网络模型，并使用add_scalar方法记录训练损失。最后使用close方法关闭SummaryWriter对象。

示例二：记录模型权重

以下是记录模型权重的示例代码：

import torch
import torch.nn as nn
from torch.utils.tensorboard import SummaryWriter

# 创建SummaryWriter对象
writer = SummaryWriter('logs')

# 定义模型
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.fc1 = nn.Linear(784, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = x.view(-1,784)
        x = self.fc1(x)
        x = nn.ReLU()(x)
        x = self.fc2(x)
        return x

model = Model()

# 记录模型权重
for name, param in model.named_parameters():
    writer.add_histogram(name, param, 0)

# 关闭SummaryWriter对象
writer.close()

上面的代码定义了一个简单的神经网络模型，并使用add_histogram方法记录模型权重。最后使用close方法关闭SummaryWriter对象。

以上是在PyTorch中简单使用TensorBoard的完整攻略，通过以上步骤和示例，我们可以轻松地使用TensorBoard可视化训练过程和模型。

Tags: