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

使用tensorboard_logger在Python中实现模型在不同批次大小下的训练过程可视化

发布时间:2024-01-09 09:33:48

TensorBoard是TensorFlow的一个可视化工具,可以帮助我们更好地理解和调试模型的训练过程。Tensorboard Logger是一个用于将训练过程信息写入TensorBoard的Python库。

下面是一个使用Tensorboard Logger的示例,展示了模型在不同批次大小下的训练过程可视化:

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

# 配置Tensorboard
configure("logs", flush_secs=5)

# 定义模型
class Net(nn.Module):
    def __init__(self):
        super(Net, 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

# 定义数据转换
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# 加载数据集
train_set = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_set = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)

# 定义数据加载器
train_loader = DataLoader(train_set, batch_size=64, shuffle=True)
test_loader = DataLoader(test_set, batch_size=64, shuffle=False)

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

# 训练模型
for epoch in range(10):
    total_loss = 0.0

    for i, data in enumerate(train_loader):
        images, labels = data

        # 前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)

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

        total_loss += loss.item()

        # 记录训练过程信息到TensorBoard
        log_value('loss', loss.item(), epoch * len(train_loader) + i + 1)
        
    # 打印每个epoch的平均损失
    print('Epoch {}/{} Avg Loss: {:.4f}'.format(epoch + 1, 10, total_loss / len(train_loader)))

# 测试模型
correct = 0
total = 0

with torch.no_grad():
    for data in test_loader:
        images, labels = data
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

# 打印准确率
print('Accuracy on test set: {:.2%}'.format(correct / total))

在上述示例中,首先我们使用configure()函数设置TensorBoard日志的保存路径。然后,我们定义了一个简单的神经网络模型(包含一个全连接层),并初始化模型、损失函数和优化器。

然后,我们使用一个嵌套的循环来进行模型的训练。外层循环控制训练的轮数(epoch),内层循环遍历每个训练批次。在每个批次中,我们先执行模型的前向传播,然后计算损失,并进行反向传播和参数优化。在每个批次的末尾,我们使用log_value()函数将当前批次的损失值写入TensorBoard日志中,并通过epoch * len(train_loader) + i + 1计算全局迭代次数。

在训练完成后,我们使用torch.no_grad()上下文来关闭梯度计算,然后对测试集进行预测,并统计准确率。最后,我们将准确率打印出来。

通过运行上述代码,我们可以在TensorBoard中查看训练过程中loss的变化情况,以及训练和测试准确率的变化情况。