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 bar：tracer.epoch_bar.close()

最后

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