Runx教程
简介
Runx是NVIDIA开源的一款深度学习实验管理工具,开源于2020年三月,是个很新的开源库。其功能是帮助深度学习研究者自动执行一些常见任务,如超参数扫描、输出日志记录、运行唯一目录创建、实验总结等操作,这是考虑到深度学习一个实验就要几个小时、几天甚至几个月,如果没有及时进行保存和记录会是一件很低效的事情。
安装
pip快速安装
pip install runx
源码安装
git clone https://github.com/NVIDIA/runx
cd runx
python setup.py install --user
模块介绍
Runx主要有三个模块,分别是runx、logx和sumx。这些模块设计之初,期待被一起使用,但如果想只使用runx也是可以的,注意的是,sumx必须配合logx进行使用。
runx
使用简单的YAML文件配置实验参数,runx会扫描这个文件并逐个执行。
logx
保存metrics、messages、checkpint以及tensorboard记录文件。
sumx
对训练结果进行总结。
runx
runx主要是实验运行的配置扫描,并逐次执行,例如有个train的Python脚本,需要学习率和学习器两个参数,只需要配置如下一个名为sweep.yml
的YAML文件即可。
cmd: 'python train.py'
hparams:
lr: [0.01, 0.02]
solver: ['sgd', 'adam']
在命令行通过python -m runx.runx sweep.yml
,之后会逐次自动执行下面的命令。它的执行参数是互相组合得到的。该模块对作业系统支持较好,当使用集群时,可以方便的进行任务提交。
python train.py --lr 0.01 --solver sgd
python train.py --lr 0.01 --solver adam
python train.py --lr 0.02 --solver sgd
python train.py --lr 0.02 --solver adam
logx
该模块是我最常使用的模块,也是后续sumx模块必要的步骤,将metric传给logx才能进行用sumx进行实验分析。同时,logx还支持标准输出的本地化、方便地tensorboard记录已经快速的模型保存。
使用该模块首先需要进行导入,得到logx对象后需要先进行初始化,初始化包含下面的参数。
logx.initialize(
logdir=None, # 日志存储目录,不存在会自动创建
coolname=False, # 是否在logdir下生成一个独有的目录
hparams=None, # 配合runx使用,超参数都存下来
tensorboard=False, # 是否写入tensorboard文件,无需手动writer
no_timestamp=False, # 是否不启用时间戳命名
global_rank=0, # 分布式训练防止多输出
eager_flush=True # 打开后即使tensorboard写入过快也不会丢失
)
初始化之后就可以正常使用logx的几个主要功能了。
logx.msg("content")
将所有的print语句换成上面的函数即可,这样所有的标准输出都会保存在logdir
目录下的一个logging.log
文件中,方便后续查看。
logx.metric(phase='val', metrics=metrics, epoch=epoch)
对传入的字典metrics
中的数据进行保存记录,若开启了tensorboard则自动add_scalar
。phase
表示不同的阶段,可选train
和val
,epoch
表示全局轮次。
logx.save_model(
save_dict, # 保存的checkpoint字典, 里面可以有`epoch`,`state_dict`等
metric, # 评估指标,logx会维护一系列的该指标值,若变好则保存新模型
epoch, # 当前轮次
higher_better=True, # 是否更高更好,如准确率
delete_old=True # 是否删除旧模型
)
sumx
用于归总多次试验运行的结果,依赖于logx产生的metric文件,避免读取tensorboard文件增加分析速度。
如执行python -m runx.sumx sweep
可得如下分析结果。
lr solver acc epoch epoch_time
run4 0.02 adam 99.1 10 5:21
run3 0.02 sgd 99.0 10 5:02
run1 0.01 sgd 98.2 10 5:40
run2 0.01 adam 98.1 10 5:25
demo
下面的代码演示训练模型拟合一元二次函数的例子。
import torch
import numpy as np
import torch.nn as nn
import torch.optim as optim
from runx.logx import logx
# dataset
x = torch.arange(16*10).view(16, 10).float()
x = x + torch.rand_like(x) / 10
y = x ** 2 + 1
# model
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.fc1 = nn.Linear(1, 10)
self.fc2 = nn.Linear(10, 1)
def forward(self, x):
x = self.fc1(x)
x = nn.functional.relu(x)
x = self.fc2(x)
return x
# log
logx.initialize("./logs/", coolname=True, tensorboard=True)
model = Model()
model.cuda()
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)
best_loss = np.inf
for epocn in range(10):
losses = 0.0 # epochs training loss
losses_valid = 0.0 # epoch validation loss
for step, (x_, y_) in enumerate(zip(x, y)):
# data
x_, y_ = x_.cuda().view(10, 1), y_.cuda().view(10, 1)
x_train, y_train = x_[:5], y_[:5]
x_val, y_val = x_[5:], y_[5:]
# forward
out = model(x_train)
# loss
loss_ = criterion(out, y_train)
losses += loss_
# backward
optimizer.zero_grad()
loss_.backward()
optimizer.step()
# valid
with torch.no_grad():
pred = model(x_val)
valid_loss = criterion(pred, y_val)
losses_valid += valid_loss
metrics = {
'loss': losses / (step + 1),
}
metrics_valid = {
'loss': losses_valid / (step+1)
}
logx.metric('train', metrics, epocn)
logx.metric('val', metrics_valid, epocn)
if losses_valid < best_loss:
best_loss = losses_valid
save_dict = {
'state_dict': model.state_dict()
}
logx.save_model(save_dict, best_loss, epocn, higher_better=False, delete_old=True)
logx.msg("epoch {}, train loss {:.6f}, valid loss {:.6f}".format(epocn, losses / (step+1), losses_valid / (step+1)))
这段训练代码执行后会生成如下结果,包含模型checkpoint,tensorboard文件,标准输出日志(logging.log)以及metric记录。
打开TensorBoard服务,可视化训练过程,结果如下,可以看到,logx已经自动记录了tensorboard文件。
补充说明
本文主要介绍了Runx的使用(主要是logx模块的使用),该包其实是对深度学习实验过程的一些需求进行了封装而避免我们写很多繁琐的代码,比如TensorBoard使用writer记录,比如loss变化判断是否模型更好等等功能,非常方便,具体详细的API文档可以查看官方文档。