使用optuna对模型的超参数进行自动优化

使用前先安装optuna.
pip install optuna
optuna适用于多种机器学习框架包括pytorch,tensorflow等。可以optuna examples github查看所有支持的框架的教程。

使用optuna优化pytorch的模型

先有基本的pytorch经验，可以更快的理解下面的代码。这里使用了gpu,如果没有GPU，可以修改最开始的DEVICE=torch.device("cuda")为DEVICE=torch.device("cpu")

import os
import optuna
from optuna.trial import TrialState
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.utils.data
from torchvision import datasets
from torchvision import transforms


"""
Optuna example that optimizes multi-layer perceptrons using PyTorch.

In this example, we optimize the validation accuracy of fashion product recognition using
PyTorch and FashionMNIST. We optimize the neural network architecture as well as the optimizer
configuration. As it is too time consuming to use the whole FashionMNIST dataset,
we here use a small subset of it.

"""


DEVICE = torch.device("cuda")
BATCHSIZE = 128 #每次训练时数据被分为小批次的大小
CLASSES = 10 #
DIR = os.getcwd()
EPOCHS = 10 #所有的小批次完成后，即为1个epoch,此处是设置总的训练的epoch次数
N_TRAIN_EXAMPLES = BATCHSIZE * 30 #训练集的样本数量
N_VALID_EXAMPLES = BATCHSIZE * 10 #测试集的样本数量

#定义神经网络模型，在里面使用了2个自优化超参数
def define_model(trial):
    # We optimize the number of layers, hidden units and dropout ratio in each layer.
    n_layers = trial.suggest_int("n_layers", 1, 3) #定义一个自由化超参数：训练的层数，从1到3
    layers = []

    in_features = 28 * 28
    #循环构建卷积块
    for i in range(n_layers):
        out_features = trial.suggest_int("n_units_l{}".format(i), 4, 128)#定义一个自优化超参数out_features，是输出feature的维度，从4到128
        layers.append(nn.Linear(in_features, out_features))#使用了上面定义的输出特征维度超参数
        layers.append(nn.ReLU())#层添加上激活函数ReLU
        p = trial.suggest_float("dropout_l{}".format(i), 0.2, 0.5)#定义一个自优化超参数p，是丢弃率，从0.2到0.5
        layers.append(nn.Dropout(p))#添加上上面的丢弃层
        in_features = out_features#再把最后的输出层维度赋值给输出层
    layers.append(nn.Linear(in_features, CLASSES))#最后添加上线性输出层
    layers.append(nn.LogSoftmax(dim=1))#再转为0-1分布函数，此处是分类模型，所以需要这个函数，转成概率值
    return nn.Sequential(*layers)#返回的就是一个多层神经网络的所有层

#在线加载FashionMNIST数据集,pytorch的基本操作
def get_mnist():
    # Load FashionMNIST dataset.
    train_loader = torch.utils.data.DataLoader(
        datasets.FashionMNIST(DIR, train=True, download=True, transform=transforms.ToTensor()),
        batch_size=BATCHSIZE,
        shuffle=True,
    )
    valid_loader = torch.utils.data.DataLoader(
        datasets.FashionMNIST(DIR, train=False, transform=transforms.ToTensor()),
        batch_size=BATCHSIZE,
        shuffle=True,
    )
    return train_loader, valid_loader


def objective(trial):
    # Generate the model.
    model = define_model(trial).to(DEVICE)

    # Generate the optimizers.
    optimizer_name = trial.suggest_categorical("optimizer", ["Adam", "RMSprop", "SGD"])#定义自优化的超参数：模型优化器
    lr = trial.suggest_float("lr", 1e-5, 1e-1, log=True) #定义自由化的超参数，学习率，从1e-5到1
    optimizer = getattr(optim, optimizer_name)(model.parameters(), lr=lr)#创建优化器对象

    # Get the FashionMNIST dataset.使用在线获取FashionMNIST数据集，前者是训练数据集，后者是测试数据集
    train_loader, valid_loader = get_mnist()

    # Training of the model.
    for epoch in range(EPOCHS):
        model.train()#训练模型
        for batch_idx, (data, target) in enumerate(train_loader):
            # 当训练的样本总数超过我们最开始设置的训练样本数之后，就停止训练
            if batch_idx * BATCHSIZE >= N_TRAIN_EXAMPLES:
                break

            #转换为特定硬件设备的张量，view是用于改变张量的形状，data.size(0)表示第1维度的大小，-1，表示自动计算其他维度的大小
            data, target = data.view(data.size(0), -1).to(DEVICE), target.to(DEVICE)

            optimizer.zero_grad()#清空优化器的梯度
            output = model(data)#得到模型的输出结果
            loss = F.nll_loss(output, target)#使用负对数似然损失（nll_loss）作为损失函数
            loss.backward()#反向传播
            optimizer.step()#根据计算得到的梯度，通过优化器更新模型的参数

        # 测试集，评估模型
        model.eval()
        correct = 0
        with torch.no_grad():
            for batch_idx, (data, target) in enumerate(valid_loader):
                # 如果测试数据量大于最开始设置的测试数据量就停止测试
                if batch_idx * BATCHSIZE >= N_VALID_EXAMPLES:
                    break
                data, target = data.view(data.size(0), -1).to(DEVICE), target.to(DEVICE)
                output = model(data)
                # 获取最大值的索引，后面的keepdim是在原来的张量相同的维度上保存。
                pred = output.argmax(dim=1, keepdim=True)
                #累加计算结果正确的所有的次数，
                correct += pred.eq(target.view_as(pred)).sum().item()
        #这里是计算准确率，
        accuracy = correct / min(len(valid_loader.dataset), N_VALID_EXAMPLES)

        trial.report(accuracy, epoch)#将当前的准确率和迭代次数传递给optuna

        #optuna评估是否需要剪枝，如果需要剪枝（说明模型的参数性能不佳），则抛出TrialPruned异常。
        if trial.should_prune():
            raise optuna.exceptions.TrialPruned()
    return accuracy

#判断当前脚本是否是作为主程序运行，只有作为主程序时，才会执行下面的代码，作为模块被引入时，不会执行下面的代码
if __name__ == "__main__":
    #创建一个optuna研究对象，名字是maxmize
    study = optuna.create_study(direction="maximize")
    #对目标函数进行优化，objective是要优化的目标函数，n_trials=100表示最多进行100次实验，timeout=600表示优化过程的超时时间为600秒。
    study.optimize(objective, n_trials=100, timeout=600)
    #获取所有被提前停止的实验
    pruned_trials = study.get_trials(deepcopy=False, states=[TrialState.PRUNED])
    #获取所有成功的实验，deepcopy=False表示返回的是原始对象的引用，而不是复制，这可以节省内存
    complete_trials = study.get_trials(deepcopy=False, states=[TrialState.COMPLETE])

    print("Study statistics: ")
    print("  Number of finished trials: ", len(study.trials))
    print("  Number of pruned trials: ", len(pruned_trials))
    print("  Number of complete trials: ", len(complete_trials))

    print("Best trial:")
    #获取结果最佳的实验
    trial = study.best_trial
    #输出最佳实验时的超参数设定
    print("  Value: ", trial.value)
    print("  Params: ")
    for key, value in trial.params.items():
        print("    {}: {}".format(key, value))

程序最终输出内容如下：

Study statistics: 
  Number of finished trials:  100
  Number of pruned trials:  61
  Number of complete trials:  39
Best trial:
  Value:  0.85078125
  Params: 
    n_layers: 1
    n_units_l0: 81
    dropout_l0: 0.22233836180755426
    optimizer: Adam
    lr: 0.0037253244556814374

即最优的超参数组合的模型的准确率是0.85078125，最优超参数设置是：

• 层数是1，
• 输出维度是81，
• 损失函数是0.22233836180755426，
• 模型优化器是Adam,
• 学习率是 0.0037253244556814374