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

TensorBoard中利用SummaryWriter()函数生成训练过程中的准确率曲线

发布时间:2023-12-25 00:00:17

在TensorBoard中,我们可以使用SummaryWriter()函数来生成训练过程中的准确率曲线。具体的步骤如下:

首先,我们需要导入必要的库:

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter

接下来,我们可以定义一个简单的网络模型,并定义损失函数和优化器:

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
    
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

net = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

然后,我们可以定义一个SummaryWriter对象,并将其与指定的日志目录相关联:

writer = SummaryWriter('./logs')

在训练过程中,每隔一定的步骤我们可以将准确率的信息写入SummaryWriter对象中:

running_loss = 0.0
total_step = len(trainloader)
for epoch in range(num_epochs):
    for i, (inputs, labels) in enumerate(trainloader, 0):
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()
        if (i+1) % 2000 == 0:
            # 计算准确率
            with torch.no_grad():
                correct = 0
                total = 0
                for images, labels in testloader:
                    outputs = net(images)
                    _, predicted = torch.max(outputs.data, 1)
                    total += labels.size(0)
                    correct += (predicted == labels).sum().item()
                    
            accuracy = 100 * correct / total
            writer.add_scalar('Accuracy', accuracy, epoch * total_step + i)
            
            print('[%d, %5d] loss: %.3f' %
                (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0

最后,我们需要在训练结束后关闭SummaryWriter对象:

writer.close()

完成以上步骤后,我们可以使用以下命令在终端中启动TensorBoard:

tensorboard --logdir=./logs

然后,我们可以在浏览器中打开生成的链接,即可看到准确率曲线图。

这是一个简单的使用TensorBoard生成训练过程准确率曲线的例子。通过SummaryWriter函数,我们可以方便地将训练过程中的诸多指标记录下来,并在TensorBoard中进行可视化和分析。这对于深度学习模型的训练和调优非常有帮助。