Source code for ax.benchmark.problems.hpo.torchvision

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from typing import Any, Dict

from ax.benchmark.problems.hpo.pytorch_cnn import (
    PyTorchCNNRunner,
    PyTorchCNNBenchmarkProblem,
)
from ax.core.runner import Runner
from ax.exceptions.core import UserInputError
from ax.utils.common.typeutils import checked_cast

try:  # We don't require TorchVision by default.
    from torchvision import transforms, datasets

    _REGISTRY = {"MNIST": datasets.MNIST, "FashionMNIST": datasets.FashionMNIST}
except ModuleNotFoundError:
    transforms = None
    datasets = None
    _REGISTRY = {}


[docs]class PyTorchCNNTorchvisionBenchmarkProblem(PyTorchCNNBenchmarkProblem):
[docs]    @classmethod
    def from_dataset_name(cls, name: str) -> "PyTorchCNNTorchvisionBenchmarkProblem":
        if name not in _REGISTRY:
            raise UserInputError(
                f"Unrecognized torchvision dataset {name}. Please ensure it is listed"
                "in PyTorchCNNTorchvisionBenchmarkProblem registry."
            )
        dataset_fn = _REGISTRY[name]

        train_set = dataset_fn(
            root="./data",
            train=True,
            download=True,
            transform=transforms.ToTensor(),
        )

        test_set = dataset_fn(
            root="./data",
            train=False,
            download=True,
            transform=transforms.ToTensor(),
        )

        problem = cls.from_datasets(name=name, train_set=train_set, test_set=test_set)
        runner = PyTorchCNNTorchvisionRunner(
            name=name, train_set=train_set, test_set=test_set
        )

        return cls(
            name=f"HPO_PyTorchCNN_Torchvision::{name}",
            search_space=problem.search_space,
            optimization_config=problem.optimization_config,
            runner=runner,
            optimal_value=problem.optimal_value,
        )


[docs]class PyTorchCNNTorchvisionRunner(PyTorchCNNRunner):
    """
    A subclass to aid in serialization. This allows us to save only the name of the
    dataset and reload it from TorchVision at deserialization time.
    """

[docs]    @classmethod
    def serialize_init_args(cls, runner: Runner) -> Dict[str, Any]:
        pytorch_cnn_runner = checked_cast(PyTorchCNNRunner, runner)

        return {"name": pytorch_cnn_runner.name}

[docs]    @classmethod
    def deserialize_init_args(cls, args: Dict[str, Any]) -> Dict[str, Any]:
        name = args["name"]

        dataset_fn = _REGISTRY[name]

        train_set = dataset_fn(
            root="./data",
            train=True,
            download=True,
            transform=transforms.ToTensor(),
        )

        test_set = dataset_fn(
            root="./data",
            train=False,
            download=True,
            transform=transforms.ToTensor(),
        )

        return {"name": name, "train_set": train_set, "test_set": test_set}