Minetorch教程

Posted 2021-06-25 周先森爱吃素

关于PyTorch的实验工具之前我介绍过NVIDIA的runx，欢迎访问，不过runx是个偏向于集群使用的一个运行实验工具，我往往只使用其中日志记录的logx模块，本文则介绍一个主攻实验控制和记录的PyTorch实验工具—Minetorch，考虑到国内好像还没有介绍该工具包的文章，因而有了这篇教程，作为一个Minecraft玩家，这个工具作者给我的感觉是很有情怀的。

Minetorch教程

简介

“在Minecraft中，火把对于挖矿是非常重要的，没有火把玩家很难获得大量的钻石等矿物。在数据科学领域也是如此，一个名为PyTorch的深度学习工具库可以帮助我们获得数据中的钻石等宝藏。Minetorch则作为研究者的工具集，帮助我们以一个更加方便的方式使用PyTorch。”这是Minetorch的官方简介，说的直白一点，这是一个针对PyTorch的工具箱，包括了训练控制、训练日志记录及可视化等诸多功能。

安装

首先是该库的安装，通过下面的命令使用pip进行安装即可。

pip install minetorch

此外，还需要安装tqdm、sklearn和jupyter等包，而且需要注意，PyTorch的版本必须大于1.7.0才能运行成功。

使用教程

首先需要说明的是，Minetorch其实对Google Sheets做了非常强大的对接，用户可以通过配置非常方便地在训练时自动将实验记录上传到电子表格中，不过由于国内Google Sheets使用不多因此本文将忽略这一部分，感兴趣的可以参考官方教程进行配置。下文主要围绕目前版本Minetorch的核心两个功能展开叙述，即实验记录和训练控制。

实验记录

首先我们来看我们使用实验记录工具最关心的实验记录功能，这就不得不提到Minetorch的一个核心类Miner，这是一个高度封装的训练器（trainer），将Model和Dataloader传给它即可方便的实现训练而不需要自己手动写循环来控制整个流程。下面就是一个Miner的构建方法，它只是将我们原本的训练流程通过这个类封装了，模型和数据的构建是不需要改动的，该对象的初始化需要的大部分参数我已经在下面的代码中详细备注了，其中drawer这个参数控制绘图器为tensorboard还是matplotlib，默认为matplotlib，它会自动绘制训练和验证损失图像。

trainer = Miner(
        alchemistic_directory='./runs',  # 日志根目录
        train_dataloader=train_loader,  # 训练dataloader
        model=model,  # 模型
        loss_func=loss_fn,  # 损失函数
        optimizer=optimizer,  # 优化器
        code="exp1",  # 实验子目录
        val_dataloader=val_loader,  # 验证dataloader
        resume=True,  # 是否自动加载之前的last模型接着训练
        eval_stride=1,  # 多少轮进行一次评估
        persist_stride=1,  # 多少轮进行一次checkpoint的保存
        drawer='tensorboard',  # 绘图器设置
        hooks=None,  # hooks配置
        logging_format=None,  # 日志格式，默认即可
        max_epochs=50,  # 训练轮数
        plugins=[
            MultiClassesClassificationMetricWithLogic(),
            NoiseSampleDetector(metric=torch.nn.CrossEntropyLoss(reduction='none'))
        ],  # 插件
        accumulated_iter=1,  # 累积多少次迭代后进行参数更新
        trival=True,  # 是否同时完成训练后随机进行验证
        in_notebook=True,  # 是否启用notebook模式，由于该库默认支持notebook，请务必打开该模式
        amp=True,  # 是否开启混合精度
        statable={}  # 实现了`state_dict`和`load_state_dict`方法的对象组成的字典，例如一个学习率调度器，trainer会自动保存这个对象的参数，注意优化器和模型默认是保存的因此不需要指定 

    )

只要调用这个trainer的train方法训练就会开始进行训练，训练其实就是类似我们自己for循环取数据进行训练，需要注意的是，这里默认是将dataloader的整个返回结果送入模型中，如果模型没有定义解包方式那么需要构建trainer的时候传入一个自定义的forward函数，参数名为forward。

训练控制

上面的这些配置只是完成了一些基本的训练流程，如果想要自定义一些新的功能那么就必须自定义一些东西，为了方便使用者控制在那个阶段干些什么，minetorch将训练分为多个hook点，用户可以注册这些hook来在不同的训练阶段做些不同的操作。

hook points	功能
after_init	trainer初始化完成调用该方法
before_epoch_start	每轮开始时调用该方法
after_epoch_end	每轮结束时调用该方法
before_train_iteration_start	每次训练迭代开始之前调用该方法
after_train_iteration_end	每次训练迭代完成后调用该方法
before_val_iteration_start	每次验证迭代开始之前调用该方法
after_val_iteration_end	每次验证迭代完成之前调用该方法
before_checkpoint_persisted	checkpoint保存前调用该方法
after_checkpoint_persisted	checkpoint保存后调用该方法
before_quit	训练完成退出前调用该方法

这些hook点的设计使得用户可以在不同的阶段嵌入自己想要增加的功能，例如想在训练每轮开始的时候打印当前轮次，那么我需要定义下面这个函数，它必须传入的是一个Miner对象用于访问miner中的所有成员。

def before_epoch_start(miner: Miner, **kwargs):
    print("start training, this is epoch {}".format(miner.current_epoch))

然后创建trainer的时候传入hook这个字典即可，那么每次之前就会输出当前轮数。

trainer = Miner(..., hooks={'before_epoch_start': before_epoch_start})

hooks的设计使得训练的拓展变得上限极高，那么插件Plugin的存在则是更具体的一种实现，它的处理流程和hook函数类似，不同的是它封装了多个hook点的功能因而实现了一个更加完善的功能，例如上面代码示例中使用的MultiClassesClassificationMetricWithLogic即实现了多分类任务指标输出和drawer（通过miner获得drawer）绘制。设计新的插件需要注册自Plugin类然后实现自己不同hook点的方法。我这里就不具体讲解了，通过阅读作者实现的异常检测的plugin，应该可以理解设计的流程，我这里给出代码。

训练演示

经过上面的讲解，其实可以发现，Minetorch是一个可拓展性非常强的工具库并且仍在不断地维护升级中，下面是我随机生成数据集进行训练流程测试的代码，由于是随机生成的数据所以拟合能力我们可以暂且忽略，可以重点看看该工具库的日志输出形式。

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from minetorch.miner import Miner
from minetorch.metrics import MultiClassesClassificationMetricWithLogic
from minetorch.plugins.noise_detector import NoiseSampleDetector

# create dataset
x = torch.zeros(1600 * 10).view(1600, 10).float()
x = x + torch.rand_like(x)
y = torch.randint(low=0, high=2, size=(1600, ))


class MyDataset(Dataset):

    def __init__(self, x, y):
        assert x.shape[0] == y.shape[0]
        self.x = x
        self.y = y

    def __getitem__(self, index):
        return self.x[index], self.y[index]

    def __len__(self):
        return self.x.shape[0]


# create model
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.fc1 = nn.Linear(10, 100)
        self.fc2 = nn.Linear(100, 2)

    def forward(self, x):
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)
        return x


def before_epoch_start(miner: Miner, **kwargs):
    print("start training, this is epoch {}".format(miner.current_epoch))


def main():
    model = Model()
    train_loader = DataLoader(MyDataset(x, y), batch_size=32, shuffle=True)
    val_loader = DataLoader(MyDataset(x, y), batch_size=32, shuffle=True)
    loss_fn = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
    # create trainer
    trainer = Miner(
        alchemistic_directory='./runs',
        train_dataloader=train_loader,
        model=model,
        loss_func=loss_fn,
        optimizer=optimizer,
        code="exp1",
        val_dataloader=val_loader,
        resume=True,
        eval_stride=1,
        persist_stride=1,
        drawer='tensorboard',
        logging_format=None,
        max_epochs=50,
        plugins=[
            MultiClassesClassificationMetricWithLogic(),
            NoiseSampleDetector(metric=torch.nn.CrossEntropyLoss(reduction='none'))
        ],
        accumulated_iter=1,
        trival=True,
        in_notebook=True,
        amp=True,
        hooks={'before_epoch_start': before_epoch_start},
    )
    trainer.train()


if __name__ == '__main__':
    main()

首先，会在trainer构建时指定的日志根目录里生成以code参数为后缀的当前实验的日志目录，然后根据使用的hook和插件以及基本的功能生成如下结构的日志目录，其中models存放保存的模型，各自插件生成的结果保存在插件名的目录下，如这里的NoiseSampleDetector和MultiClassesClassificationMetricWithLogic目录（这是基本设计，不是一定要生成在这里，比如MultiClassesClassificationMetricWithLogic的绘图结果就在其他目录里面），和上级目录同名的目录里是各种指标的绘图结果，损失是基本项，其他的均由插件支持生成。此外，还有一个名为log.txt的文本文件记录输出的日志，便于我们回顾训练的各个步骤的执行情况。

└─trival_exp1
    ├─models
    ├─MultiClassesClassificationMetricWithLogic
    ├─NoiseSampleDetector
    └─trival_exp1
        ├─accuracy
        │  └─accuracy
        ├─kappa_score
        │  └─kappa_score
        └─loss
            ├─train
            └─val

由于我这里采用的tensorboard进行绘制，通过可视化可以看到如下图所示的各个指标的绘图结果。

总结

本文主要介绍了一个在竞赛和项目中比较使用的PyTorch工具库—Minetorch，它让基于PyTorch的训练变得更加方便，欢迎大家的尝试使用。