欢迎访问宙启技术站
智能推送

使用tensorboard_logger记录训练过程中的损失和准确率

发布时间:2024-01-14 07:18:10

TensorBoard是TensorFlow提供的一个可视化工具,可以帮助我们实时监控、可视化训练过程中的各类指标,比如损失、准确率、梯度等。

在TensorFlow中,我们可以使用tensorboard_logger库将训练过程中的损失和准确率等信息保存到TensorBoard的日志文件中,以便后续进行可视化分析。

下面我们通过一个示例来演示如何使用tensorboard_logger记录训练过程中的损失和准确率。

首先,我们需要安装tensorboard_logger库。可以通过以下命令进行安装:

pip install tensorboard_logger

我们首先定义一个简单的分类模型(使用了PyTorch):

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

# 定义模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(784, 10)

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

model = SimpleModel()

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

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

# 定义TensorBoard日志路径
log_dir = './logs'
writer = SummaryWriter(log_dir)

# 训练模型
epochs = 5
for epoch in range(epochs):
    for i, (inputs, targets) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        # 记录损失和准确率到TensorBoard
        writer.add_scalar('Loss/train', loss.item(), epoch * len(train_loader) + i)
        _, predicted = outputs.max(1)
        correct = predicted.eq(targets).sum().item()
        accuracy = correct / targets.size(0)
        writer.add_scalar('Accuracy/train', accuracy, epoch * len(train_loader) + i)
        
    print(f'Epoch: {epoch+1}/{epochs}, Loss: {loss.item()}, Accuracy: {accuracy}')

在上述示例中,我们首先导入了需要的库,并定义了一个简单的分类模型SimpleModel。然后,我们加载了MNIST数据集,定义了优化器和损失函数。

接下来,我们使用SummaryWriter创建了一个TensorBoard的日志写入器,指定了日志的保存路径log_dir

在每个epoch的训练过程中,我们通过在相关位置调用writer.add_scalar函数,将损失和准确率信息写入到TensorBoard日志文件中。其中,每个损失和准确率的标签(tag)使用了"Loss/train""Accuracy/train"

最后,我们通过调用writer.close()函数关闭了TensorBoard的日志写入器。

在命令行中,可以通过以下命令启动TensorBoard服务器:

tensorboard --logdir=./logs

然后,在浏览器中访问http://localhost:6006/,就可以看到TensorBoard的可视化界面了。

在可视化界面中,可以切换到"Scalars"选项卡,就能看到记录的损失和准确率曲线了。

以上就是使用tensorboard_logger记录训练过程中的损失和准确率的一个例子。通过使用tensorboard_logger,我们可以方便地将训练过程中的各类指标记录到TensorBoard的日志文件中,并通过TensorBoard进行可视化分析,以帮助我们更全面地了解和优化模型的训练过程。