Trainer

This module provides classes and methods to facilitate the configuration, data handling, training, and evaluation of machine learning models using PyTorch and OpenML datasets. The functionalities include: - Generation of default configurations for models. - Handling of image and tabular data. - Training and evaluating machine learning models. - Exporting trained models to ONNX format. - Managing data transformations and loaders.

This module provides classes and methods to facilitate the configuration, data handling, training, and evaluation of machine learning models using PyTorch and OpenML datasets. The functionalities include: - Generation of default configurations for models. - Handling of image and tabular data. - Training and evaluating machine learning models. - Exporting trained models to ONNX format. - Managing data transformations and loaders.

BaseDataHandler

BaseDataHandler class is an abstract base class for data handling operations.

Source code in openml_pytorch/trainer.py

class BaseDataHandler:
    """
    BaseDataHandler class is an abstract base class for data handling operations.
    """

    def prepare_data(
        self, X_train, y_train, X_val, y_val, data_config: SimpleNamespace
    ):
        raise NotImplementedError

    def prepare_test_data(self, X_test, data_config: SimpleNamespace):
        raise NotImplementedError

BasicTrainer

BasicTrainer class provides a simple training loop for PyTorch models.You pass in the model, loss function, optimizer, data loaders, and device. The fit method trains the model for the specified number of epochs.

Source code in openml_pytorch/trainer.py

class BasicTrainer:
    """
    BasicTrainer class provides a simple training loop for PyTorch models.You pass in the model, loss function, optimizer, data loaders, and device. The fit method trains the model for the specified number of epochs.
    """

    def __init__(
        self,
        model: Any,
        loss_fn: Any,
        opt: Any,
        dataloader_train: torch.utils.data.DataLoader,
        dataloader_test: torch.utils.data.DataLoader,
        device: torch.device,
    ):
        self.device = device
        self.model = model.to(self.device)
        self.loss_fn = loss_fn
        self.opt = opt(self.model.parameters())
        self.dataloader_train = dataloader_train
        self.dataloader_test = dataloader_test
        self.losses = {"train": [], "test": []}

    def train_step(self, x, y):
        self.model.train()
        self.opt.zero_grad()
        yhat = self.model(x)
        loss = self.loss_fn(yhat, y)
        loss.backward()
        self.opt.step()
        return loss.item()

    def test_step(self, x, y):
        self.model.eval()
        with torch.no_grad():
            yhat = self.model(x)
            loss = self.loss_fn(yhat, y)
        return loss.item()

    def fit(self, epochs):
        if self.dataloader_train is None:
            raise ValueError("dataloader_train is not set")
        if self.dataloader_test is None:
            raise ValueError("dataloader_test is not set")
        bar = tqdm(range(epochs), desc="Epochs")
        for epoch in bar:
            # train
            for x, y in self.dataloader_train:
                x, y = x.to(self.device), y.to(self.device)
                loss = self.train_step(x, y)
                self.losses["train"].append(loss)
            # test
            test_loss = 0
            for x, y in self.dataloader_test:
                x, y = x.to(self.device), y.to(self.device)
                test_loss += self.test_step(x, y)
                self.losses["test"].append(test_loss)
            bar.set_postfix(
                {"Train loss": loss, "Test loss": test_loss, "Epoch": epoch + 1}
            )

DataContainer

class DataContainer

A class to contain the training, validation, and test data loaders. This just makes it easier to access them when required.

Attributes: train_dl: DataLoader object for the training data. valid_dl: DataLoader object for the validation data. test_dl: Optional DataLoader object for the test data.

Source code in openml_pytorch/trainer.py

class DataContainer:
    """
    class DataContainer:
        A class to contain the training, validation, and test data loaders. This just makes it easier to access them when required.

        Attributes:
        train_dl: DataLoader object for the training data.
        valid_dl: DataLoader object for the validation data.
        test_dl: Optional DataLoader object for the test data.
    """

    def __init__(
        self,
        train_dl: DataLoader,
        valid_dl: torch.utils.data.DataLoader,
        test_dl: torch.utils.data.DataLoader = None,
    ):  # type: ignore
        self.train_dl, self.valid_dl = train_dl, valid_dl
        self.test_dl = test_dl

    @property
    def train_ds(self):
        return self.train_dl.dataset

    @property
    def valid_ds(self):
        return self.valid_dl.dataset

    @property
    def test_ds(self):
        return self.test_dl.dataset

DefaultConfigGenerator

DefaultConfigGenerator class provides various methods to generate default configurations.

Source code in openml_pytorch/trainer.py

class DefaultConfigGenerator:
    """
    DefaultConfigGenerator class provides various methods to generate default configurations.
    """

    @staticmethod
    def _default_loss_fn_gen(task: OpenMLTask) -> torch.nn.Module:
        """
        _default_loss_fn_gen returns a loss fn based on the task type - regressions use
        torch.nn.SmoothL1Loss while classifications use torch.nn.CrossEntropyLoss
        """
        if isinstance(task, OpenMLRegressionTask):
            return torch.nn.SmoothL1Loss()
        elif isinstance(task, OpenMLClassificationTask):
            return torch.nn.CrossEntropyLoss()
        else:
            raise ValueError(task)

    @staticmethod
    def _default_predict(output: torch.Tensor, task: OpenMLTask) -> torch.Tensor:
        """
        _default_predict turns the outputs into predictions by returning the argmax of the output tensor for classification tasks, and by flattening the prediction in case of the regression
        """
        output_axis = output.dim() - 1
        if isinstance(task, OpenMLClassificationTask):
            output = torch.argmax(output, dim=output_axis)
        elif isinstance(task, OpenMLRegressionTask):
            output = output.view(-1)
        else:
            raise ValueError(task)
        return output

    @staticmethod
    def _default_predict_proba(output: torch.Tensor, task: OpenMLTask) -> torch.Tensor:
        """
        _default_predict_proba turns the outputs into probabilities using softmax
        """
        output_axis = output.dim() - 1
        output = output.softmax(dim=output_axis)
        return output

    @staticmethod
    def _default_sanitize(tensor: torch.Tensor) -> torch.Tensor:
        """
        _default sanitizer replaces NaNs with 1e-6
        """
        tensor = torch.where(
            torch.isnan(tensor), torch.ones_like(tensor) * torch.tensor(1e-6), tensor
        )
        return tensor

    @staticmethod
    def _default_retype_labels(tensor: torch.Tensor, task: OpenMLTask) -> torch.Tensor:
        """
        _default_retype_labels changes the type of the tensor to long for classification tasks and to float for regression tasks
        """
        if isinstance(task, OpenMLClassificationTask):
            return tensor.long()
        elif isinstance(task, OpenMLRegressionTask):
            return tensor.float()
        else:
            raise ValueError(task)

    def get_device(
        self,
    ):
        """
        Checks if a GPU is available and returns the device to be used for training (cuda, mps or cpu)
        """
        if torch.cuda.is_available():
            device = torch.device("cuda")
        elif torch.backends.mps.is_available() and torch.backends.mps.is_built():
            device = torch.device("mps")
        else:
            device = torch.device("cpu")

        return device

    def default_image_transform(self):
        return Compose(
            [
                ToPILImage(),  # Convert tensor to PIL Image to ensure PIL Image operations can be applied.
                Lambda(convert_to_rgb),  # Convert PIL Image to RGB if it's not already.
                Resize((128, 128)),  # Resize the image.
                ToTensor(),  # Convert the PIL Image back to a tensor.
            ]
        )

    def default_image_transform_test(self):
        return Compose(
            [
                ToPILImage(),  # Convert tensor to PIL Image to ensure PIL Image operations can be applied.
                Lambda(convert_to_rgb),  # Convert PIL Image to RGB if it's not already.
                Resize((128, 128)),  # Resize the image.
                ToTensor(),  # Convert the PIL Image back to a tensor.
            ]
        )

    def return_model_config(self):
        """
        Returns a configuration object for the model
        """

        return SimpleNamespace(
            device=self.get_device(),
            loss_fn=self._default_loss_fn_gen,
            # predict turns the outputs of the model into actual predictions
            predict=self._default_predict,  # type: Callable[[torch.Tensor, OpenMLTask], torch.Tensor]
            # predict_proba turns the outputs of the model into probabilities for each class
            predict_proba=self._default_predict_proba,  # type: Callable[[torch.Tensor], torch.Tensor]
            # epoch_count represents the number of epochs the model should be trained for
            epoch_count=3,  # type: int,
            verbose=True,
            schduler=None,
            opt_kwargs={},
            scheduler_kwargs={},
        )

    def return_data_config(self):
        """
        Returns a configuration object for the data
        """
        return SimpleNamespace(
            type_of_data="image",
            perform_validation=False,
            # progress_callback is called when a training step is finished, in order to report the current progress
            # sanitize sanitizes the input data in order to ensure that models can be trained safely
            sanitize=self._default_sanitize,  # type: Callable[[torch.Tensor], torch.Tensor]
            # retype_labels changes the types of the labels in order to ensure type compatibility
            retype_labels=(
                self._default_retype_labels
            ),  # type: Callable[[torch.Tensor, OpenMLTask], torch.Tensor]
            # image_size is the size of the images that are fed into the model
            image_size=128,
            # batch_size represents the processing batch size for training
            batch_size=64,  # type: int
            data_augmentation=None,
            validation_split=0.1,
            transform=self.default_image_transform(),
            transform_test=self.default_image_transform_test(),
        )

get_device()

Checks if a GPU is available and returns the device to be used for training (cuda, mps or cpu)

Source code in openml_pytorch/trainer.py

def get_device(
    self,
):
    """
    Checks if a GPU is available and returns the device to be used for training (cuda, mps or cpu)
    """
    if torch.cuda.is_available():
        device = torch.device("cuda")
    elif torch.backends.mps.is_available() and torch.backends.mps.is_built():
        device = torch.device("mps")
    else:
        device = torch.device("cpu")

    return device

return_data_config()

Returns a configuration object for the data

Source code in openml_pytorch/trainer.py

def return_data_config(self):
    """
    Returns a configuration object for the data
    """
    return SimpleNamespace(
        type_of_data="image",
        perform_validation=False,
        # progress_callback is called when a training step is finished, in order to report the current progress
        # sanitize sanitizes the input data in order to ensure that models can be trained safely
        sanitize=self._default_sanitize,  # type: Callable[[torch.Tensor], torch.Tensor]
        # retype_labels changes the types of the labels in order to ensure type compatibility
        retype_labels=(
            self._default_retype_labels
        ),  # type: Callable[[torch.Tensor, OpenMLTask], torch.Tensor]
        # image_size is the size of the images that are fed into the model
        image_size=128,
        # batch_size represents the processing batch size for training
        batch_size=64,  # type: int
        data_augmentation=None,
        validation_split=0.1,
        transform=self.default_image_transform(),
        transform_test=self.default_image_transform_test(),
    )

return_model_config()

Returns a configuration object for the model

Source code in openml_pytorch/trainer.py

def return_model_config(self):
    """
    Returns a configuration object for the model
    """

    return SimpleNamespace(
        device=self.get_device(),
        loss_fn=self._default_loss_fn_gen,
        # predict turns the outputs of the model into actual predictions
        predict=self._default_predict,  # type: Callable[[torch.Tensor, OpenMLTask], torch.Tensor]
        # predict_proba turns the outputs of the model into probabilities for each class
        predict_proba=self._default_predict_proba,  # type: Callable[[torch.Tensor], torch.Tensor]
        # epoch_count represents the number of epochs the model should be trained for
        epoch_count=3,  # type: int,
        verbose=True,
        schduler=None,
        opt_kwargs={},
        scheduler_kwargs={},
    )

Learner

A class to store the model, optimizer, loss_fn, and data loaders for training and evaluation.

Source code in openml_pytorch/trainer.py

class Learner:
    """
    A class to store the model, optimizer, loss_fn, and data loaders for training and evaluation.
    """

    def __init__(
        self, model, opt, loss_fn, scheduler, data, model_classes, device="cpu"
    ):
        (
            self.model,
            self.opt,
            self.loss_fn,
            self.scheduler,
            self.data,
            self.model_classes,
            self.device,
        ) = (model, opt, loss_fn, scheduler, data, model_classes, device)

OpenMLDataModule

Source code in openml_pytorch/trainer.py

class OpenMLDataModule:
    def __init__(
        self,
        type_of_data="image",
        filename_col="Filename",
        file_dir="images",
        target_mode="categorical",
        transform=None,
        transform_test=None,
        target_column="encoded_labels",
        **kwargs,
    ):
        self.config_gen = DefaultConfigGenerator()
        self.data_config = self.config_gen.return_data_config()
        self.data_config.type_of_data = type_of_data
        self.data_config.filename_col = filename_col
        self.data_config.file_dir = file_dir
        self.data_config.target_mode = target_mode
        self.data_config.target_column = target_column
        self.handler: BaseDataHandler | None = data_handlers.get(type_of_data)

        if transform is not None:
            self.data_config.transform = transform
        if transform_test is not None:
            self.data_config.transform_test = transform_test

        if not self.handler:
            raise ValueError(f"Data type {type_of_data} not supported.")

    def get_data(
        self,
        X_train: pd.DataFrame,
        y_train: Optional[pd.Series],
        X_test: pd.DataFrame,
        task,
    ):
        # Split the training data
        X_train_train, X_val, y_train_train, y_val = self.split_training_data(
            X_train, y_train
        )

        y_train_train, y_val, model_classes = self.encode_labels(y_train_train, y_val)

        # Use handler to prepare datasets
        train_loader, val_loader = self.prepare_datasets_for_training_and_validation(
            X_train_train, X_val, y_train_train, y_val
        )

        # Prepare test data
        test_loader = self.process_test_data(X_test)

        return DataContainer(train_loader, val_loader, test_loader), model_classes

    def process_test_data(self, X_test):
        test = self.handler.prepare_test_data(X_test, self.data_config)
        test_loader = DataLoader(
            test, batch_size=self.data_config.batch_size, shuffle=False
        )

        return test_loader

    def prepare_datasets_for_training_and_validation(
        self, X_train_train, X_val, y_train_train, y_val
    ):
        if self.handler is None:
            raise ValueError(
                f"Data type {self.data_config.type_of_data} not supported."
            )

        train, val = self.handler.prepare_data(
            X_train_train, y_train_train, X_val, y_val, self.data_config
        )

        train_loader = DataLoader(
            train, batch_size=self.data_config.batch_size, shuffle=True
        )
        val_loader = DataLoader(
            val, batch_size=self.data_config.batch_size, shuffle=False
        )

        return train_loader, val_loader

    def encode_labels(self, y_train_train, y_val):
        """
        Encode the labels for categorical data
        """
        if self.data_config.target_mode == "categorical":
            label_encoder = preprocessing.LabelEncoder().fit(y_train_train)
            y_train_train = pd.Series(label_encoder.transform(y_train_train))
            y_val = pd.Series(label_encoder.transform(y_val))
            # Determine model classes
            model_classes = (
                label_encoder.classes_
                if self.data_config.target_mode == "categorical"
                else None
            )

        return y_train_train, y_val, model_classes

    def split_training_data(self, X_train, y_train):
        if type(y_train) != pd.Series:
            y_train = pd.Series(y_train)

        X_train_train, X_val, y_train_train, y_val = train_test_split(
            X_train,
            y_train,
            test_size=self.data_config.validation_split,
            shuffle=True,
            stratify=y_train,
            random_state=0,
        )

        return X_train_train, X_val, y_train_train, y_val

encode_labels(y_train_train, y_val)

Encode the labels for categorical data

Source code in openml_pytorch/trainer.py

def encode_labels(self, y_train_train, y_val):
    """
    Encode the labels for categorical data
    """
    if self.data_config.target_mode == "categorical":
        label_encoder = preprocessing.LabelEncoder().fit(y_train_train)
        y_train_train = pd.Series(label_encoder.transform(y_train_train))
        y_val = pd.Series(label_encoder.transform(y_val))
        # Determine model classes
        model_classes = (
            label_encoder.classes_
            if self.data_config.target_mode == "categorical"
            else None
        )

    return y_train_train, y_val, model_classes

OpenMLImageHandler

Bases: BaseDataHandler

OpenMLImageHandler is a class that extends BaseDataHandler to handle image data from OpenML datasets.

Source code in openml_pytorch/trainer.py

class OpenMLImageHandler(BaseDataHandler):
    """
    OpenMLImageHandler is a class that extends BaseDataHandler to handle image data from OpenML datasets.
    """

    def prepare_data(self, X_train, y_train, X_val, y_val, data_config=None):
        train = OpenMLImageDataset(
            image_dir=data_config.file_dir,
            X=X_train,
            y=y_train,
            transform_x=data_config.transform,
            image_size=data_config.image_size,
        )
        val = OpenMLImageDataset(
            image_dir=data_config.file_dir,
            X=X_val,
            y=y_val,
            transform_x=data_config.transform_test,
            image_size=data_config.image_size,
        )
        return train, val

    def prepare_test_data(self, X_test, data_config=None):
        test = OpenMLImageDataset(
            image_dir=data_config.file_dir,
            X=X_test,
            y=None,
            transform_x=data_config.transform_test,
            image_size=data_config.image_size,
        )
        return test

OpenMLTabularHandler

Bases: BaseDataHandler

OpenMLTabularHandler is a class that extends BaseDataHandler to handle tabular data from OpenML datasets.

Source code in openml_pytorch/trainer.py

class OpenMLTabularHandler(BaseDataHandler):
    """
    OpenMLTabularHandler is a class that extends BaseDataHandler to handle tabular data from OpenML datasets.
    """

    def prepare_data(self, X_train, y_train, X_val, y_val, data_config=None):
        train = OpenMLTabularDataset(X=X_train, y=y_train)
        val = OpenMLTabularDataset(X=X_val, y=y_val)
        return train, val

    def prepare_test_data(self, X_test, data_config=None):
        test = OpenMLTabularDataset(X=X_test, y=None)
        return test

OpenMLTrainerModule

Source code in openml_pytorch/trainer.py

class OpenMLTrainerModule:
    def _default_progress_callback(
        self, fold: int, rep: int, epoch: int, step: int, loss: float, accuracy: float
    ):
        # todo : move this into callback
        """
                _default_progress_callback reports the current fold, rep, epoch, step and loss for every
        training iteration to the default logger
        """
        self.logger.info(
            "[%d, %d, %d, %d] loss: %.4f, accuracy: %.4f"
            % (fold, rep, epoch, step, loss, accuracy)
        )

    def __init__(
        self,
        experiment_name: str,
        data_module: OpenMLDataModule,
        opt: Callable = torch.optim.AdamW,
        opt_kwargs: dict = {"lr": 3e-4, "weight_decay": 1e-4},
        loss_fn: Callable = torch.nn.CrossEntropyLoss,
        loss_fn_kwargs: dict = {},
        callbacks: List[Callback] = [],
        use_tensorboard: bool = True,
        metrics: List[Callable] = [],
        scheduler=torch.optim.lr_scheduler.CosineAnnealingLR,
        scheduler_kwargs: dict = {"T_max": 50, "eta_min": 1e-6},
        **kwargs,
    ):
        self.experiment_name = experiment_name
        self.config_gen = DefaultConfigGenerator()
        self.model_config = self.config_gen.return_model_config()
        self.data_module = data_module
        self.callbacks = callbacks
        self.metrics = metrics

        self.config = SimpleNamespace(
            **{**self.model_config.__dict__, **self.data_module.data_config.__dict__}
        )
        # update the config with the user defined values
        self.config.__dict__.update(kwargs)
        self.config.opt = opt
        self.config.opt_kwargs = opt_kwargs
        self.config.loss_fn_kwargs = loss_fn_kwargs
        if loss_fn is not None:
            self.loss_fn = loss_fn(**self.config.loss_fn_kwargs)
        self.config.progress_callback = self._default_progress_callback
        self.logger: logging.Logger = logging.getLogger(__name__)

        self.user_defined_measures = OrderedDict()
        # Tensorboard support
        timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        self.tensorboard_writer = None

        if use_tensorboard:
            self.tensorboard_writer = SummaryWriter(
                comment=experiment_name,
                log_dir=f"tensorboard_logs/{experiment_name}/{timestamp}",
            )

        self.loss = 0
        self.training_state = True

        self.phases = [0.2, 0.8]
        self.config.scheduler = scheduler
        self.config.scheduler_kwargs = scheduler_kwargs

        # Add default callbacks
        self.cbfs = [
            Recorder,
            partial(AvgStatsCallback, self.metrics),
            # partial(ParamScheduler, "lr", self.scheds),
            partial(PutDataOnDeviceCallback, self.config.device),
        ]
        if self.tensorboard_writer is not None:
            self.cbfs.append(partial(TensorBoardCallback, self.tensorboard_writer))

        self.add_callbacks()

    def export_to_onnx(self, model_copy):
        """
        Converts the model to ONNX format. Uses a hack for now (global variable) to get the sample input.
        """
        global sample_input
        f = io.BytesIO()
        torch.onnx.export(model_copy, sample_input, f)
        onnx_model = onnx.load_model_from_string(f.getvalue())
        onnx_ = onnx_model.SerializeToString()
        global last_models
        last_models = onnx_
        return onnx_

    def export_to_netron(self, onnx_file_name: str = f"model.onnx"):
        """
        Exports the model to ONNX format and serves it using netron.
        """
        if self.onnx_model is None:
            try:
                self.onnx_model = self.export_to_onnx(self.model)
            except Exception as e:
                raise ValueError("Model is not defined")
            raise ValueError("Model is not defined")

        # write the onnx model to a file
        with open(onnx_file_name, "wb") as f:
            f.write(self.onnx_model)
            print(f"Writing onnx model to {onnx_file_name}. Delete if neeeded")

        # serve with netro
        netron.start(onnx_file_name)

    def run_model_on_fold(
        self,
        model: torch.nn.Module,
        task: OpenMLTask,
        X_train: pd.DataFrame,
        rep_no: int,
        fold_no: int,
        y_train: Optional[pd.Series],
        X_test: pd.DataFrame,
    ):
        # if task has no class labels, we assign the class labels to be the unique values in the training set
        if task.class_labels is None:
            task.class_labels = y_train.unique()

        # Add the user defined callbacks

        self.model = copy.deepcopy(model)

        try:
            data, model_classes = self.run_training(task, X_train, y_train, X_test)

        except AttributeError as e:
            # typically happens when training a regressor8 on classification task
            raise PyOpenMLError(str(e))

        # In supervised learning this returns the predictions for Y
        pred_y, proba_y = self.run_evaluation(task, data, model_classes)

        # Convert predictions to class labels
        if task.class_labels is not None:
            pred_y = [task.class_labels[i] for i in pred_y]

        # Convert model to onnx
        onnx_ = self.export_to_onnx(self.model)

        # Hack to store the last model for ONNX conversion
        global last_models
        # last_models = onnx_
        self.onnx_model = onnx_

        return pred_y, proba_y, self.user_defined_measures, None

    def check_config(self):
        raise NotImplementedError

    def _prediction_to_probabilities(
        self, y: np.ndarray, classes: List[Any]
    ) -> np.ndarray:
        """Transforms predicted probabilities to match with OpenML class indices.

        Parameters
        ----------
        y : np.ndarray
            Predicted probabilities (possibly omitting classes if they were not present in the
            training data).
        model_classes : list
            List of classes known_predicted by the model, ordered by their index.

        Returns
        -------
        np.ndarray
        """
        # y: list or numpy array of predictions
        # model_classes: mapping from original array id to
        # prediction index id
        if not isinstance(classes, list):
            raise ValueError(
                "please convert model classes to list prior to calling this fn"
            )
        result = np.zeros((len(y), len(classes)), dtype=np.float32)
        for obs, prediction_idx in enumerate(y):
            result[obs][prediction_idx] = 1.0
        return result

    def run_evaluation(self, task, data, model_classes):
        if isinstance(task, OpenMLSupervisedTask):
            self.model.eval()
            pred_y = self.pred_test(task, self.model, data.test_dl, self.config.predict)
        else:
            raise ValueError(task)

        if isinstance(task, OpenMLClassificationTask):
            try:
                self.model.eval()
                proba_y = self.pred_test(
                    task, self.model, data.test_dl, self.config.predict_proba
                )

            except AttributeError:
                if task.class_labels is not None:
                    proba_y = self._prediction_to_probabilities(
                        pred_y, list(task.class_labels)
                    )
                else:
                    raise ValueError("The task has no class labels")

            if task.class_labels is not None:
                if proba_y.shape[1] != len(task.class_labels):
                    # Remap the probabilities in case there was a class missing
                    # at training time. By default, the classification targets
                    # are mapped to be zero-based indices to the actual classes.
                    # Therefore, the model_classes contain the correct indices to
                    # the correct probability array. Example:
                    # classes in the dataset: 0, 1, 2, 3, 4, 5
                    # classes in the training set: 0, 1, 2, 4, 5
                    # then we need to add a column full of zeros into the probabilities
                    # for class 3 because the rest of the library expects that the
                    # probabilities are ordered the same way as the classes are ordered).
                    print("The size of the tensor of predicted probabilities does not match the number of classes. Check the shape of the output of your model.")
                    proba_y_new = np.zeros((proba_y.shape[0], len(task.class_labels)))
                    for idx, model_class in enumerate(model_classes):
                        proba_y_new[:, model_class] = proba_y[:, idx]
                    proba_y = proba_y_new

                if proba_y.shape[1] != len(task.class_labels):
                    message = "Estimator only predicted for {}/{} classes!".format(
                        proba_y.shape[1],
                        len(task.class_labels),
                    )
                    warnings.warn(message)
                    self.logger.warning(message)
            else:
                raise ValueError("The task has no class labels")

        elif isinstance(task, OpenMLRegressionTask):
            proba_y = None

        else:
            raise TypeError(type(task))
        return pred_y, proba_y

    def run_training(self, task, X_train, y_train, X_test):
        if isinstance(task, OpenMLSupervisedTask) or isinstance(
            task, OpenMLClassificationTask
        ):
            self.opt = self.config.opt(
                self.model.parameters(), **self.config.opt_kwargs
            )
            self.scheduler = self.config.scheduler(
                optimizer=self.opt, **self.config.scheduler_kwargs
            )
            if self.loss_fn is None:
                self.loss_fn = self.config.loss_fn(task)
            self.device = self.config.device

            if self.config.device != "cpu":
                self.loss_fn = self.loss_fn.to(self.config.device)

            self.data, self.model_classes = self.data_module.get_data(
                X_train, y_train, X_test, task
            )
            self.learn = Learner(
                model=self.model,
                opt=self.opt,
                loss_fn=self.loss_fn,
                scheduler=self.scheduler,
                data=self.data,
                model_classes=self.model_classes,
            )
            self.learn.device = self.device
            self.learn.model.to(self.device)
            gc.collect()

            self.runner = ModelRunner(cb_funcs=self.cbfs)

            # some additional default callbacks
            self.stats = self.runner.cbs[1]
            self.plot_loss = self.stats.plot_loss
            self.plot_lr = self.stats.plot_lr
            self.plot_metric = self.stats.plot_metric
            self.plot_all_metrics = self.stats.plot_all_metrics

            self.learn.model.train()
            self.runner.fit(epochs=self.config.epoch_count, learn=self.learn)
            self.learn.model.eval()

            self.lrs = self.runner.cbs[1].lrs

            print("Loss", self.runner.loss)
        else:
            raise Exception("OpenML Task type not supported")
        return self.data, self.model_classes

    def add_callbacks(self):
        """
        Adds the user-defined callbacks to the list of callbacks
        """
        if self.callbacks is not None and len(self.callbacks) > 0:
            for callback in self.callbacks:
                if callback not in self.cbfs:
                    self.cbfs.append(callback)
                else:
                    # replace the callback with the new one in the same position
                    self.cbfs[self.cbfs.index(callback)] = callback

    def pred_test(self, task, model_copy, test_loader, predict_func):
        probabilities = []
        for batch_idx, inputs in enumerate(test_loader):
            inputs = self.config.sanitize(inputs)
            # if torch.cuda.is_available():
            inputs = inputs.to(self.config.device)

            # Perform inference on the batch
            pred_y_batch = model_copy(inputs)
            pred_y_batch = predict_func(pred_y_batch, task)
            pred_y_batch = pred_y_batch.cpu().detach().numpy()

            probabilities.append(pred_y_batch)

            # Concatenate probabilities from all batches
        pred_y = np.concatenate(probabilities, axis=0)
        return pred_y

add_callbacks()

Adds the user-defined callbacks to the list of callbacks

Source code in openml_pytorch/trainer.py

def add_callbacks(self):
    """
    Adds the user-defined callbacks to the list of callbacks
    """
    if self.callbacks is not None and len(self.callbacks) > 0:
        for callback in self.callbacks:
            if callback not in self.cbfs:
                self.cbfs.append(callback)
            else:
                # replace the callback with the new one in the same position
                self.cbfs[self.cbfs.index(callback)] = callback

export_to_netron(onnx_file_name=f'model.onnx')

Exports the model to ONNX format and serves it using netron.

Source code in openml_pytorch/trainer.py

def export_to_netron(self, onnx_file_name: str = f"model.onnx"):
    """
    Exports the model to ONNX format and serves it using netron.
    """
    if self.onnx_model is None:
        try:
            self.onnx_model = self.export_to_onnx(self.model)
        except Exception as e:
            raise ValueError("Model is not defined")
        raise ValueError("Model is not defined")

    # write the onnx model to a file
    with open(onnx_file_name, "wb") as f:
        f.write(self.onnx_model)
        print(f"Writing onnx model to {onnx_file_name}. Delete if neeeded")

    # serve with netro
    netron.start(onnx_file_name)

export_to_onnx(model_copy)

Converts the model to ONNX format. Uses a hack for now (global variable) to get the sample input.

Source code in openml_pytorch/trainer.py

def export_to_onnx(self, model_copy):
    """
    Converts the model to ONNX format. Uses a hack for now (global variable) to get the sample input.
    """
    global sample_input
    f = io.BytesIO()
    torch.onnx.export(model_copy, sample_input, f)
    onnx_model = onnx.load_model_from_string(f.getvalue())
    onnx_ = onnx_model.SerializeToString()
    global last_models
    last_models = onnx_
    return onnx_

convert_to_rgb(image)

Converts an image to RGB mode if it is not already in that mode.

Parameters: image (PIL.Image): The image to be converted.

Returns: PIL.Image: The converted image in RGB mode.

Source code in openml_pytorch/trainer.py

def convert_to_rgb(image):
    """
    Converts an image to RGB mode if it is not already in that mode.

    Parameters:
    image (PIL.Image): The image to be converted.

    Returns:
    PIL.Image: The converted image in RGB mode.
    """
    if image.mode != "RGB":
        return image.convert("RGB")
    return image