如果觉得下面哪部分讲的不够清楚,欢迎评论指出。
2022.5.30 更新:支持CI自动化测试,代码质量更加可靠
2022.5.29 更新:新增详细的API文档
为什么需要Trainer?
在江湖上混,没有一把趁手的兵器怎么能行。对于一个深度学习算法工程师来说,一个强大的Trainer就是趁手的兵器,在复现顶会论文的时候不用每次都要写大量的模版代码,可以更加关注模型实现本身。我平时也开源了一些深度学习项目,累计获赞400+,最近整理了一些之前的代码,也看了很多优秀的开源实现,然后从0实现了一个小而美的trainer,主要有以下两个优点吧:
- 轻量级:代码量很少,易上手也易改动。
- 可拓展:能够满足工程师各种魔改网络/训练流程的需求。
为什么不选择pytorch-lightning / detectron2 / mmcv?
现有的detectron2、mmcv、pytorch-lightning中的trainer虽然也很优雅,但是看了源码就能感受到代码中的抽象层次太多,看的有些吃力。例如下面这个文件是detectron2中保存checkpoint的代码,共594行,考虑到了很多场景,功能非常完善。
fvcore/checkpoint.py at main · facebookresearch/fvcore
但是对于我们平时使用来说,可能需要的只是如下几行代码。
此外detectron2是基于迭代而不是基于epoch来训练的,和我平时的习惯有些许出入。mmcv虽然是基于epoch的,但是整体设计上要略微逊色于detectron2。
最重要的一点:虽然以上这些开源库提供的trainer足够强大,学会之后能满足所有需求,但是,我还是想造个轮子。一方面是实现过程中会深入思考一些工程细节,对于代码水平会有提高。另一方面是写一个trainer也不咋耗时间,最近没什么事情,一周还是拿得出来的。
代码已经放到github上了,仓库名字叫:Core-PyTorch-Utils,简称CPU。代码注释和文档都非常详细。
GitHub - serend1p1ty/core-pytorch-utils: Core components for deep learning.
里面大多数文件都不会依赖其它模块,可以直接copy作为小工具使用。下面将逐个介绍这里的每个模块。
[1/6]灵活的配置系统:config_argparse.py
一般在一个深度学习项目中,有数十上百个需要调节的参数,cpu/config_argparse.py这个文件提供了一个简易灵活的配置系统,实现细节参考自pytorch-image-models。
目前有两种比较流行的参数配置系统:
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--batch-size', type=int, default=1)
parser.add_argument('--learning-rate', type=float, default=0.001)
print(parser.parse_args())
- 利用
yacs
库创建默认配置,然后从命令行传递--config-file
参数来覆盖默认配置。
default_config.py
保存了项目的默认配置。
from yacs.config import CfgNode as CN
_C = CN()
_C.BATCH_SIZE = 1
_C.LEARNING_RATE = 0.001
def get_default_cfg():
"""返回一个默认配置的备份,防止默认配置被修改。"""
return _C.clone()
config.yaml
为某次实验使用的配置文件,用来覆盖默认配置。
BATCH_SIZE: 5
main.py
中创建一个包含 --config-file
的 ArgumentParser
来实现加载配置文件。
import argparse
from default_config import get_default_cfg
parser = argparse.ArgumentParser()
parser.add_argument("--config-file", help="配置文件的路径。")
parser.add_argument("opts", nargs=argparse.REMAINDER, help="通过命令行修改配置。")
args = parser.parse_args()
# 最终的配置 = 默认配置 + 配置文件
cfg = get_default_cfg()
cfg.merge_from_file(args.cfg_file)
cfg.merge_from_list(args.opts)
cfg.freeze()
print(cfg)
然后运行 python main.py --config-file config.yaml
,最终得到的配置为
BATCH_SIZE = 5
LEARNING_RATE = 0.001
如果配置文件中遗漏了一些内容,我们也可以在命令行中实时地修改配置。运行 python main.py --config-file config.yaml LEARNING_RATE 0.002
,最终得到的配置为
BATCH_SIZE = 5
LEARNING_RATE = 0.002
第一种方式的缺点是无法加载 yaml
配置文件,所有参数都要在命令行中输入(虽然 argparse
自带 fromfile_prefix_chars
特性,但不好用)。第二种方法的缺点是,一个配置系统被割裂成了两个部分, yacs
部分+ argparse
部分,且在命令行中通过 python main.py --help
无法查看每个参数的说明(只能显示由 argparse
管理的那部分参数,由 yacs
管理的那部分参数则无能为力)。
因此我们的 config_argparse.py
就实现了一个可以加载配置的文件的 ArgumentParser
类。使用方式和 argparse
完全相同,只是多了一个 --config
的选项用来指定配置文件的路径。下面讲讲一些实现细节。
我们先定义了一个内部 parser
用来解析用户指定的配置文件的路径。
self.config_parser.add_argument("-c", "--config", default=None, metavar="FILE",
help="where to load YAML configuration")
然后根据用户指定的配置文件覆盖相应选项的默认值。
# 先解析出用户指定的配置文件的路径,保存在res.config中
res, remaining_argv = self.config_parser.parse_known_args(args)
if res.config is not None:
# 加载配置文件
with open(res.config, "r") as f:
config_vars = yaml.safe_load(f)
# 判断配置文件中是否含有错误的键值对
for key in config_vars:
if key not in self.option_names:
self.error(f"unexpected configuration entry: {key}")
# ----将配置文件提供的值用作默认值----
# 例如在定义某个选项的时候默认值为1,parser.add_argument("--batch-size", default=1)
# 配置文件的内容为"batch_size: 2"
# 那么就用配置文件提供的值覆盖该选项的默认值,即目前batch_size的默认值为2
self.set_defaults(**config_vars)
# parse其余命令行参数
# 如果在命令中没有指定batch_size,那么最终的batch_size就是覆盖之后的默认值2
# 如果在命令行中指定了--batch-size=3,那么最终的batch_size就是3,因为命令行具有最高优先级!
return super().parse_args(remaining_argv)
目前我们实现的 ConfigArgumentParser
类,还有一些小问题,例如:
如果指定了一个必选参数, parser.add_argument("--batch-size", required=True)
,且在配置文件中也已经设置了 batch_size
的值,这个时候在命令行中理应不需要再设置 batch_size
了。但是目前我们的实现仍然需要在命令行中指定 batch_size
的值。这是因为我们将配置文件视为选项的默认值,而不是用户的实际输入值。
如果非常介意这个问题,可以选用ConfigArgParse这个库(实际上我们的实现方式就是ConfigArgParse的早期实现的简单版本)。但是在常用的深度学习系统中,我们提供的这个类已经足够了。
[2/6]将信息显示在终端和文件中:logger.py
在深度学习项目中经常需要把终端显示的内容保存到文件中,这里我们选择使用python自带的 logging
模块来作为信息记录系统。
我们在 logger.py
中实现了一个 setup_logger()
函数,用来对 logging
模块返回的 logger
做了一些简单配置
- 设置终端中显示的信息的格式
- 终端信息高亮,让warning和error信息更加醒目
- 将终端显示的信息同步保存到文件中
下面是核心代码的讲解。
# 保证一个logger只会被初始化一次
if name in logger_initialized:
return logger_initialized[name]
# 如果name等于None,则返回的是全局的root logger
# 初始化完root logger之后,所有logger的信息(除非这个logger的propagate为False)都会被传播到root logger,
# 然后root logger再根据配置好的格式将消息显示在终端中。
# 所以最常见的用法就是在程序的最开始调用一次setup_logger()配置好root logger,
# 然后其它所有logger都免配置直接使用。
logger = logging.getLogger(name)
logger.setLevel(log_level)
# 如果同时创建了子logger和父logger,
# 那么子logger产生的信息会向上传递给父logger,
# 终端中就会把同一段信息显示两次!因此在初始化logger的时候要设置logger.propagate=False,
# 防止子logger的信息向上传递。
logger.propagate = False
plain_formatter = logging.Formatter(
"[%(asctime)s %(name)s %(levelname)s]: %(message)s", datefmt="%m/%d %H:%M:%S"
)
# 在分布式训练中只初始化主进程的logger
if rank == 0:
ch = logging.StreamHandler(stream=sys.stdout)
ch.setLevel(log_level)
if color:
# _ColorfulFormatter的功能是,如果当前信息的级别是INFO,
# 那么就不展示级别,如果是WARNING、DEBUG、ERROR,
# 就在信息的开头显示对应级别。这样的好处是,因为绝大多数信息都是INFO,
# 这样以来可以节省空间显示更多文字。
formatter = _ColorfulFormatter(
colored("[%(asctime)s %(name)s]: ", "green") + "%(message)s",
datefmt="%m/%d %H:%M:%S",
)
else:
formatter = plain_formatter
ch.setFormatter(formatter)
logger.addHandler(ch)
if output is not None:
if output.endswith(".txt") or output.endswith(".log"):
filename = output
else:
filename = os.path.join(output, "log.txt")
# 如果单个文件名作为参数传递,则os.path.dirname将返回一个空字符串。
# 例如,os.path.dirname("log.txt")== ""。这将导致os.makedirs()出错。
# 所以我们需要用绝对路径来获取目录名。
os.makedirs(os.path.dirname(os.path.abspath(filename)), exist_ok=True)
fh = logging.FileHandler(filename)
fh.setLevel(log_level)
fh.setFormatter(plain_formatter)
logger.addHandler(fh)
logger_initialized[name] = logger
return logger
[3/6]平滑显示,让loss不要再上下乱跳了:history_buffer.py
如果想在终端中输出训练过程的loss值,相信有些同学会像下面这么写:
for i, batch in enumerate(data_loader):
output = model(batch)
loss = criterion(output, label)
if i % 10 == 0:
print(f"loss: {loss.item()}")
但是这样一来输出的loss跳变很大不利于观察,可能你看到的是这样一组数字:0.5、0.8、0.3…
为了更加平滑的显示loss,我们不应该直接输出loss的当前值,而是应该输出过去一段时间内loss的平均值。假设在5次迭代中loss的值分别为0.5、0.8、0.3、0.4、0.3,窗口大小为2(我们每次输出的是过去两个迭代的loss的平均值)。
- 在第一次迭代时,loss队列如下:
0.5
数量太少,无法填满窗口,直接输出0.5
- 第二次迭代,loss队列如下,窗口为0.5和0.8。
--------
0.5 0.8
--------
因此输出窗口内的平均值0.65
- 第二次迭代,loss队列如下,窗口右移(因为窗口的含义是最近两个迭代)
--------
0.5 0.8 0.3
--------
输出窗口内的平均值0.55
- 剩下的迭代以此类推。
history_buffer中提供了一个HistoryBuffer类,可以帮助我们输出窗口内的平均值(也能输出全局平均值、全局和、最近的值)。用法如下:
loss_buffer = HistoryBuffer()
for i, batch in enumerate(data_loader):
output = model(batch)
loss = criterion(output, label)
loss_buffer.update(loss.item())
if i % 10 == 0:
print(f"loss: {loss_buffer.avg}")
pytorch-image-models中也有类似的用法,不过使用的是只能输出全局平均值的 AverageMeter
。个人认为局部平均值比全局平均值更好,能够更明显地反映出当前的变化。所以不太理解为什么pytorch-image-models为什么这样做,知道的朋友烦请评论区告知。
[4/6]让PyTorch自带的lr scheduler支持warmup:lr_scheduler.py
这部分内容太长了,单独放到一篇文章中。
[5/6]巧用钩子赋予Trainer更强的可拓展性
CPU库的代码是写死的,而用户会有各种各样的需求,比如想每隔3个epoch就打印一次当前的学习率。为了给Trainer带来更强的拓展性,CPU采用了钩子(hook)机制。
我们首先定义了钩子的基类Hook,包含了几个内置的方法。以 before_epoch()
为例,这个方法将在每个epoch开始之前调用。
class Hook:
def before_train(self) -> None:
"""Called before the first iteration."""
pass
def after_train(self) -> None:
"""Called after the last iteration."""
pass
def before_epoch(self) -> None:
"""Called before each epoch."""
pass
def after_epoch(self) -> None:
"""Called after each epoch."""
pass
def before_iter(self) -> None:
"""Called before each iteration."""
pass
如果想要实现自定义的功能(例如每隔三个epoch保存一次模型的权重),用户可以继承这个基类。
class CheckpointHook(Hook):
def after_epoch(self):
# self.trainer是在钩子注册的时候定义的
epoch = self.trainer.epoch
if epoch % 3 == 0:
save_checkpoint()
再写一个简单的Trainer类。
class Trainer:
def __init__(self):
self._hooks = []
def register_hooks(self, hooks):
# 将钩子注册到trainer中,实际上就是放到trainer的_hooks列表里以便后续调用
for hook in hooks:
# 这里为每个钩子创建一个类变量,指向当前trainer。
# 这样就可以访问trainer内部的model、optimizer、epoch,iter等。
hook.trainer = self
self._hooks.append(hook)
def before_train(self):
# trainer的before_train()函数就是
# 调用每个注册的hook的before_train()函数
for hook in self._hooks:
hook.before_train()
def after_train(self):
for hook in self._hooks:
hook.after_train()
def before_epoch(self):
for hook in self._hooks:
hook.before_epoch()
def after_epoch(self):
for hook in self._hooks:
hook.after_epoch()
def before_iter(self):
for hook in self._hooks:
hook.before_iter()
def after_iter(self):
for hook in self._hooks:
hook.after_iter()
def train(self):
self.before_train()
for self.epoch in range(max_epochs):
self.before_epoch()
for self.iter, data in enumerate(data_loader):
self.before_iter()
self.train_one_iter(data)
self.after_iter()
self.after_epoch()
self.after_train()
然后实例化我们刚才写的 CheckpointHook
类,并将其注册到 Trainer
中。
ckpt_hook = CheckpointHook()
trainer = Trainer()
trainer.register_hooks([ckpt_hook])
trainer.train()
上面这些是伪代码,不能运行,只是为了方便描述钩子的原理。
[6/6]开始写 Trainer
本节我们将逐行讲解 trainer.py
的内容。先来看构造函数。
class Trainer:
def __init__(
self,
model: nn.Module,
optimizer: optim.Optimizer,
lr_scheduler: optim.lr_scheduler._LRScheduler,
data_loader: DataLoader,
max_epochs: int,
work_dir: str = "work_dir",
max_num_checkpoints: int = None,
checkpoint_period: int = 1,
log_period: int = 50,
clip_grad_norm: float = 0.0,
enable_amp: bool = False,
# 以下是关于lr warmup的参数
by_epoch: bool = True,
warmup_t: int = 0,
warmup_by_epoch: bool = False,
warmup_mode: str = "fix",
warmup_factor: float = 0.0,
):
"""
Args:
model (torch.nn.Module)
optimizer (torch.optim.Optimizer)
lr_scheduler (optim.lr_scheduler._LRScheduler)
data_loader (torch.utils.data.DataLoader): 训练数据。
max_epochs (int): 训练多少个epoch。
work_dir (str): 保存checkpoint和log文件的目录,默认是"work_dir"。
max_num_checkpoints (int): 最多保存多少个checkpoint,如果是None,则保存所有的checkpoint。
checkpoint_period (int): 每隔多少个epoch保存一次checkpoint,默认是1。
log_period (int): 每隔多少次迭代输出一次log,默认是50。
clip_grad_norm (float): 梯度裁剪使用的范数阈值,如果<=0,则不使用梯度裁剪。
enable_amp (bool): 是否开启混合精度训练,默认不开启。
by_epoch, warmup_t, warmup_by_epoch, warmup_mode, warmup_factor: 参考 lr_scheduler.py的文档。
"""
self.model = model
self.optimizer = optimizer
# 将pytorch自带的lr scheduler转化为带warmup的scheduler
self.lr_scheduler = LRWarmupScheduler(
torch_scheduler=lr_scheduler, by_epoch=by_epoch, epoch_len=len(data_loader),
warmup_t=warmup_t, warmup_by_epoch=warmup_by_epoch, warmup_mode=warmup_mode,
warmup_factor=warmup_factor)
self.data_loader = data_loader
self.work_dir = work_dir
# 我们使用HistoryBuff来平滑训练过程中产生的指标。
# MetricStorage等下会详细讲,这里可以简单地理解成是一个字典。
# 字典的key是指标的名字,value是保存这个指标的HistoryBuffer。
self.metric_storage = MetricStorage()
# 一些有用的计数器,在hook中可能会用到
self.inner_iter: int # [0, epoch_len - 1]
self.epoch: int # [0, max_epochs - 1]
self.start_epoch = 0 # [0, max_epochs - 1]
self.max_epochs = max_epochs
# 保存所有已注册的hook
self._hooks: List[HookBase] = []
# 因为想记录从data_loader中加载每个batch需要的时间,所以我们
# 没有选择用for循环,而是用迭代器的方式来遍历data_loader
self._data_iter = iter(data_loader)
self._max_num_checkpoints = max_num_checkpoints
self._checkpoint_period = checkpoint_period
self._log_period = log_period
self._clip_grad_norm = clip_grad_norm
self._enable_amp = enable_amp
# 注册一些默认的hooks
self.register_hooks(self._build_default_hooks())
logger.info(f"Registered default hooks: {self.registered_hook_names}")
if self._enable_amp:
logger.info("Automatic Mixed Precision (AMP) training is on.")
self._grad_scaler = GradScaler()
然后我们又定义了一些有用的属性,注意这里使用了python的 @property
装饰器。以 self.max_iters
为例,当更改 self.max_epochs
时, self.max_iters
的数值也会自动改变,不需要我们手动更新。
@property
def lr(self) -> float:
return self.optimizer.param_groups[0]["lr"]
@property
def epoch_len(self) -> int:
"""每个epoch包含多少个iter(iteration,迭代)"""
return len(self.data_loader)
@property
def max_iters(self) -> int:
"""一共需要训练多少个iter"""
return self.max_epochs * self.epoch_len
@property
def cur_iter(self) -> int:
"""当前iter的索引,范围是[0, max_iters - 1]"""
return self.epoch * self.epoch_len + self.inner_iter
@property
def start_iter(self) -> int:
"""训练从哪个iter开始,最小值是0"""
return self.start_epoch * self.epoch_len
@property
def ckpt_dir(self) -> str:
"""保存checkpoints的目录"""
return osp.join(self.work_dir, "checkpoints")
@property
def tb_log_dir(self) -> str:
"""保存tensorboard log文件的目录"""
return osp.join(self.work_dir, "tb_logs")
@property
def model_or_module(self) -> nn.Module:
if isinstance(self.model, (DistributedDataParallel, DataParallel)):
return self.model.module
return self.model
@property
def registered_hook_names(self) -> List[str]:
"""所有已经注册的钩子的名称"""
return [h.__class__.__name__ for h in self._hooks]
然后是一些和钩子相关的方法。
def register_hooks(self, hooks: List[Optional[HookBase]]) -> None:
hooks = [h for h in hooks if h is not None]
for h in hooks:
assert isinstance(h, HookBase)
# 这里使用weakref是为了避免循环引用导致的内存泄漏
h.trainer = weakref.proxy(self)
# LoggerHook旨在记录训练过程中产生的一些指标(模型的loss、其它钩子产生的指标)
# 为了不遗漏其它钩子产生的信息,LoggerHook应该被最后执行(注册的时候放到末尾)
if self._hooks and isinstance(self._hooks[-1], LoggerHook):
self._hooks.insert(len(self._hooks) - 1, h)
else:
self._hooks.append(h)
def _call_hooks(self, stage: str) -> None:
"""_call_hooks("before_epoch")可以调用所有钩子的before_epoch()方法"""
for h in self._hooks:
getattr(h, stage)()
def _build_default_hooks(self) -> List[HookBase]:
return [
LRUpdateHook(), # 更新学习率应该在iter结束之后立马执行,因此要放在第一位
CheckpointerHook(self._checkpoint_period, self._max_num_checkpoints),
LoggerHook(self._log_period, tb_log_dir=self.tb_log_dir),
]
然后就是核心的训练环节。
def train(self) -> None:
logger.info(f"Start training from epoch {self.start_epoch}")
self._prepare_for_training()
self._call_hooks("before_train")
for self.epoch in range(self.start_epoch, self.max_epochs):
self._call_hooks("before_epoch")
self._train_one_epoch()
self._call_hooks("after_epoch")
self._call_hooks("after_train")
def _train_one_epoch(self) -> None:
# 测试hook会把模型设置为eval模式,因此这里要设置为train模式
self.model.train()
for self.inner_iter in range(self.epoch_len):
self._call_hooks("before_iter")
self.train_one_iter()
self._call_hooks("after_iter")
# 这里要手动更新迭代器,不然迭代到data_loader末尾时会引发异常
self._data_iter = iter(self.data_loader)
def train_one_iter(self) -> None:
iter_start_time = time.perf_counter()
# 1. 从data_loader中加载batch
# 为了计算数据加载时间,我们选择通过迭代器读取数据,而不是"for data in data_loader"
start = time.perf_counter()
batch = next(self._data_iter)
data_time = time.perf_counter() - start
# 2. 计算loss
if self._enable_amp:
with autocast():
loss_dict = self.model(batch)
else:
loss_dict = self.model(batch)
if isinstance(loss_dict, torch.Tensor):
losses = loss_dict
loss_dict = {"total_loss": loss_dict}
else:
losses = sum(loss_dict.values())
# 3. 计算梯度
self.optimizer.zero_grad()
if self._enable_amp:
self._grad_scaler.scale(losses).backward()
else:
losses.backward()
if self._clip_grad_norm > 0:
if self._enable_amp:
self._grad_scaler.unscale_(self.optimizer)
clip_grad_norm_(self.model.parameters(), self._clip_grad_norm)
# 4. 更新模型权重
if self._enable_amp:
self._grad_scaler.step(self.optimizer)
self._grad_scaler.update()
else:
self.optimizer.step()
# 记录此次迭代产生的各项指标
self._log_iter_metrics(loss_dict, data_time, time.perf_counter() - iter_start_time, self.lr)
接下来详细讲讲 Trainer
中如何自动记录训练过程中产生的各项指标。首先回顾一下 HistoryBuffer
的功能,如果内容为
--------
0.5 0.8 0.3
--------
HistoryBuffer
同时提供了 avg()
方法用来获得窗口内的平均值0.55,和 latest()
方法获得最新的数值0.3。
随着训练的进行,会把每次迭代的指标保存到字典 metric_storage
中。假设已经进行了3次迭代,此时 metric_storage
的内容为
"loss":
--------
0.5 0.8 0.3
--------
"lr":
---------
0.1 0.1 0.08
---------
LoggerHook
会每隔几个iter来从 metric_storage
中取出各项指标的值,并打印至终端。此时对于需要平滑的指标"loss", LoggerHook
打印的是 metric_storage["loss"].avg
,对于不需要平滑的指标"lr",打印的是 metric_storage["lr"].latest
。
那么 LoggerHook
又该如何知道哪个指标需要平滑呢?这就需要我们在更新 metric_storage
的时候指定一个 smooth
参数,如果 smooth
为True代表该指标需要平滑,否则不需要。
metric_storage.update(loss=loss_value, smooth=True)
metric_storage.update(lr=lr_value, smooth=False)
我们没有采用简单的字典,而是设计了 MetricStorage
类,就是为了可以自动确定哪些指标需要平滑。 MetricStorage
的代码如下。
class MetricStorage(dict):
"""该类在训练过程中存储多个度量的值(其中一些可能有噪声,例如损失、批处理时间),并提
供对平滑值的访问,以便更好地记录。
该类是为自动记录指标而设计的。用户在调用update()方法时应该指定smooth参数,
以便可以确定在执行日志记录时应该平滑哪些指标。
Example::
>>> metric_storage = MetricStorage()
>>> metric_storage.update(iter=0, loss=0.2)
>>> metric_storage.update(iter=0, lr=0.01, smooth=False)
>>> metric_storage.update(iter=1, loss=0.1)
>>> metric_storage.update(iter=1, lr=0.001, smooth=False)
>>> # loss将会被平滑, 但lr不会
>>> metric_storage.values_maybe_smooth
{"loss": (1, 0.15), "lr": (1, 0.001)}
>>> # 类似于字典,可以直接用字符串索引
>>> metric_storage["loss"].avg
0.15
"""
def __init__(self, window_size: int = 20) -> None:
self._window_size = window_size
self._history: Dict[str, HistoryBuffer] = self
self._smooth: Dict[str, bool] = {}
self._latest_iter: Dict[str, int] = {}
def update(self, iter: Optional[int] = None, smooth: bool = True, **kwargs) -> None:
"""添加在特定迭代中生成的多个指标的数值。
Args:
iter (int): 这些数值是在哪个迭代生成的。如果为None,则使用内置的从0开始的计数器。
smooth (bool): 如果为True,则在调用values_maybe_smooth()时返回这些指标的平滑值。
否则,返回最新的值。在对update()的不同调用中,同一指标必须具有相同的smooth参数值.
"""
for key, value in kwargs.items():
if key in self._smooth:
assert self._smooth[key] == smooth
else:
self._smooth[key] = smooth
self._history[key] = HistoryBuffer(window_size=self._window_size)
self._latest_iter[key] = -1
if iter is not None:
assert iter > self._latest_iter[key]
self._latest_iter[key] = iter
else:
self._latest_iter[key] += 1
self._history[key].update(value)
@property
def values_maybe_smooth(self) -> Dict[str, Tuple[int, float]]:
"""返回平滑值或最新值。具体行为取决于更新指标时的smooth参数。"""
return {
key: (self._latest_iter[key], his_buf.avg if self._smooth[key] else his_buf.latest)
for key, his_buf in self._history.items()
}
在每次迭代结束时,我们都使用 _log_iter_metrics()
函数来记录本次迭代产生的一些信息。
def log(self, *args, **kwargs) -> None:
"""调用metric_storage字典的update函数"""
self.metric_storage.update(*args, **kwargs)
def _log_iter_metrics(self, loss_dict: Dict[str, torch.Tensor], data_time: float,
iter_time: float, lr: float) -> None:
"""
Args:
loss_dict (dict): 模型产生的损失,默认是一个字典,{"loss1": 0.5, "loss2": 0.6}.
data_time (float): 从data_loader中加载batch需要的时间.
iter_time (float): 完成一次iter需要的时间.
lr (float): 本次迭代使用的学习率.
"""
self.log(self.cur_iter, data_time=data_time, iter_time=iter_time)
self.log(self.cur_iter, lr=lr, smooth=False)
loss_dict = {k: v.detach().cpu().item() for k, v in loss_dict.items()}
loss_value = sum(loss_dict.values())
if not np.isfinite(loss_value):
raise FloatingPointError(
f"Loss became infinite or NaN at epoch={self.epoch}! loss_dict = {loss_dict}."
)
self.log(self.cur_iter, total_loss=loss_value)
if len(loss_dict) > 1:
self.log(self.cur_iter, **loss_dict)
总结
感谢诸位读者的时间,你们的阅读让我的创作有了意义。断断续续这么久,终于更新完了。文章只能挑一些重点来讲,不可避免的会遗漏一些细节。感兴趣的朋友可以去github中查看完整的最新代码,欢迎批评指正+点赞fork!
路漫漫其修远,希望CPU这个库能在日后的工作学习中陪着我一起披荆斩棘,炼出有用的丹药。