TensorBoard中使用SummaryWriter()函数进行模型保存和重载

TensorBoard 提供了一种可视化工具，用于解释和监测神经网络模型。可以使用 SummaryWriter() 函数来保存和重载模型。

以下是使用 SummaryWriter() 进行模型保存和重载的示例代码：

1. 导入相关的库：

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

2. 创建模型：

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(5, 2)

    def forward(self, x):
        x = self.fc1(x)
        x = torch.relu(x)
        x = self.fc2(x)
        return x

model = MyModel()

3. 定义损失函数和优化器：

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)

4. 创建一个 SummaryWriter 对象：

writer = SummaryWriter()

5. 训练模型并保存相关信息：

for epoch in range(10):
    # 前向传播
    inputs = torch.randn(32, 10)
    outputs = model(inputs)
    
    # 计算损失
    targets = torch.randint(0, 2, (32,))
    loss = criterion(outputs, targets)
    
    # 反向传播和优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    # 保存相关信息到 TensorBoard
    writer.add_scalar('Loss/train', loss.item(), epoch)
    writer.add_histogram('FC1/weights', model.fc1.weight, epoch)
    writer.add_histogram('FC1/biases', model.fc1.bias, epoch)

6. 关闭 SummaryWriter 对象：

writer.close()

在上述示例中，我们使用 writer.add_scalar() 函数来记录训练过程中的损失，使用 writer.add_histogram() 函数来记录模型中某些参数的分布情况。这些信息将会被写入到 TensorBoard 的日志文件中。

保存的日志文件可以使用以下命令在终端中启动 TensorBoard 进行可视化：

tensorboard --logdir=./logs

这将在浏览器中打开一个地址，显示模型的训练过程和其他相关信息。

要加载之前保存的模型并进行推理，可以使用以下代码：

# 加载模型
checkpoint = torch.load('path/to/model_checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])

# 运行推理
input_data = ...
output = model(input_data)

以上是使用 SummaryWriter() 函数进行模型保存和重载的例子。通过将相关信息记录到 TensorBoard，我们可以方便地监测和分析神经网络模型的训练过程。