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

使用tensorboard_logger记录模型在不同数据集上的性能

发布时间:2024-01-14 07:26:25

Tensorboard_logger是一个python库,可以用来记录和可视化模型的性能指标。它是Tensorboard的一个轻量级替代品,并且非常容易使用。下面是一个使用tensorboard_logger记录模型在不同数据集上性能的例子。

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

pip install tensorboard_logger

接下来,我们将使用一个简单的分类模型作为示例,并记录在训练集和测试集上的准确率。

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from tensorboard_logger import configure, log_value

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

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

# 配置tensorboard_logger
configure("logs")

# 加载MNIST数据集
train_dataset = MNIST(root="./data", train=True, transform=transforms.ToTensor(), download=True)
test_dataset = MNIST(root="./data", train=False, transform=transforms.ToTensor(), download=True)

# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64)

# 初始化模型、损失函数和优化器
model = Classifier()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

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

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

        # 每100步记录一次训练集的准确率
        if (i + 1) % 100 == 0:
            total = labels.size(0)
            _, predicted = torch.max(outputs.data, 1)
            correct = (predicted == labels).sum().item()
            accuracy = 100 * correct / total

            # 使用tensorboard_logger记录训练集的准确率
            log_value('train_accuracy', accuracy, step=(epoch * len(train_loader) + i))

    # 在测试集上评估模型性能
    total = 0
    correct = 0
    with torch.no_grad():
        for images, labels in test_loader:
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total

    # 使用tensorboard_logger记录测试集的准确率
    log_value('test_accuracy', accuracy, step=epoch)

# 关闭tensorboard_logger
configure()

在上述示例中,我们首先导入所需的库,然后定义了一个简单的分类模型。然后我们使用configure函数初始化tensorboard_logger,并指定日志输出的目录。接下来,我们加载MNIST数据集,并创建训练集和测试集的数据加载器。

然后,我们定义了模型的损失函数和优化器,并开始训练模型。在训练过程中,我们使用log_value函数记录每个epoch的训练集准确率。我们还在每个epoch结束时,在测试集上评估模型的性能,并使用log_value函数记录测试集准确率。

最后,我们使用configure()函数关闭tensorboard_logger。在训练过程中,tensorboard_logger会将记录的指标保存在指定的日志目录中。我们可以使用tensorboard命令来启动Tensorboard服务器,并在浏览器中查看记录的指标。

tensorboard --logdir=logs

在浏览器中打开http://localhost:6006即可查看可视化的指标图表。

总结来说,使用tensorboard_logger记录模型在不同数据集上的性能非常简单。我们只需要使用log_value函数记录准确率等指标,并在训练过程中定期调用该函数即可。同时,我们还可以配置tensorboard_logger的输出目录,并使用tensorboard命令来查看记录的指标图表。