引言
深度学习领域我们常用jupyter来演练代码，但实际生产环境中不可能像jupyter一样，所有代码逻辑都在面向过程编程，这会导致代码可复用性差，维护难度高。
前面这篇文章 基于pytorch+可视化重学线性回归模型 已经封装了数据加载器，本文我们将要对整个训练循环的逻辑进行重构，采用封装的方式来提升代码的可复用性，降低维护难度。
步骤大概是：
- 封装小批量单次训练
- 封装小批量单次测试
- 封装训练循环
- 封装损失数据的收集和可视化
- 封装参数和梯度变化的数据可视化
- 封装保存和加载模型
首先，导入需要的包
import torch
import numpy as np
import matplotlib.pyplot as plt
import torch.nn as nn
import torch.optim as optim1. 数据准备
1.1 数据生成
鉴于正式工程中不会自己生成数据，所以数据生成部分始终会保持不变。
true_w = 2
true_b = 1
N = 100np.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_splitratio = 0.8
batch_size = 8dataset = TensorDataset(x_tensor, y_tensor)train_size = int(len(dataset) * ratio)
test_size = len(dataset) - train_sizetrain_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.2def 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, optimizermodel, 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 = lossfnself.optimizer = optimizerself.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_loaderself.test_loader = test_loaderprint(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_stepdef 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_stepsetattr(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 = Nonestep_fn = Noneif test:data_loader = self.test_loaderstep_fn = self.build_test_step()else:data_loader = self.train_loaderstep_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 losssetattr(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 = Truetorch.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，这里有两条命令：
- 第一条命令：是用于为jupyter notebook加载tensorboard扩展。
- 第二条命令：将在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 tensorboardReusing 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 shutildef 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 Noneself.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 Nonesetattr(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 osdef reset(self, model, lossfn, optimizer):if hasattr(self, "model"):self.model.cpu() if self.model != None else Nonedel self.modeldel self.optimizerself.model = modelself.lossfn = lossfnself.optimizer = optimizerprint(f"reset model and optimizer: {self.model.state_dict()}, {self.optimizer.state_dict()}") if self.verbose else Nonesetattr(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 Nonesetattr(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 Nonesetattr(LinearTrainer, "load_checkpoint", load_checkpoint)
devicetrain_loadertest_loadertrain_lossestest_losseswriterdebugoptimizer
为了与前面的训练结果完全隔离开，我们重新创建一个训练器，一个新训练器需要进行的初始化总共包括以下几项：
- 模型、损失函数和优化器
- 训练数据集和测试数据集的加载器
- 随机数种子
- 设置训练数据的收集位置，便于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}")
deviceverbosetrain_loadertest_loaderwriteroptimizer
这些后面在方法中添加的字段，由于初始化的顺序不同，很容易引发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 = modelself.lossfn = lossfnself.optimizer = optimizerself.verbose = Falseself.writer = Noneself.train_loader = Noneself.test_loader = Nonesetattr(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 npclass 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 = lossfnself.optimizer = optimizerself.verbose = Falseself.writer = Noneself.train_loader = Noneself.test_loader = Nonedef set_loader(self, train_loader, test_loader=None):self.train_loader = train_loaderself.test_loader = test_loaderprint(f'set train_loader: {self.train_loader}\ntest_loader: {self.test_loader}') if self.verbose else Nonedef 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_stepdef 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_stepdef mini_batch(self, test=False):data_loader = Nonestep_fn = Noneif test:data_loader = self.test_loaderstep_fn = self.build_test_step()else:data_loader = self.train_loaderstep_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 lossdef 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 = Truetorch.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 Noneself.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 Nonedef 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 Nonedef 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 Nonedef reset(self, model, lossfn, optimizer):if hasattr(self, "model"):self.model.cpu() if self.model != None else Nonedel self.modeldel self.optimizerself.model = modelself.lossfn = lossfnself.optimizer = optimizerprint(f"reset model and optimizer: {self.model.state_dict()}, {self.optimizer.state_dict()}") if self.verbose else None参考资料
- 基于pytorch+可视化重学线性回归模型
- 基于numpy演练可视化梯度下降
