pytorch 连接tensorboard

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

# 简单的线性模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(1, 1)

    def forward(self, x):
        return self.fc(x)

# 数据生成
def generate_data():
    x = torch.linspace(-1, 1, 100).unsqueeze(1)  # 输入特征
    y = 2 * x + 1 + 0.2 * torch.randn_like(x)  # 目标带噪声
    return x, y

# 设置
model = SimpleModel()
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.1)

# 创建 TensorBoard 的 SummaryWriter
writer = SummaryWriter(log_dir="/root/tf-logs")

# 数据
x_train, y_train = generate_data()

# 训练过程
for epoch in range(100):
    model.train()

    # 前向传播
    y_pred = model(x_train)
    loss = criterion(y_pred, y_train)

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

    # 将损失写入 TensorBoard
    writer.add_scalar('Loss/train', loss.item(), epoch)

    # 可视化模型权重和偏置
    for name, param in model.named_parameters():
        writer.add_histogram(f'{name}_grad', param.grad, epoch)
        writer.add_histogram(f'{name}_data', param.data, epoch)

    # 打印训练信息
    if (epoch + 1) % 10 == 0:
        print(f"Epoch [{epoch + 1}/100], Loss: {loss.item():.4f}")

# 保存模型并记录到 TensorBoard
torch.save(model.state_dict(), "simple_model.pth")
writer.add_text("Model", "Model training complete and saved.")
writer.close()

print("训练完成！使用 `tensorboard --logdir=runs` 启动 TensorBoard 查看结果。")

测试效果