引言
深度学习领域我们常用jupyter来演练代码,但实际生产环境中不可能像jupyter一样,所有代码逻辑都在面向过程编程,这会导致代码可复用性差,维护难度高。
前面这篇文章 基于pytorch+可视化重学线性回归模型 已经封装了数据加载器,本文我们将要对整个训练循环的逻辑进行重构,采用封装的方式来提升代码的可复用性,降低维护难度。
步骤大概是:
- 封装小批量单次训练
- 封装小批量单次测试
- 封装训练循环
- 封装损失数据的收集和可视化
- 封装参数和梯度变化的数据可视化
- 封装保存和加载模型
首先,导入需要的包
python
import torch
import numpy as np
import matplotlib.pyplot as plt
import torch.nn as nn
import torch.optim as optim
1. 数据准备
1.1 数据生成
鉴于正式工程中不会自己生成数据,所以数据生成部分始终会保持不变。
python
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显存的使用。
python
x_tensor = torch.as_tensor(x).float()
y_tensor = torch.as_tensor(y).float()
对于单纯的tensor数据可以直接使用pytorch内置的TensorDataset类来封装数据集, 并使用random_split来划分训练集和测试集。
python
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来代替自定义。
python
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 创建训练器
为了实现训练器的高内聚、低耦合,分离动态与静态,我们使用面向对象的方法对其进行重构,对于模型训练这个业务来说:
- 变化的内容应该是模型、数据源、损失函数、优化器、随机数种子等,这些内容应该由外部传入;
- 不变的内容应该是训练循环、小批量迭代训练、评估模型损失、模型的保存和加载、预测计算等,这些内容应该由内部封装。
首先,我们定义一个训练器类,它包含以下功能:
python
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定义的数据加载器,所以需要给训练器添加一个设置数据加载器的方法。
python
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()
,但这样一来是无法记录和观测梯度值的,给排查问题造成阻碍,所以这里将梯度清零的步骤移到下一次训练之前,目的是允许主循环获取当前梯度值。
python
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 添加小批量迭代方法
在小批量迭代训练过程中,是训练和测试两个环节交叉进行。这两个环节的逻辑很相似,都是输入数据输出损失,不同之处在于所使用的数据加载器和单次迭代函数不同。我们可以封装一个统一的小批量迭代方法,来屏蔽这个差别。
python
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配合使用。
python
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决定。
python
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显示出来,以观察训练和测试两条损失曲线是否随着迭代次数而稳定下降。
python
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: 收集音频数据,可以用来观察训练过程。
python
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只适用于一个图标一组数据的情况。
python
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中内嵌一个网页来访问此服务。
python
%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内置到训练器中,这样我们就可以在训练过程中收集数据了。
python
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: 专门用于收集参数的方法,包括参数值本身和梯度。
python
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()
这个函数调用来完成,主循环反而变得更简单清晰:
注:在训练之前,先收集原始参数值,是为了保证原始参数值也被收集,并在图表中显示出来。
python
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函数来重置模型,主要功能是将模型和优化器重置,并删除旧的训练日志。
python
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)
python
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方法重新开始训练。
python
trainer.train(100)
trainer.model.state_dict()
OrderedDict([('weight', tensor([[1.8748]])), ('bias', tensor([1.0477]))])
python
# %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()函数保存到文件中。
python
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()将之前保存在文件中的模型和参数加载进来。
python
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)
python
device
train_loader
test_loader
train_losses
test_losses
writer
debug
optimizer
为了与前面的训练结果完全隔离开,我们重新创建一个训练器,一个新训练器需要进行的初始化总共包括以下几项:
- 模型、损失函数和优化器
- 训练数据集和测试数据集的加载器
- 随机数种子
- 设置训练数据的收集位置,便于tensorboard可视化
python
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加载模型参数,可以看到之前的训练结果已经加载进新的训练器。
python
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]))])
接着之前的训练结果继续训练
python
trainer2.train(100)
python
%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. 训练器封装结果
到目前为止,给训练器添加的所有方法汇总如下:
python
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()
给训练器添加的所有字段汇总如下:
python
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__
方法进行改造,以便对这些字段提前初始化。
python
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类的完整代码:
python
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