MLOps极致细节：20. MLFlow Pytorch 的使用案例2：数据分类超参数自动优化

破浪会有时

已于 2022-05-07 15:20:14 修改

阅读量502

点赞数 1

分类专栏： mlops 文章标签： pytorch 分类深度学习

于 2022-05-07 15:18:05 首次发布

本文链接：https://blog.csdn.net/zyctimes/article/details/124631415

版权

mlops 专栏收录该内容

28 篇文章 16 订阅

订阅专栏

MLOps极致细节：20. MLFlow Pytorch 的使用案例2：数据分类超参数自动优化

本案例解释了如何在Pytorch中使用MLFlow，在 IRIS 的数据分类超参数自动优化案例。

解释mlflow.pytorch的具体使用方式；
解释pl.LightningModule与pl.LightningDataModule的具体使用方式；
解释mlflow run以及mlflow ui的具体使用方式。解释如何使用parent run和child run；
通过 IRIS FLOWER 的数据分类超参数自动优化案例解释如何在pyTorch中使用mlFlow以及AxClient，并呈现结果。

运行环境：

平台：Win10。
IDE：Visual Studio Code
需要预装：Anaconda3
MLFlow当前版本：1.25.1
代码

觉得写的可以的话点个赞，收藏，加关注哦。

文章目录

MLOps极致细节：20. MLFlow Pytorch 的使用案例2：数据分类超参数自动优化

1 关于 MLFlow

MLFlow是一个能够覆盖机器学习全流程（从数据准备到模型训练到最终部署）的新平台。它一共有四大模块（如下为官网的原文以及翻译）：

MLflow Tracking：如何通过API的形式管理实验的参数、代码、结果，并且通过UI的形式做对比。
MLflow Projects：以可重用、可复制的形式打包ML代码，以便与其他数据科学家共享或部署到生产环境（MLflow项目）。
MLflow Models：管理和部署从各种ML库到各种模型服务和推理平台（MLflow模型）的模型。
MLflow Model Registry：提供一个中央模型存储，以协同管理MLflow模型的整个生命周期，包括模型版本控制、阶段转换和注释（MLflow模型注册表）。

在这个系列的前半部分，我们对MLFlow做了详细的介绍，以及每一个模块的案例讲解，这里不再赘述。

2 关于如何在PyTorch中使用MLFlow

mlflow.pytorch模块提供了一个用于记录和加载 PyTorch 模型的 API。

需要注意的是，MLFlow 无法直接和PyTorch一起使用，我们需要先装一下pytorch_lightning，当我们调用 pytorch_lightning.Trainer() 的 fit 方法时会执行自动记录（mlflow.pytorch.autolog）。

3 关于代码的运行

有两种种运行代码的方式，这里我们也会一一列举。如果我们是初学者，建议先尝试第一种方式。

3.1 第一种运行代码方式：本地创建虚拟环境运行

首先，我们在Windows的平台下安装Anaconda3。具体的安装步骤此处略过，参见Anaconda的官方文档。

安装完后，新建虚拟环境。在VSCode，使用conda create -n your_env_name python=X.X（2.7、3.6等)命令创建python版本为X.X、名字为your_env_name的虚拟环境。

这里我们输入conda create -n mlFlowEx python=3.8.2。

安装完默认的依赖后，我们进入虚拟环境：conda activate mlFlowEx。注意，如果需要退出，则输入conda deactivate。另外，如果Terminal没有成功切换到虚拟环境，可以尝试conda init powershell，然后重启terminal。

然后，我们在虚拟环境中下载好相关依赖：pip3 install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple。

这个案例的依赖包括：

mlflow==1.25.1
torchvision>=0.9.1
torch>=1.9.0
pytorch-lightning==1.6.1
ax-platform==0.2.5.1

这里，ax-platform这个库是用来进行超参数优化。

我们将代码下载到本地：git clone https://gitee.com/yichaoyyds/mlflow-ex-pytorch.git。进入文件夹AxHyperOptimizationPTL后，我们看到几个.py文件，输入python ax_hpo_iris.py运行代码。

3.2 第二种运行代码方式：mlflow run 指令运行

这里我们不需要在本地新建一个虚拟环境（mlflow run 指令会自动新建），我们将代码下载到本地，进入文件夹AxHyperOptimizationPTL后，我们在terminal运行mlflow run .。注意，最后有一个.不要忘记，这个.的意思是当前文件夹路径。在之前的系列文章中，我们解释过，运行mlflow run之后，系统会去寻找指定文件夹下的MLproject文件。此文件包含了我们需要运行的脚本指令。

4 概念解释

在解释详细代码之前（其实代码很简单），我们有四个重要的概念（函数）需要解释。

4.1 mlflow.pytorch.autolog

mlflow.pytorch的其他函数我们都可以先不看，只要在train代码之前加上这行，mlflow就可以自动开始运行，包括保存artifacts，metrics，params，tag值。有很多值是自动生成的。在结果这个章节中，我们会详细解释。这里我们先解释mlflow.pytorch.autolog函数。

一般在代码中，我们直接使用mlflow.pytorch.autolog()即可，因为函数中的参数都有默认值，我们一般使用默认值就可以。这里我们来详细过一遍其中重要的参数。完整的函数定义如下：

mlflow.pytorch.autolog(log_every_n_epoch=1, log_every_n_step=None, log_models=True, disable=False, exclusive=False, disable_for_unsupported_versions=False, silent=False, registered_model_name=None)

log_every_n_epoch: 如果指定，则每 n 个 epoch 记录一次 metric 值。默认情况下，每个 epoch 后都会记录 metric 值；
log_models：如果为 True，则经过训练的模型将记录在 MLflow artifacts 路径下。如果为 False，则不记录经过训练的模型。注意，这里只会记录一个model，应该是性能最好的那个model。在Microsoft Azure MLOps中，它会记录每一个epoch的model。所以从性能角度上，Azure MLOps工具确实强大，但我们可以根据实际需要进行选择，毕竟Azure MLOps是付费的，而MLFlow是免费的。如果对 Azure MLOps 感兴趣，也可以翻看这个系列的其他文章。
disable：如果为 True，则禁用 PyTorch Lightning 自动日志记录集成功能。如果为 False，则启用。当PyTorch Lightning 在进行模型训练（进行初始化）的时候，Lightning 在后台使用 TensorBoard 记录器，并将日志存储到目录中（默认情况下在 Lightning_logs/ 中）。相关链接。我们可以将这个默认日志功能关闭。这里我们建议开着自动日志功能，因为如果关闭，MLFlow的artifacts，metrics，params，tag默认保存的参数就无法保存；

4.2 pl.LightningModule

对于 PyTorch Lightning，有两个函数是至关重要，一个是pl.LightningModule，一个是pl.LightningDataModule。前者的包含了训练/验证/预测/优化的所有模块，后者则是数据集读取模块。我们通过PyTorch Lightning进行模型训练的时候，通常会继承这两个类。目前我对 PyTorch Lightning 不是很了解，所以这里我作为一个初学者的角度，针对这个案例进行一些相关的解读。

关于pl.LightningModule，和我们这个案例相关的函数包括：

forward，作用和torch.nn.Module.forward()一样，这里我们不再赘述；
training_step，我们计算并返回训练损失和一些额外的metrics。
validation_step，我们计算并返回验证损失和一些额外的metrics。
test_step，我们计算并返回测试损失和一些额外的metrics。
validation_epoch_end，在验证epoch结束后，计算这个epoch的平均验证accuracy。
test_epoch_end，在测试epoch结束后，计算计算这个epoch的平均测试accuracy。
configure_optimizers，选择要在优化中使用的优化器和学习率调度器。

此网页有详细的描述，这里不再赘述。

4.3 pl.LightningDataModule

pl.LightningDataModule 标准化了训练、验证、测试集的拆分、数据准备和转换。主要优点是一致的数据拆分、数据准备和跨模型转换，一个例子如下：

class MyDataModule(LightningDataModule):
    def __init__(self):
        super().__init__()
    def prepare_data(self):
        # download, split, etc...
        # only called on 1 GPU/TPU in distributed
    def setup(self, stage):
        # make assignments here (val/train/test split)
        # called on every process in DDP
    def train_dataloader(self):
        train_split = Dataset(...)
        return DataLoader(train_split)
    def val_dataloader(self):
        val_split = Dataset(...)
        return DataLoader(val_split)
    def test_dataloader(self):
        test_split = Dataset(...)
        return DataLoader(test_split)
    def teardown(self):
        # clean up after fit or test
        # called on every process in DDP

4.4 MLproject 以及 conda.yaml 文件

如果我们要使用mlflow run指令，那么我们就需要明白MLproject以及conda.yaml文件的作用。

MLproject文件：

name: ax-hpo-mnist

conda_env: conda.yaml

entry_points:
  main:
    parameters:
      epoch: {type: int, default: 50}
      total_trials: {type: int, default: 30}
      num_workers: {type: int, default: 2}
      batch_size: {type: int, default: 4}
    command: "python ax_hpo_iris.py
            --epoch {epoch}
            --total_trials {total_trials}
            --num_workers {num_workers}
            --batch_size {batch_size}"

当我们执行mlflow run .的时候，相当于系统读取当前文件夹下MLproject中的指令。如果参数我们都取默认值，那么mlflow run .就足够。但如果我们希望在terminal中输入参数值，可以使用诸如mlflow run . -P max_epochs=3 -P total_trials=3。

conda.yaml文件

channels:
- conda-forge
- https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
- https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free
- https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/r
- https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/pro
- https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/msys2
dependencies:
- python=3.8.2
- pip
- pip:
  - mlflow
  - pytorch-lightning==1.6.1
  - ax-platform
  - torchvision>=0.9.1
  - torch>=1.9.0
  - -i https://pypi.tuna.tsinghua.edu.cn/simple

包含了这个项目的依赖项。当我们第一次运行mlflow run .的时候，系统会自动新建一个虚拟环境，安装对应的依赖（对应conda.yaml文件），最后运行对应的代码（对应MLproject文件）。

5 IRIS 分类超参数自动优化案例

5.1 sklearn.datasets Iris 数据集

我们先来看一下这个Iris数据集，其实这是一个非常简单的数据集，我们这里用它来做一个例子：

from sklearn.datasets import load_iris
iris = load_iris()
dataset = iris.data
#print(dataset)
print("dataset scale: {}".format(dataset.shape))
print("Label: {}".format(iris.target_names))
print("feature names: {}".format(iris.feature_names))

结果如下：

dataset scale: (150, 4)
Label: ['setosa' 'versicolor' 'virginica']
feature names: ['sepal length (cm)', 'sepal width (cm)', 'petal length (cm)', 'petal width (cm)']

所以说，这个数据集很小，只有150个数据，4列说明只有4个特征，3个分类，分别是：‘setosa’，‘versicolor’以及’virginica’。

5.2 IrisDataModule 数据集读取与预处理模块

这个模块的代码放在iris_data_module.py中。IrisDataModule继承了pl.LightningDataModule类（这个类的说明参见上一章节）。

class IrisDataModule(pl.LightningDataModule):
    def __init__(self):
        super().__init__()
        self.columns = None

    def _get_iris_as_tensor_dataset(self):
        iris = load_iris()
        df = iris.data
        self.columns = iris.feature_names
        target = iris["target"]
        data = torch.Tensor(df).float()
        labels = torch.Tensor(target).long()
        data_set = TensorDataset(data, labels)
        return data_set

    def setup(self, stage=None):

        # Assign train/val datasets for use in dataloaders
        if stage == "fit" or stage is None:
            iris_full = self._get_iris_as_tensor_dataset()
            self.train_set, self.val_set = random_split(iris_full, [130, 20])

        # Assign test dataset for use in dataloader(s)
        if stage == "test" or stage is None:
            self.train_set, self.test_set = random_split(self.train_set, [110, 20])

    def train_dataloader(self):
        return DataLoader(self.train_set, batch_size=4)

    def val_dataloader(self):
        return DataLoader(self.val_set, batch_size=4)

    def test_dataloader(self):
        return DataLoader(self.test_set, batch_size=4)

5.3 IrisClassification 模型训练、验证、测试模块

这个模块的代码放在iris.py中。IrisClassification继承了pl.LightningModule类（这个类的说明参见上一章节）。

class IrisClassification(pl.LightningModule):
    def __init__(self, **kwargs):
        super().__init__()

        self.train_acc = Accuracy()
        self.val_acc = Accuracy()
        self.test_acc = Accuracy()
        self.args = kwargs

        self.fc1 = nn.Linear(4, 10)
        self.fc2 = nn.Linear(10, 10)
        self.fc3 = nn.Linear(10, 3)

        self.lr = kwargs.get("lr", 0.001)
        self.momentum = kwargs.get("momentum", 0.7)
        self.weight_decay = kwargs.get("weight_decay", 0.01)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        # probability distribution over labels
        x = torch.log_softmax(x, dim=1)
        return x

    def configure_optimizers(self):
        return torch.optim.SGD(
            self.parameters(), lr=self.lr, momentum=self.momentum, weight_decay=self.weight_decay
        )

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self.forward(x)
        loss = self.cross_entropy_loss(logits, y)
        self.train_acc(torch.argmax(logits, dim=1), y)
        self.log("train_acc", self.train_acc.compute(), on_step=False, on_epoch=True)
        self.log("loss", loss)
        return {"loss": loss}

    def cross_entropy_loss(self, logits, labels):
        """
        Initializes the loss function

        :return: output - Initialized cross entropy loss function
        """
        return F.nll_loss(logits, labels)

    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self.forward(x)
        loss = F.cross_entropy(logits, y)
        self.val_acc(torch.argmax(logits, dim=1), y)
        self.log("val_acc", self.val_acc.compute())
        self.log("val_loss", loss, sync_dist=True)

    def test_step(self, batch, batch_idx):
        x, y = batch
        logits = self.forward(x)
        loss = F.cross_entropy(logits, y)
        self.test_acc(torch.argmax(logits, dim=1), y)
        self.log("test_loss", loss)
        self.log("test_acc", self.test_acc.compute())

其实逻辑还是非常直白的，在初始化函数__init__中，我们设置了三层神经网络，分别是：

self.fc1 = nn.Linear(4, 10)
self.fc2 = nn.Linear(10, 10)
self.fc3 = nn.Linear(10, 3)

我们可以看到，第一层神经网络的输入是4，因为这个IRIS数据集只有4个特征；最后一层输出是3，因为IRIS的label一共只有三个。

5.4 超参数优化

在主程序ax_hpo_iris.py中，我们首先读取输入的参数：

    parser = argparse.ArgumentParser()
    parser = pl.Trainer.add_argparse_args(parent_parser=parser)

    parser.add_argument(
        "--total_trials",
        type=int,
        default=10,
        help="umber of trials to be run for the optimization experiment",
    )

    parser.add_argument(
        "--epoch",
        type=int,
        default=50,
        help="umber of trials to be run for the optimization experiment",
    )

    parser.add_argument(
            "--num_workers",
            type=int,
            default=2,
            metavar="N",
            help="number of workers (default: 3)",
    )

    parser.add_argument(
            "--batch_size",
            type=int,
            default=5,
            metavar="N",
            help="input batch size for training (default: 64)",
    )

    args = parser.parse_args()

    if "max_epochs" in args:
        args.max_epochs = args.epoch
        max_epochs = args.max_epochs
    else:
        max_epochs = args.epoch

    num_workers = args.num_workers
    batch_size = args.batch_size

    params = {"lr": 0.001, "momentum": 0.7, "weight_decay": 0.0001}

然后，我们在函数model_training_hyperparameter_tuning中新建run，使用ax.service.ax_client通过不同超参数的选择，不断训练模型，最终得到性能最好的模型。先附上代码：

def train_evaluate(params, max_epochs=100, num_workers=2, batch_size=4):
    model = IrisClassification(**params)
    dm = IrisDataModule(num_workers, batch_size)
    dm.setup(stage="fit")
    trainer = pl.Trainer(max_epochs=max_epochs)
    mlflow.pytorch.autolog()
    trainer.fit(model, dm)
    trainer.test(datamodule=dm)
    test_accuracy = trainer.callback_metrics.get("test_acc")
    return test_accuracy


def model_training_hyperparameter_tuning(max_epochs, total_trials, params, num_workers, batch_size):
    """
     This function takes input params max_epochs, total_trials, params
     and creates a nested run in Mlflow. The parameters, metrics, model and summary are dumped into their
     respective mlflow-run ids. The best parameters are dumped along with the baseline model.

    :param max_epochs: Max epochs used for training the model. Type:int
    :param total_trials: Number of ax-client experimental trials. Type:int
    :param params: Model parameters. Type:dict
    """
    with mlflow.start_run(run_name="Parent Run"):
        train_evaluate(params=params, max_epochs=max_epochs)

        ax_client = AxClient()
        ax_client.create_experiment(
            parameters=[
                {"name": "lr", "type": "range", "bounds": [1e-4, 1e-2], "log_scale": True},
                {"name": "weight_decay", "type": "range", "bounds": [1e-4, 1e-3]},
                {"name": "momentum", "type": "range", "bounds": [0.7, 1.0]},
            ],
            objective_name="test_accuracy",
        )

        for i in range(total_trials):
            with mlflow.start_run(nested=True, run_name="Trial " + str(i)) as child_run:
                parameters, trial_index = ax_client.get_next_trial()
                test_accuracy = train_evaluate(parameters, max_epochs, num_workers, batch_size)

                # completion of trial
                ax_client.complete_trial(trial_index=trial_index, raw_data=test_accuracy.item())

        best_parameters, metrics = ax_client.get_best_parameters()
        for param_name, value in best_parameters.items():
            mlflow.log_param("optimum_" + param_name, value)

代码逻辑：

我们首先新建一个parent run，所有训练的运行都在这个run下面：mlflow.start_run(run_name="Parent Run"):；

初始化AxClient，并且设定可以调节的几个超参数的数值范围：

  ax_client = AxClient()
      ax_client.create_experiment(
          parameters=[
              {"name": "lr", "type": "range", "bounds": [1e-3, 0.15], "log_scale": True},
              {"name": "weight_decay", "type": "range", "bounds": [1e-4, 1e-3]},
              {"name": "momentum", "type": "range", "bounds": [0.7, 1.0]},
          ],
          objective_name="test_accuracy",
      )

我们开始训练，每一个循环，我们会选择一套超参数，循环的次数由total_trials决定。对于每一个循环，我们会新建一个child run：with mlflow.start_run(nested=True, run_name="Trial " + str(i)) as child_run:；

每一个循环的训练/验证代码在train_evaluate函数中。

  def train_evaluate(params, max_epochs=100, num_workers=2, batch_size=4):
  model = IrisClassification(**params)
  dm = IrisDataModule(num_workers, batch_size)
  dm.setup(stage="fit")
  trainer = pl.Trainer(max_epochs=max_epochs)
  mlflow.pytorch.autolog()
  trainer.fit(model, dm)
  trainer.test(datamodule=dm)
  test_accuracy = trainer.callback_metrics.get("test_acc")
  return test_accuracy

在循环结束后，我们找到最好的run，并把最优的超参数保存：

  best_parameters, metrics = ax_client.get_best_parameters()
  for param_name, value in best_parameters.items():
      mlflow.log_param("optimum_" + param_name, value)

6 结果

我们可以（在当前路径下）在虚拟环境中通过python ax_hpo_iris.py（如加参数，则可以这么写：python ax_hpo_iris.py --epochs 3 --total_trials 3 --num_workers 2 --batch_size 4），或者mlflow run .（如加参数，则可以这么写：mlflow run . -P epoch=3 -P total_trials=3 -P num_workers=2 -P batch_size=4）来自动创建虚拟环境并运行代码。

代码运行结束后，mlflow ui打开前端网页。

在这里插入图片描述

6.1 Parent Run

我们打开parent run，由于在代码中，我们在parent run里面也训练了一次，所以我们会Parameters以及Metrics里面有很多参数。这里我们只关注optimum_lr，optimum_momentum以及optimum_weight_decay。这三个参数是所有循环结束后，系统找到的精度最高的那个模型对应的超参数的值。见下图：

在这里插入图片描述