Pytorch中使用Tensorboard可视化训练过程

一.Tensorboard函数功能

1.add_scalar

        主要用途是将一个标量值记录到 TensorBoard 中，以便在训练或评估模型时可视化模型的性能指标或其他相关数值。

相关参数：

• tag（字符串类型）：这是一个用于标识该标量数据的标签。
• scalar_value（数值类型）：这是要记录的实际标量值。
• global_step（整数类型）：这个参数表示全局的步数或者时间步。
• walltime（可选，浮点数类型）：这个参数表示事件发生的时间。如果不提供，TensorBoard 会自动使用当前时间。

def add_scalar_demo():

    # 将信息写入logs文件夹，可以供TensorBoard消费，来可视化
    writer = SummaryWriter("logs/add_scalar")

    # 绘制 y = 2x 实例
    x = range(100)
    for i in x:
        writer.add_scalar('add_scalar实例：y=2x', i * 2, i)

    # 关闭
    writer.close()
add_scalar_demo()

2.add_graph

        主要用途是将模型的计算图结构可视化在 TensorBoard 中。

相关参数：

• model：这是要可视化的计算模型。
• input_to_model：这是一个输入数据的示例，用于展示模型在处理实际数据时的计算流程。

def add_graph_demo():
    # 将信息写入logs文件夹，可以供TensorBoard消费，来可视化
    writer = SummaryWriter("logs/add_graph")

    img = torch.rand([1, 3, 64, 64], dtype=torch.float32)
    model = torchvision.models.AlexNet(num_classes=10)
    #print(model)
    writer.add_graph(model, input_to_model=img)
    # 关闭
    writer.close()

add_graph_demo()

3.add_text

        主要用途在模型构建阶段，可以使用add_text来记录模型的详细架构信息。

相关参数：

• tag：这是一个字符串，它在 TensorBoard 中起到标识文本的作用。
• text_string：这是要记录的文本内容。
• global_step（可选）：这是一个整数，表示全局的步骤。

 

def add_text_demo():
    # 将信息写入logs文件夹，可以供TensorBoard消费，来可视化
    writer = SummaryWriter("logs/add_text")

    writer.add_text('lstm', 'This is an lstm', 0)
    writer.add_text('rnn', 'This is an rnn', 10)
    # 关闭
    writer.close()

add_text_demo()

二.使用tensorboard可视化训练过程（手写体数字识别）

1.数据集下载

train_set = datasets.MNIST("data", train=True, download=True, transform=pipeline)
test_set = datasets.MNIST("data", train=False, download=True, transform=pipeline)

train_loader = DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True)

test_loader = DataLoader(test_set, batch_size=BATCH_SIZE, shuffle=True)

2.模型搭建

class Digit(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 10, 5)
        self.conv2 = nn.Conv2d(10, 20, 3) #输入 16 * 10 * 24 * 24 输出 16 * 20 * 22 * 22
        self.fc1 = nn.Linear(20*10*10, 500)
        self.fc2 = nn.Linear(500, 10)
    def forward(self, x):
        input_size = x.size(0)
        x = self.conv1(x)                                       #输入 16 * 1 * 28 * 28 输出 16 * 10 * 24 * 24
        x = F.relu(x)
        x = F.max_pool2d(x, 2, 2)               #输入 16 * 10 * 24 * 24 输出 16 * 10 * 12 * 12
        x = self.conv2(x)                                       #输入 16 * 10 * 12 * 12 输出 16 * 20 * 10 * 10
        x = F.relu(x)

        x = x.view(input_size, -1)
        x = self.fc1(x) #输入 16 * 2000 输出 16 * 500
        x = F.relu(x)
        x = self.fc2(x) #输入 16 * 500 输出 16 * 10
        output = F.log_softmax(x,dim=1)
        return output

3.train_model 和 test_model

def train_model(model, device, train_loader, optimizer, epoch):
    model.train()
    train_loss=0
    for batch_index, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.cross_entropy(output, target)

        loss.backward()
        optimizer.step()
        train_loss=loss.item()
        if batch_index % 3000 == 0:
            print('Train Epoch : {} \t Loss : {:.6f}'.format(epoch, loss.item()))
    train_loss=train_loss/len(train_loader.dataset)
    return train_loss

def test_model(model, device, test_loader):
    model.eval()
    correct=0.0
    test_loss = 0.0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.cross_entropy(output, target).item()
            pred = output.max(1, keepdim=True)[1]
            correct += pred.eq(target.view_as(pred)).sum().item()
        test_loss /= len(test_loader.dataset)
        correct=100.0 * correct / len(test_loader.dataset)
        print('Test -- Average loss : {:.4f}, Accuracy : {:.3f}\n'.format(test_loss,
                                                                          100.0 * correct / len(test_loader.dataset)))
    return test_loss,correct

4.可视化

4.1代码部分

writer = SummaryWriter("logs/add_scalar")
model=Digit().to(DEVICE)
optimizer = optim.Adam(model.parameters())
#print(model)
writer = SummaryWriter("logs/add_graph")
#print(model)
img=torch.rand([1, 1, 28, 28], dtype=torch.float32).to(DEVICE)
writer.add_graph(model, input_to_model=img)
flops, params = profile(model, inputs=(img,))
writer.add_text('FLOPs', "flops:{}".format(flops))
writer.add_text('PARAMS', "params:{}".format(params))

print(flops, params)
writer.add_scalar('Parameters', params)
for epoch in range(1, EPOCHS + 1):
    train_loss = train_model(model, DEVICE, train_loader, optimizer, epoch)
    test_loss, correct = test_model(model, DEVICE, test_loader)
    writer.add_scalars('loss and correct', {"train_loss": train_loss, "test_loss": test_loss, "correct": correct},  epoch)
writer.close()

4.2 flops and params

4.3 model

4.4 tran_loss,test_loss and correct