1. 代码实现

波士顿房价数据集下载

1.1 一元线性回归模型的训练

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader, random_split
from tensorboardX import SummaryWriter

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
batch_size = 2
num_epochs = 200

writer = SummaryWriter()

model = nn.Linear(1, 1).to(device)
nn.init.normal_(model.weight, mean=0, std=0.01)
nn.init.constant_(model.bias, 0)
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4)

data = np.load('dataset/boston_housing/boston_housing.npz')
X = torch.tensor(data['x'][:, 0].reshape(-1, len(model.weight)), dtype=torch.float, device=device)
y = torch.tensor(data['y'].reshape(-1, 1), dtype=torch.float, device=device)
dataset = TensorDataset(X, y)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

for epoch in range(num_epochs):
    for _X, _y in dataloader:
        _X, _y = _X.to(device), _y.to(device)
        loss = criterion(model(_X), _y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    loss = criterion(model(X), y)
    torch.save(model.state_dict(), 'model/linearRegression.pt')
    model.load_state_dict(torch.load('model/linearRegression.pt'))
    writer.add_scalar('Loss/train', loss, epoch)
    writer.add_scalar('W/train', model.weight, epoch)
    writer.add_scalar('b/train', model.bias, epoch)
writer.close()

2. 代码解读

2.1. tensorboardX

tensorboardX是一种能将训练过程可视化的工具

2.1.1. tensorboardX的安装

安装命令：

pip install tensorboardX

VSCode集成了TensorBoard支持，不过事先要安装torch-tb-profiler，安装命令：

pip install torch-tb-profiler

安装完成后，在Python源文件中tensorboardX模块导入处，点击“启动TensorBoard会话”按钮，然后选择运行事件所在目录，默认选择当前目录即可，tensorboard会自动在当前目录查找运行事件，由此即可启动TensorBoard。


此外，也可以通过以下命令在浏览器查看tensorboard可视化结果：

# logdir为运行事件所在目录
> tensorboard logdir=runs
TensorFlow installation not found - running with reduced feature set.
I1202 20:37:50.824767 15412 plugin.py:429] Monitor runs begin
Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all
TensorBoard 2.14.0 at http://localhost:6006/ (Press CTRL+C to quit)
# 手动打开命令输出提供的本地服务器地址，如http://localhost:6006/

2.1.2. tensorboardX的使用

• 直接创建对象

from tensorboardX import SummaryWriter
writer = SummaryWriter()
# writer.add_scalar()：添加监控变量
writer.close()

• 使用上下文管理器

from tensorboardX import SummaryWriter
with SummaryWriter() as writer:
	# writer.add_scalar()：添加监控变量