【他山之石】Ray和Pytorch Lightning 使用指北

“他山之石，可以攻玉”，站在巨人的肩膀才能看得更高，走得更远。在科研的道路上，更需借助东风才能更快前行。为此，我们特别搜集整理了一些实用的代码链接，数据集，软件，编程技巧等，开辟“他山之石”专栏，助你乘风破浪，一路奋勇向前，敬请关注。

作者：知乎—Chea Sim

地址：https://www.zhihu.com/people/chea-sim

本文翻译来源：

https://docs.ray.io/en/master/tune/tutorials/tune-pytorch-lightning.html

本文默认已对pytorch lightning较为熟悉。

01

介绍

Ray是一个超参搜索调优框架，可以在只改少量代码的情况下调优我们的模型。而Pytorch lightning是一个模型设计框架，可以大大减少我们在设计实验设计模型设计中的一些dirty work。那么两样东西加在一起，会变得更加快乐吗？

首先使用Pytorch Lightning 这个库我有一个习惯（从别人代码偷来的），将里面的训练类单独拎出来，作为train_model来训练我们真正要训练的model。一般我们写深度学习任务的代码有三个步骤：

1. 进行数据预处理，比如清洗数据，数据增强等等。对于NLP任务就是进行分词，将文本输入转化为计算机能够理解的向量格式。
2. 模型选择，选择合适的模型，使用Pytorch框架的话就是继承一个nn.module类。
3. 模型训练，使用各种花里胡哨的训练方式，比如分布式，半精度，各种优化器轮着来。而这部分往往是论文不太涉及的部分，也是我们一般称为dirty work的部分。因为训练模型，超参搜索，都是没有什么创新但是又需要很大的功夫去写这些重复的代码。

而pytorch lightning这个框架主要解决的就是上述提到的第三步。将模型训练的大部分代码整合好，你只需要设定各种参数，即会自动进行训练，不需要写各种分布式代码，各种if else花里胡哨的。

众所周知，我们调参的时候，会使用控制变量法，在固定其他变量的时候，针对某一个变量进行改变，再对结果进行分析，从而找到一个较好的超参。而如果超参很多，那么像我这种使用脚本--lr 3e-5作为超参输入就会很烦，因为要测试超参就要在bash脚本中写个循环，不够优雅并且不支持多gpu同时超参搜索。最最重要的是，像我这种比较懒得，对于结果文件保存都是使用时间戳或者随便当场想一个名字，等到明天训练结束后，我都忘记了这个模型的改进点还有超参设置是什么了。所以，我需要一个能够自动统计超参，并且能够保存文件后显示最好的超参选择是什么，并且对应的模型文件在哪里。

有了Ray之后，我们需要一个能够灵活而不是写死的模型，模型的所有要变化的超参都由一个config字典存储。在模型初始化的时候设置模型超参。

02

怎么改？

假设我们已经有了一个用pytorch lightning写好的模型代码，也就是我们已经有一个可以运行，可以训练的代码。那么我们需要加入啥，让ray可以自动帮我们调参呢？

2.1 设置参数，训练，汇报结果

前面我们讲了，可以使用config设定模型超参，那么我们需要让Ray调参的时候知道每个模型训练的结果。在pytorch lightning中我们使用callback嵌入我们需要修改的代码，比如在每个epoch结束的时候，输出结果保存checkpoint。而Ray这个框架贴心地写好了callback，只需要引用一下，加一个callback就可以记录我们每个超参的结果

from ray.tune.integration.pytorch_lightning import TuneReportCallback
callback = TuneReportCallback(
    {
        "loss": "val_loss",
        "mean_accuracy": "val_accuracy"
    },
    on="validation_end")

2.2 修改argparse和训练函数

在Ray中可能会使用其他参数，比如ip地址等，所以我们在使用argparse的时候需要使用这个args = parser.parse_known_args()将设置好的参数赋值即可，其余ray的不用管。

如果你对pytorch lightning熟悉的话，应该知道在这个框架中对训练过程使用了一个LightningModule类来替代，你只需要将trainer.fit(model)就可以完成训练，而Ray这个框架需要你将训练的过程放进一个函数，之后使用tune.run(func,...)进行训练。所以我们把训练fit提出来，弄成一个函数。

2.3 塞进run

具体请看VCR。不对是栗子。

栗子

在一般情况下，我们需要定义一个训练类，继承自LightningModule。在这个类中，我们将所有需要改变的超参都使用config这个字典来读取和存储。

class LightningMNISTClassifier(pl.LightningModule):
    def __init__(self, config, data_dir=None):
        super(LightningMNISTClassifier, self).__init__()

        self.data_dir = data_dir or os.getcwd()

        self.layer_1_size = config["layer_1_size"]
        self.layer_2_size = config["layer_2_size"]
        self.lr = config["lr"]
        self.batch_size = config["batch_size"]

        # mnist images are (1, 28, 28) (channels, width, height)
        self.layer_1 = torch.nn.Linear(28 * 28, self.layer_1_size)
        self.layer_2 = torch.nn.Linear(self.layer_1_size, self.layer_2_size)
        self.layer_3 = torch.nn.Linear(self.layer_2_size, 10)

    def forward(self, x):
        batch_size, channels, width, height = x.size()
        x = x.view(batch_size, -1)

        x = self.layer_1(x)
        x = torch.relu(x)

        x = self.layer_2(x)
        x = torch.relu(x)

        x = self.layer_3(x)
        x = torch.log_softmax(x, dim=1)

        return x

    def cross_entropy_loss(self, logits, labels):
        return F.nll_loss(logits, labels)

    def accuracy(self, logits, labels):
        _, predicted = torch.max(logits.data, 1)
        correct = (predicted == labels).sum().item()
        accuracy = correct / len(labels)
        return torch.tensor(accuracy)

    def training_step(self, train_batch, batch_idx):
        x, y = train_batch
        logits = self.forward(x)
        loss = self.cross_entropy_loss(logits, y)
        accuracy = self.accuracy(logits, y)

        self.log("ptl/train_loss", loss)
        self.log("ptl/train_accuracy", accuracy)
        return loss

    def validation_step(self, val_batch, batch_idx):
        x, y = val_batch
        logits = self.forward(x)
        loss = self.cross_entropy_loss(logits, y)
        accuracy = self.accuracy(logits, y)
        return {"val_loss": loss, "val_accuracy": accuracy}

    def validation_epoch_end(self, outputs):
        avg_loss = torch.stack([x["val_loss"] for x in outputs]).mean()
        avg_acc = torch.stack([x["val_accuracy"] for x in outputs]).mean()
        self.log("ptl/val_loss", avg_loss)
        self.log("ptl/val_accuracy", avg_acc)

    @staticmethod
    def download_data(data_dir):
        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307, ), (0.3081, ))
        ])
        return MNIST(data_dir, train=True, download=True, transform=transform)

    def prepare_data(self):
        mnist_train = self.download_data(self.data_dir)

        self.mnist_train, self.mnist_val = random_split(
            mnist_train, [55000, 5000])

    def train_dataloader(self):
        return DataLoader(self.mnist_train, batch_size=int(self.batch_size))

    def val_dataloader(self):
        return DataLoader(self.mnist_val, batch_size=int(self.batch_size))

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)
        return optimizer


def train_mnist(config):
    model = LightningMNISTClassifier(config)
    trainer = pl.Trainer(max_epochs=10, show_progress_bar=False)

    trainer.fit(model)

将训练过程封装一下。让config变成变量传入。

def train_mnist_tune(config, data_dir=None, num_epochs=10, num_gpus=0):
    model = LightningMNISTClassifier(config, data_dir)
    trainer = pl.Trainer(
        max_epochs=num_epochs,
        gpus=num_gpus,
        logger=TensorBoardLogger(
            save_dir=tune.get_trial_dir(), name="", version="."),
        progress_bar_refresh_rate=0,
        callbacks=[
            TuneReportCallback(
                {
                    "loss": "ptl/val_loss",
                    "mean_accuracy": "ptl/val_accuracy"
                },
                on="validation_end")
        ])
    trainer.fit(model)

之后就需要使用Ray这个库了，Ray整合了pytorch lightning中的一些组件，可以即插即用。pytorch lightning使用了callbacks的机制在不影响原来代码的基础上加入自己想要的操作。比如在每次验证完毕后记录一些准确率等指标。

from ray.tune.integration.pytorch_lightning import TuneReportCallback
callback = TuneReportCallback({
    "loss": "avg_val_loss",
    "mean_accuracy": "avg_val_accuracy"
}, on="validation_end")

之后我们选取一下超参数的范围。

config = {
    "layer_1_size": tune.choice([32, 64, 128]),
    "layer_2_size": tune.choice([64, 128, 256]),
    "lr": tune.loguniform(1e-4, 1e-1),
    "batch_size": tune.choice([32, 64, 128]),
}

Ray有一个好东西，就是超参数搜索中的计划表，可以在有些明显不好的超参中停止训练。

scheduler = ASHAScheduler(
    max_t=num_epochs,
    grace_period=1,
    reduction_factor=2
)

在最终界面中，我们可以选择想要关注的变量。

reporter = CLIReporter(    
    metric_columns=["loss", "mean_accuracy", "training_iteration"])

最后启动我们的tune.run

analysis = tune.run(
        tune.with_parameters(
            train_mnist_tune,
            data_dir=data_dir,
            num_epochs=num_epochs,
            num_gpus=gpus_per_trial),
        resources_per_trial={
            "cpu": 1,
            "gpu": gpus_per_trial
        },
        metric="loss",
        mode="min",
        config=config,
        num_samples=num_samples,
        scheduler=scheduler,
        progress_reporter=reporter,
        name="tune_mnist_asha")

一个可以运行的小demo已经放在colab上了，可以运行一下，更清晰的理解ray和pytorchlightning。

https://colab.research.google.com/drive/1yW32SOhsCoLsbzgbL8kn-GO8o8Uru3jp?usp=sharing

03

各种错误

3.1 argparse设置

default_worker.py: error: unrecognized arguments

可能是由于argparse没有使用args = parser.parse_known_args()或者 parse放在函数里面，这个是需要放在每个训练函数外面的。

3.2 启动个数太少

设定tune.run(num_samples=10)。

本文目的在于学术交流，并不代表本公众号赞同其观点或对其内容真实性负责，版权归原作者所有，如有侵权请告知删除。

“他山之石”历史文章

更多他山之石专栏文章，

请点击文章底部“阅读原文”查看

分享、点赞、在看，给个三连击呗！