基于面向对象重构模型训练器

引言

深度学习领域我们常用jupyter来演练代码，但实际生产环境中不可能像jupyter一样，所有代码逻辑都在面向过程编程，这会导致代码可复用性差，维护难度高。

前面这篇文章 基于pytorch+可视化重学线性回归模型 已经封装了数据加载器，本文我们将要对整个训练循环的逻辑进行重构，采用封装的方式来提升代码的可复用性，降低维护难度。

步骤大概是：

1. 封装小批量单次训练
2. 封装小批量单次测试
3. 封装训练循环
4. 封装损失数据的收集和可视化
5. 封装参数和梯度变化的数据可视化
6. 封装保存和加载模型

首先，导入需要的包

import torch
import numpy as np
import matplotlib.pyplot as plt
import torch.nn as nn
import torch.optim as optim

1. 数据准备

1.1 数据生成

鉴于正式工程中不会自己生成数据，所以数据生成部分始终会保持不变。

true_w = 2
true_b = 1
N = 100

np.random.seed(42)
x = np.random.rand(N, 1)
eplison = 0.1 * np.random.randn(N, 1)
y = true_w * x + true_b + eplison
x.shape, y.shape, eplison.shape

((100, 1), (100, 1), (100, 1))

1.2 数据拆分改造

将数据集转换为张量，这里将不作发送到设备to(device)的操作，而是推迟到小批量训练时再将数据发送到设备上，以节省和优化GPU显存的使用。

x_tensor = torch.as_tensor(x).float()
y_tensor = torch.as_tensor(y).float()

对于单纯的tensor数据可以直接使用pytorch内置的TensorDataset类来封装数据集, 并使用random_split来划分训练集和测试集。

from torch.utils.data import TensorDataset, DataLoader, random_split

ratio = 0.8
batch_size = 8

dataset = TensorDataset(x_tensor, y_tensor)

train_size = int(len(dataset) * ratio)
test_size = len(dataset) - train_size

train_dataset, test_dataset = random_split(dataset, [train_size, test_size])

train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=True)

len(train_dataset), len(test_dataset), next(iter(train_loader))[0].shape, next(iter(test_loader))[0].shape

(80, 20, torch.Size([8, 1]), torch.Size([8, 1]))

2. 训练器的面向对象改造

2.1 定义模型的基本组件

基本组件目前固定是模型、损失函数和优化器，为方便后续复用，这里定义一个函数来生成这些组件。

线性回归模型在pytorch中已经有封装，这里直接使用nn.Linear来代替自定义。

lr = 0.2

def make_model_components(lr):
    torch.manual_seed(42)
    model = nn.Linear(1, 1)
    lossfn = nn.MSELoss(reduction='mean')
    optimizer = optim.SGD(model.parameters(), lr=lr)
    return model, lossfn, optimizer

model, lossfn, optimizer = make_model_components(lr)
model.state_dict()

OrderedDict([('weight', tensor([[0.7645]])), ('bias', tensor([0.8300]))])

2.2 创建训练器

为了实现训练器的高内聚、低耦合，分离动态与静态，我们使用面向对象的方法对其进行重构，对于模型训练这个业务来说：

• 变化的内容应该是模型、数据源、损失函数、优化器、随机数种子等，这些内容应该由外部传入；
• 不变的内容应该是训练循环、小批量迭代训练、评估模型损失、模型的保存和加载、预测计算等，这些内容应该由内部封装。

首先，我们定义一个训练器类，它包含以下功能：

class LinearTrainer:

    def __init__(self, model, lossfn, optimizer, verbose=False):
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.model = model.to(self.device)
        self.lossfn = lossfn
        self.optimizer = optimizer
        self.verbose = verbose   # 用于调试模式打印日志
    
trainer = LinearTrainer(model, lossfn, optimizer, verbose=True)
trainer.model.state_dict(), trainer.device, trainer.lossfn, trainer.optimizer.state_dict()

(OrderedDict([('weight', tensor([[0.7645]])), ('bias', tensor([0.8300]))]),
 'cpu',
 MSELoss(),
 {'state': {},
  'param_groups': [{'lr': 0.2,
    'momentum': 0,
    'dampening': 0,
    'weight_decay': 0,
    'nesterov': False,
    'maximize': False,
    'foreach': None,
    'differentiable': False,
    'params': [0, 1]}]})

2.3 设置数据加载器

数据源是变化的，但训练逻辑其实只需要依赖符合pytorch定义的数据加载器，所以需要给训练器添加一个设置数据加载器的方法。

def set_loader(self, train_loader, test_loader=None):
    self.train_loader = train_loader
    self.test_loader = test_loader
    print(f'set train_loader: {self.train_loader}\ntest_loader: {self.test_loader}')

setattr(LinearTrainer, 'set_loader', set_loader)
trainer.set_loader(train_loader, test_loader)

set train_loader: <torch.utils.data.dataloader.DataLoader object at 0x14ad551e0>
test_loader: <torch.utils.data.dataloader.DataLoader object at 0x1150a5900>

2.4 添加单次迭代构建器

给训练器添加一个单次迭代构建器，用于构建单次迭代训练函数和单次迭代测试函数。

• build_train_step： 构建单次迭代训练函数，返回一个能够完成单次迭代训练的函数train_step。
• build_test_step： 构建单次迭代测试函数，返回一个能够完成单次迭代测试的函数test_step。

注：关于梯度清零，常规做法是放在optimizer.step()更新参数之后调用optimizer.zero_grad()，但这样一来是无法记录和观测梯度值的，给排查问题造成阻碍，所以这里将梯度清零的步骤移到下一次训练之前，目的是允许主循环获取当前梯度值。

def build_train_step(self):
    def train_step(x, y):
        # 切换模型为训练模式
        self.model.train()
        # 将梯度清零的步骤移到下一次训练之前，目的是允许主循环获取当前梯度值
        self.optimizer.zero_grad()
        # 计算预测值
        yhat = self.model(x)
        # 计算损失
        loss = self.lossfn(yhat, y)
        # 反向传播计算梯度
        loss.backward()
        # 使用优化器更新参数
        self.optimizer.step()
        return loss.item()
    return train_step

def build_test_step(self):
    def test_step(x, y):
        # 切换模型为测试模式
        self.model.eval()
        # 计算预测值
        yhat = self.model(x)
        # 计算损失
        loss = self.lossfn(yhat, y)
        return loss.item()
    return test_step

setattr(LinearTrainer, 'build_train_step', build_train_step)
setattr(LinearTrainer, 'build_test_step', build_test_step)
trainer.build_train_step(), trainer.build_test_step()

(<function __main__.build_train_step.<locals>.train_step(x, y)>,
 <function __main__.build_test_step.<locals>.test_step(x, y)>)

2.5 添加小批量迭代方法

在小批量迭代训练过程中，是训练和测试两个环节交叉进行。这两个环节的逻辑很相似，都是输入数据输出损失，不同之处在于所使用的数据加载器和单次迭代函数不同。我们可以封装一个统一的小批量迭代方法，来屏蔽这个差别。

def mini_batch(self, test=False):
    data_loader = None
    step_fn = None
    if test:
        data_loader = self.test_loader
        step_fn = self.build_test_step()
    else:
        data_loader = self.train_loader
        step_fn = self.build_train_step()

    if data_loader is None:
        raise ValueError("No data loader")

    x_batch, y_batch = next(iter(data_loader))
    x = x_batch.to(self.device)
    y = y_batch.to(self.device)
    loss = step_fn(x, y)
    return loss

setattr(LinearTrainer, "mini_batch", mini_batch)

LinearTrainer.mini_batch

<function __main__.mini_batch(self, test=False)>

2.6 设置随机数种子

为了确保结果的可复现性，我们需要为numpy和torch指定随机种子。除此之外，还需要设置cudnn的确定性

• torch.backends.cudnn.deterministic： 当设置为True时，这个选项会确保cuDNN算法是确定性的，对于相同的输入和配置，它们将总是产生相同的输出。但是此选项可能会降低性能。
• torch.backends.cudnn.benchmark：当设置为True时，cuDNN将会花费一些时间来“基准测试”各种可能的算法，并选择一个最快的。而设置为False时，则始终使用一种确定的算法，常和deterministic配合使用。

def set_seed(self, seed):
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False    

setattr(LinearTrainer, 'set_seed', set_seed)
trainer.set_seed(42)

2.7 添加主训练方法

此方法主要完成一个循环迭代的训练过程，每次迭代都会执行一个小批量训练和小批量测试，并实时收集训练损失和测试损失用于观察，迭代的次数由参数epoch_n决定。

def train(self, epoch_n):
    self.train_losses = []
    self.test_losses = []

    for i in range(epoch_n):
        loss = self.mini_batch(test=False)
        self.train_losses.append(loss)

        with torch.no_grad():
            test_loss = self.mini_batch(test=True)
            self.test_losses.append(test_loss)
    print(f'train loss: {self.train_losses[-1]}')
        
setattr(LinearTrainer, 'train', train)
trainer.train(100)
trainer.model.state_dict()

train loss: 0.006767683196812868

2.8 显示损失曲线

将生成的损失数据用matplotlib显示出来，以观察训练和测试两条损失曲线是否随着迭代次数而稳定下降。

def show_losses(self):
    fig, ax = plt.subplots(1, 1, figsize=(6, 4))
    ax.plot(self.train_losses, label='train losses', color='blue')
    ax.plot(self.test_losses, label='test losses', color='red')
    ax.legend(loc='upper right')
    ax.set_title('Loss descent')
    ax.set_xlabel('epochs')
    ax.set_ylabel('loss')
    ax.set_yscale('log')

    plt.show()

setattr(LinearTrainer, 'show_losses', show_losses)
trainer.show_losses()

我们每次为了可视化数据，都要手动记录损失数据，并手动写函数来绘制损失曲线。是否有更简单的方法呢？答案是有的，那就是tensorboard。

3. tensorboard

Tensorboard 是一个来自Tensorflow的可视化工具，但pytorch也提供了类和方法集成和使用它，可见它有多么的好用。

手动收集数据过于麻烦，而且每次画图都要写一个绘图函数，而且数据量和参数很多的时候，将需要写很多函数。

3.1 Tensorboard的基本使用

Tensorboard的使用分为两个部分：收集数据和显示数据。

• 收集数据：主要靠SummaryWriter类，集成到pytorch中使用。
• 显示数据：主要靠tensorboard命令，类似jupyter一样启动一个服务，然后通过浏览器访问。

SummaryWriter类提供了很多常用的方法来收集数据：
- add_graph: 收集模型的网络结构。
- add_scalar/add_scalars：收集标量数据，像损失函数值，准确率等。
- add_image/add_images：收集图片数据，像输入图片，输出图片等。
- add_text: 收集文本数据，可以记录一些文字。
- add_histogram: 收集直方图数据，可以用来观察参数分布。
- add_video: 收集视频数据，可以用来观察训练过程。
- add_embedding: 收集嵌入数据，可以用来观察数据分布。
- add_audio: 收集音频数据，可以用来观察训练过程。

from torch.utils.tensorboard import SummaryWriter
# 告诉tensorboard，要将日志记录到哪个文件夹
writer = SummaryWriter("../log/tensorboard_test")
# 取一个样例数据，连同model一起传给add_graph函数，它将能够从这个样例数据的预测过程中，收集到模型的计算图
x_sample, y_sample = next(iter(train_loader))
writer.add_graph(model, input_to_model=x_sample)

将模型的损失收集到tensorboard中，add_scalars可以将多组数据添加到一个图表中（训练和测试同图显示），而add_scalar只适用于一个图标一组数据的情况。

for i in range(len(trainer.train_losses)):
    writer.add_scalars("loss", {"train": trainer.train_losses[i], "test": trainer.test_losses[i]}, i)

运行Tensorboard，这里有两条命令：

1. 第一条命令：是用于为jupyter notebook加载tensorboard扩展。
2. 第二条命令：将在6006端口上启动一个服务器，并自动在当前jupyter notebook中内嵌一个网页来访问此服务。

%load_ext tensorboard
# 告诉tensorboard在logdir指定的文件夹中查找日志
%tensorboard --logdir "../log/tensorboard_test" --port 6006

The tensorboard extension is already loaded. To reload it, use:
  %reload_ext tensorboard



Reusing TensorBoard on port 6006 (pid 15193), started 0:00:04 ago. (Use '!kill 15193' to kill it.)

这个图上是可以点击进行操作的，可以在graphs、scalars、histoogam、images页签间切换。

3.2 使用tensorboard来改造训练器

tensorboard将收集数据与显示数据的工作分离，这样我们就不用等到训练完再查看数据，可以训练模型时，单开一个任务来可视化观察训练过程。

首先，我们需要一个设置tensorboard的方法，将SummaryWriter内置到训练器中，这样我们就可以在训练过程中收集数据了。

import os
import shutil

def set_tensorboard(self, name, log_dir, clear=True):
    log_file_path = f"{log_dir}/{name}"
        # 删除训练日志
    if clear == True and os.path.exists(log_file_path):
        shutil.rmtree(log_file_path)
        print(f"clear tensorboard path: {log_file_path}") if self.verbose else None

    self.writer = SummaryWriter(log_file_path)
    if hasattr(self, "train_loader") and self.train_loader is not None:
        sample_x, _ = next(iter(self.train_loader))
        self.writer.add_graph(self.model, sample_x)

    print(f"Tensorboard log dir: {self.writer.log_dir}") if self.verbose else None

setattr(LinearTrainer, "set_tensorboard", set_tensorboard)

具体收集的数据，除了之前的损失值外，我们还有必要收集参数的值和梯度，这对于排查损失不下降的原因很有帮助。

为避免收集数据的代码污染主循环，我们单独封装两个方法用来收集数据，分别是：

• record_train_data: 收集训练数据的主方法，包括收集数据和执行flush操作。
• record_parameters: 专门用于收集参数的方法，包括参数值本身和梯度。

def record_parameters(self, epoch_idx):
     for name, param in self.model.named_parameters():
          self.writer.add_scalar(name, param.data, epoch_idx)
          if param.grad is not None:
               self.writer.add_scalar(name+"/grad", param.grad.item(), epoch_idx)
          if self.verbose:
               print(f"epoch_idx={epoch_idx}, name={name}, param.data: {param.data}, param.grad.item: {param.grad.item() if param.grad is not None else 'None'}")

def record_train_data(self, train_loss, test_loss, epoch_idx):
     # 记录损失数据，训练损失和验证损失对比显示
     self.writer.add_scalars('loss', {'train': train_loss, 'test': test_loss}, epoch_idx)
     if self.verbose:
          print(f"epoch_idx={epoch_idx}, train_loss: {train_loss}, test_loss: {test_loss}")
     
     # 记录模型的所有参数变化，以及参数梯度的变化过程
     self.record_parameters(epoch_idx)
     self.writer.flush()

setattr(LinearTrainer, 'record_parameters', record_parameters)
setattr(LinearTrainer, 'record_train_data', record_train_data)

是时候改造训练主循环了，我们将收集数据的操作统一放到record_train_data()这个函数调用来完成，主循环反而变得更简单清晰：

注：在训练之前，先收集原始参数值，是为了保证原始参数值也被收集，并在图表中显示出来。

def train(self, eporch_n):
    # 收集原始参数
    self.record_parameters(0)
    # 开始训练
    for i in range(eporch_n):
        train_loss = self.mini_batch(test=False)

        with torch.no_grad():
            test_loss = self.mini_batch(test=True)
        # 记录训练数据
        self.record_train_data(train_loss, test_loss, i+1)
        
 
setattr(LinearTrainer, 'train', train)

由于刚才已经训练过一次，所以需要重置下模型的参数，从头开始训练并收集中间过程中的数据。

注：考虑到训练是反复进行的，为了后续方便，封装一个reset函数来重置模型，主要功能是将模型和优化器重置，并删除旧的训练日志。

import shutil
import os

def reset(self, model, lossfn, optimizer):
    if hasattr(self, "model"):
        self.model.cpu() if self.model != None else None
        del self.model
        del self.optimizer

    self.model = model
    self.lossfn = lossfn
    self.optimizer = optimizer
    print(f"reset model and optimizer: {self.model.state_dict()}, {self.optimizer.state_dict()}") if self.verbose else None

setattr(LinearTrainer, "reset", reset)

model, lossfn, optimizer = make_model_components(lr)
trainer.reset(model, lossfn, optimizer)
trainer.set_seed(42)
trainer.set_tensorboard(name="linear_objected-1", log_dir="../log")
trainer.model.state_dict()

OrderedDict([('weight', tensor([[0.7645]])), ('bias', tensor([0.8300]))])

可以看到，经过重置后，参数又恢复了原始值，下面调用train方法重新开始训练。

trainer.train(100)
trainer.model.state_dict()

OrderedDict([('weight', tensor([[1.8748]])), ('bias', tensor([1.0477]))])

# %load_ext tensorboard
%tensorboard --logdir "../log/linear_objected-1" --port 6007

Reusing TensorBoard on port 6007 (pid 26888), started 0:00:03 ago. (Use '!kill 26888' to kill it.)

4. 保存和加载模型

我们这个场景使用的是最简单的线性回归模型，所以保存和加载模型非常快。但实际中，我们可能使用更复杂的模型，这些模型可能包含很多层，每层参数都可能非常多，整个训练过程可能需要几个小时甚至几天，所以保存训练结果就显得非常重要了。

4.1 保存模型

保存模型本质上是保存模型的状态，包括模型参数、优化器状态、损失值等。这些数据都保存包裹到一个dict中，然后使用torhc.save()函数保存到文件中。

def save_checkpoint(self, checkpoint_path):
    checkpoint = {
        "model_state_dict": self.model.state_dict(),
        "optimizer_state_dict": self.optimizer.state_dict(),
    }
    torch.save(checkpoint, checkpoint_path)
    print(f"save checkpoint: {self.model.state_dict()}") if self.verbose else None


setattr(LinearTrainer, "save_checkpoint", save_checkpoint)

checkpoint_path = "../checkpoint/torch_linear-1.pth"
trainer.save_checkpoint(checkpoint_path)

4.2 加载模型

当我们需要部署模型进行数据预测，或者重新开始未完成的训练时，就需要使用torch.load()将之前保存在文件中的模型和参数加载进来。

def load_checkpoint(self, checkpoint_path):
    checkpoint = torch.load(checkpoint_path)
    self.model.load_state_dict(checkpoint["model_state_dict"])
    self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    print(f"load checkpoint: {self.model.state_dict()}") if self.verbose else None

setattr(LinearTrainer, "load_checkpoint", load_checkpoint)

  device
  train_loader
  test_loader
  train_losses
  test_losses
  writer
  debug
  optimizer

为了与前面的训练结果完全隔离开，我们重新创建一个训练器，一个新训练器需要进行的初始化总共包括以下几项：

1. 模型、损失函数和优化器
2. 训练数据集和测试数据集的加载器
3. 随机数种子
4. 设置训练数据的收集位置，便于tensorboard可视化

model, lossfn, optimizer = make_model_components(lr)
trainer2 = LinearTrainer(model, lossfn, optimizer, verbose=True)
trainer2.set_seed(42)
trainer2.set_loader(train_loader, test_loader)
trainer2.set_tensorboard(name='linear_objected-2', log_dir="../log")

set train_loader: <torch.utils.data.dataloader.DataLoader object at 0x14ad551e0>
test_loader: <torch.utils.data.dataloader.DataLoader object at 0x1150a5900>
Tensorboard log dir: ../log/linear_objected-2

然后从checkpoint加载模型参数，可以看到之前的训练结果已经加载进新的训练器。

print(f"before load: {trainer2.model.state_dict()}")
trainer2.load_checkpoint(checkpoint_path)
print(f"after load: {trainer2.model.state_dict()}")

before load: OrderedDict([('weight', tensor([[0.7645]])), ('bias', tensor([0.8300]))])
load checkpoint: OrderedDict([('weight', tensor([[1.8748]])), ('bias', tensor([1.0477]))])
after load: OrderedDict([('weight', tensor([[1.8748]])), ('bias', tensor([1.0477]))])

接着之前的训练结果继续训练

trainer2.train(100)

%tensorboard --logdir "../log/linear_objected-2" --port 6009

Reusing TensorBoard on port 6009 (pid 32774), started 0:00:03 ago. (Use '!kill 32774' to kill it.)

可以看到，经过又一轮的训练后，权重weight从1.87学习到了1.9159，离真实值2更接近了。

5. 训练器封装结果

到目前为止，给训练器添加的所有方法汇总如下：

for key, value in vars(LinearTrainer).items():
    if callable(value) and not key.startswith("__"):  # 忽略内置或特殊方法
        print(f"  {key}()")

  set_loader()
  build_train_step()
  build_test_step()
  mini_batch()
  set_seed()
  set_tensorboard()
  train()
  reset()
  record_parameters()
  record_train_data()
  save_checkpoint()
  load_checkpoint()

给训练器添加的所有字段汇总如下：

for key, value in vars(trainer).items():
    if not callable(value) and not key.startswith("__"):  # 忽略内置或特殊方法
        print(f"  {key}")

  device
  verbose
  train_loader
  test_loader
  writer
  optimizer

这些后面在方法中添加的字段，由于初始化的顺序不同，很容易引发AttributeError: object has no attribute 'xxx'，所以需要对__init__方法进行改造，以便对这些字段提前初始化。

def __init__(self, model, lossfn, optimizer):
    self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
    self.model = model
    self.lossfn = lossfn
    self.optimizer = optimizer
    self.verbose = False
    self.writer = None
    self.train_loader = None
    self.test_loader = None

setattr(LinearTrainer, '__init__', __init__)
test_trainer = LinearTrainer(model, lossfn, optimizer)
test_trainer.writer

通过初始化的改造后，上面新创建的test_trainer虽然没有调用set_tensorboard，但是仍然可以访问.writer字段而不报错。

最后LinearTrainer类的完整代码：

import os
import shutil
import torch
import numpy as np

class LinearTrainer:

    def __init__(self, model, lossfn, optimizer, verbose=False):
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.model = model.to(self.device)
        self.lossfn = lossfn
        self.optimizer = optimizer
        self.verbose = False
        self.writer = None
        self.train_loader = None
        self.test_loader = None

    def set_loader(self, train_loader, test_loader=None):
        self.train_loader = train_loader
        self.test_loader = test_loader
        print(f'set train_loader: {self.train_loader}\ntest_loader: {self.test_loader}') if self.verbose else None

    def build_train_step(self):
        def train_step(x, y):
            # 切换模型为训练模式
            self.model.train()
            # 将梯度清零的步骤移到下一次训练之前，目的是允许主循环获取当前梯度值
            self.optimizer.zero_grad()
            # 计算预测值
            yhat = self.model(x)
            # 计算损失
            loss = self.lossfn(yhat, y)
            # 反向传播计算梯度
            loss.backward()
            # 使用优化器更新参数
            self.optimizer.step()
            return loss.item()
        return train_step

    def build_test_step(self):
        def test_step(x, y):
            # 切换模型为测试模式
            self.model.eval()
            # 计算预测值
            yhat = self.model(x)
            # 计算损失
            loss = self.lossfn(yhat, y)
            return loss.item()
        return test_step
    
    def mini_batch(self, test=False):
        data_loader = None
        step_fn = None
        if test:
            data_loader = self.test_loader
            step_fn = self.build_test_step()
        else:
            data_loader = self.train_loader
            step_fn = self.build_train_step()

        if data_loader is None:
            raise ValueError("No data loader")

        x_batch, y_batch = next(iter(data_loader))
        x = x_batch.to(self.device)
        y = y_batch.to(self.device)
        loss = step_fn(x, y)
        return loss
    
    def train(self, eporch_n):
        # 收集原始参数
        self.record_parameters(0)
        # 开始训练
        for i in range(eporch_n):
            train_loss = self.mini_batch(test=False)

            with torch.no_grad():
                test_loss = self.mini_batch(test=True)
            # 记录训练数据
            self.record_train_data(train_loss, test_loss, i+1)
    
    def set_seed(self, seed):
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False    

    def set_tensorboard(self, name, log_dir, clear=True):
        log_file_path = f"{log_dir}/{name}"
            # 删除训练日志
        if clear == True and os.path.exists(log_file_path):
            shutil.rmtree(log_file_path)
            print(f"clear tensorboard path: {log_file_path}") if self.verbose else None

        self.writer = SummaryWriter(log_file_path)
        if hasattr(self, "train_loader") and self.train_loader is not None:
            sample_x, _ = next(iter(self.train_loader))
            self.writer.add_graph(self.model, sample_x)

        print(f"Tensorboard log dir: {self.writer.log_dir}") if self.verbose else None

    def record_parameters(self, epoch_idx):
        for name, param in self.model.named_parameters():
            self.writer.add_scalar(name, param.data, epoch_idx)
            if param.grad is not None:
                self.writer.add_scalar(name+"/grad", param.grad.item(), epoch_idx)
            if self.verbose:
                print(f"epoch_idx={epoch_idx}, name={name}, param.data: {param.data}, param.grad.item: {param.grad.item() if param.grad is not None else 'None'}")

    def record_train_data(self, train_loss, test_loss, epoch_idx):
        # 记录损失数据，训练损失和验证损失对比显示
        self.writer.add_scalars('loss', {'train': train_loss, 'test': test_loss}, epoch_idx)
        if self.verbose:
            print(f"epoch_idx={epoch_idx}, train_loss: {train_loss}, test_loss: {test_loss}")
        
        # 记录模型的所有参数变化，以及参数梯度的变化过程
        self.record_parameters(epoch_idx)
        self.writer.flush()

    def save_checkpoint(self, checkpoint_path):
        checkpoint = {
            "model_state_dict": self.model.state_dict(),
            "optimizer_state_dict": self.optimizer.state_dict(),
        }
        torch.save(checkpoint, checkpoint_path)
        print(f"save checkpoint: {self.model.state_dict()}") if self.verbose else None

    def load_checkpoint(self, checkpoint_path):
        checkpoint = torch.load(checkpoint_path)
        self.model.load_state_dict(checkpoint["model_state_dict"])
        self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        print(f"load checkpoint: {self.model.state_dict()}") if self.verbose else None

    def reset(self, model, lossfn, optimizer):
        if hasattr(self, "model"):
            self.model.cpu() if self.model != None else None
            del self.model
            del self.optimizer

        self.model = model
        self.lossfn = lossfn
        self.optimizer = optimizer
        print(f"reset model and optimizer: {self.model.state_dict()}, {self.optimizer.state_dict()}") if self.verbose else None

参考资料

• 基于pytorch+可视化重学线性回归模型
• 基于numpy演练可视化梯度下降

原文地址：https://blog.csdn.net/xiaojia1001/article/details/140650210

免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！