torchtracer:一个管理 PyTorch AI 实验项目的工具

在使用 pytorch 框架进行机器学习(尤其是深度学习)实验时,经常需要考虑如何保存以下实验数据: - 模型的 checkpoints - 每次训练的 hyper-parameters - 训练过程中的各种变化参数及其图像(loss, accuracy, learning-rate 等) 除此之外,[Keras](https://keras.io/) 之类的其他框架在 fit 时会有一个表示训练进度的进度条,而 pytorch 原生并没有。

torchtracer:一个管理 PyTorch AI 实验项目的工具

项目地址

Github: https://github.com/OIdiotLin/torchtracer/

PyPI: https://pypi.org/project/torchtracer/

开发 torchtracer 的初衷

在使用 pytorch 框架进行机器学习(尤其是深度学习)实验时,经常需要考虑如何保存以下实验数据:

  • 模型的 checkpoints
  • 每次训练的 hyper-parameters
  • 训练过程中的各种变化参数及其图像(loss, accuracy, learning-rate 等)

除此之外,Keras 之类的其他框架在 fit 时会有一个表示训练进度的进度条,而 pytorch 原生并没有。

其实上述的这些功能完全可以在 tensorboardX 中找到,而且 UI/UX 效果也非常好,但是配置 tensorboardX 太复杂了我太菜了,于是想要自己做个简单的小工具满足上述需求。

如何使用

使用方法在项目地址里有,这里再整理一下。

安装

pip install torchtracer

创建 Tracer 实例

所有操作都是基于 Tracer 类的实例,它是所有数据的控制者。比如我们需要在 checkpoints 目录下新建一个任务 lmmnb

from torchtracer import Tracer

tracer = Tracer('checkpoints').attach('lmmnb')

需要注意的是,checkpoints 根目录需要提前创建好。为了避免用户一不小心覆盖了某个已存在的任务文件夹,当任务文件夹(比如上面的 lmmnb)已存在时,再尝试创建则会报错。

如果 attach() 不指定任务名,则会以当前的 datetime 作为任务名。

tracer = Tracer('checkpoints').attach()

保存实验配置/超参

原始配置应该以 dict 形式传入 Config 构造函数,并用 tracer.store() 方法保存之。

from torchtracer.data import Config

# `net` is a defined nn.Module
args = {'epoch_n': 120,
        'batch_size': 10,
        'criterion': nn.MSELoss(),
        'optimizer': torch.optim.RMSprop(net.parameters(), lr=1e-3)}

tracer.store(Config(args))

这一步会在 ./checkpoints/lmmnb 下创建一个 config.json,内容如下:

{
  "epoch_n": 120,
  "batch_size": 10,
  "criterion": "MSELoss",
  "optimizer": {
    "lr": 0.001,
    "momentum": 0,
    "alpha": 0.99,
    "eps": 1e-08,
    "centered": false,
    "weight_decay": 0,
    "name": "RMSprop"
  }
}

输出日志

在训练时(实际上什么时候都可以),我们可以使用 tracer.log(msg, file) 方法来输出日志到文件。如果不指定 file 参数,则会输出到 ./checkpoints/lmmnb/log 文件中去,否则会输出到 ./checkpoints/lmmnb/something.log

tracer.log(msg='Epoch #{:03d}\ttrain_loss: {:.4f}\tvalid_loss: {:.4f}'.format(epoch, train_loss, valid_loss),
           file='losses')

上面这段代码会在 ./checkpoints/lmmnb/ 中创建 losses.log,其中的日志信息如下:

Epoch #001	train_loss: 18.6356	valid_loss: 21.3882
Epoch #002	train_loss: 19.1731	valid_loss: 17.8482
Epoch #003	train_loss: 19.6756	valid_loss: 19.1418
Epoch #004	train_loss: 20.0638	valid_loss: 18.3875
Epoch #005	train_loss: 18.4679	valid_loss: 19.6304
...

保存模型

Config 一样,传递给 tracer 的模型也需要构造成 torchtracer.data.Model

tracer.store(Model(model), file='somename')

如果不指定 file,则模型名缺省为 model

上面这段代码会在 ./checkpoints/lmmnb/ 中创建两个文件:

  • 模型结构描述文件 somename.txt
Sequential(
  (0): Linear(in_features=1, out_features=6, bias=True)
  (1): ReLU()
  (2): Linear(in_features=6, out_features=12, bias=True)
  (3): ReLU()
  (4): Linear(in_features=12, out_features=12, bias=True)
  (5): ReLU()
  (6): Linear(in_features=12, out_features=1, bias=True)
)
  • 模型参数文件 somename.pth

保存 matplotlib 图像

训练过程中产生的 matplotlib 图像,我们也希望合理地、有结构地保存下来。使用 tracer.store(figure, file)images 目录下保存图片。

# assume that `train_losses` and `valid_losses` are lists of losses. 
# create figure manually.
plt.plot(train_losses, label='train loss', c='b')
plt.plot(valid_losses, label='valid loss', c='r')
plt.title('Demo Learning on SQRT')
plt.legend()
# save figure. remember to call `plt.gcf()`
tracer.store(plt.gcf(), 'losses.png')

上面这段代码会在 ./checkpoints/lmmnb/images/ 中产生一个 losses.png,表达了 loss 变化曲线。

训练进度条

在训练开始前,使用 tracer.epoch_bar_init(total) 初始化进度条。

tracer.epoch_bar_init(epoch_n)

在训练过程中,使用 tracer.epoch_bar.update(n, **param) 方法来更新进度条和进度条之后的参数。

# this runs in the end of each epoch.
tracer.epoch_bar.update(train_loss=train_loss, valid_loss=train_loss)
Tracer start at /home/oidiotlin/projects/torchtracer/checkpoints
Tracer attached with task: lmmnb
Epoch: 100%|█████████| 120/120 [00:02<00:00, 41.75it/s, train_loss=0.417, valid_loss=0.417]

最后,千万别忘了关闭 progress bartracer.epoch_bar.close()

最后

如果你使用了我的 torchtracer,并有什么新的需求/意见/建议,欢迎提交 Issue。当然,也欢迎 Star。万分感谢。