optuna_v2_release_blog_hpi.py

import os

import optuna
from optuna import visualization
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

DEVICE = torch.device('cpu')
CLASSES = 10
# DIR = os.getcwd()
DIR = '../optuna'
EPOCHS = 20


def define_model(trial):
    layers = []

    in_features = 28 * 28
    for i in range(2):
        out_features = trial.suggest_int('n_units_l{}'.format(i), 4, 128)
        activation = trial.suggest_categorical(
            'activation_l{}'.format(i), choices=['relu', 'tanh'])
        p = trial.suggest_uniform('dropout_l{}'.format(i), 0.2, 0.5)

        layers.append(nn.Linear(in_features, out_features))
        if activation == 'relu':
            layers.append(nn.ReLU())
        elif activation == 'tanh':
            layers.append(nn.Tanh())
        else:
            assert False
        layers.append(nn.Dropout(p))

        in_features = out_features
    layers.append(nn.Linear(in_features, CLASSES))
    layers.append(nn.LogSoftmax(dim=1))

    return nn.Sequential(*layers)


def get_mnist(trial):
    batch_size = trial.suggest_categorical(
        'batch_size', choices=[8, 16, 32, 64, 128])

    train_loader = torch.utils.data.DataLoader(
        datasets.MNIST(
            DIR, train=True, download=True, transform=transforms.ToTensor()),
        batch_size=batch_size,
        shuffle=True,
    )
    test_loader = torch.utils.data.DataLoader(
        datasets.MNIST(
            DIR, train=False, transform=transforms.ToTensor()),
        batch_size=batch_size,
        shuffle=True,
    )
    return train_loader, test_loader


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

    optimizer_name = trial.suggest_categorical(
        'optimizer', ['Adam', 'RMSprop', 'SGD'])
    lr = trial.suggest_loguniform('lr', 1e-5, 1e-1)
    optimizer = getattr(optim, optimizer_name)(model.parameters(), lr=lr)

    train_loader, test_loader = get_mnist(trial)

    model.train()
    for epoch in range(EPOCHS):
        epoch_loss = 0
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.view(-1, 28 * 28).to(DEVICE), target.to(DEVICE)
            optimizer.zero_grad()
            output = model(data)
            loss = F.nll_loss(output, target)
            epoch_loss += loss.data
            loss.backward()
            optimizer.step()

        epoch_loss /= len(train_loader.dataset)
        print('Epoch {}: {}'.format(epoch, epoch_loss))
        trial.report(epoch_loss, step=epoch)

    model.eval()
    correct = 0
    with torch.no_grad():
        for batch_idx, (data, target) in enumerate(test_loader):
            data, target = data.view(-1, 28 * 28).to(DEVICE), target.to(DEVICE)
            output = model(data)
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    accuracy = correct / len(test_loader.dataset)
    return accuracy


if __name__ == '__main__':
    study = optuna.create_study(
        sampler=optuna.samplers.RandomSampler(),
        direction='maximize',
    )

    print('Number of trials', len(study.trials))
    for t in study.trials:
        print(t.number, t.params)

    study.optimize(objective, n_trials=100)
    
    evaluator = optuna.importance.MeanDecreaseImpurityImportanceEvaluator()
    importances = optuna.importance.get_param_importances(study, evaluator=evaluator)
    print('Importances', importances)
    
    fig = optuna.visualization.plot_param_importances(study, evaluator=evaluator)
    fig.show()