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 temp_dir/pytorch/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 temp_dir/pytorch/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 temp_dir/pytorch/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 temp_dir/pytorch/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:
        """
        Converts model outputs to predicted labels.
        For classification: uses argmax.
        For regression: flattens output.
        """
        if isinstance(task, OpenMLClassificationTask):
            return torch.argmax(output, dim=-1)
        elif isinstance(task, OpenMLRegressionTask):
            return output.view(-1)
        else:
            raise ValueError(f"Unsupported task type: {type(task)}")

    @staticmethod
    def _default_predict_proba(output: torch.Tensor, task: OpenMLTask) -> torch.Tensor:
        """
        Converts model outputs to probabilities using softmax.
        """
        if not isinstance(task, OpenMLClassificationTask):
            raise ValueError("predict_proba is only valid for classification tasks")

        return torch.nn.functional.softmax(output, dim=-1)

    @staticmethod
    def _default_sanitize(tensor: torch.Tensor) -> torch.Tensor:
        """
        Replaces NaNs with 1e-6 in-place if possible.
        """
        nan_mask = torch.isnan(tensor)
        if nan_mask.any():
            tensor = tensor.clone()  # clone only if modification is needed
            tensor[nan_mask] = 1e-6
        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,
            scheduler=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 temp_dir/pytorch/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 temp_dir/pytorch/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 temp_dir/pytorch/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,
        scheduler=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 temp_dir/pytorch/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=torch.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 temp_dir/pytorch/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",
        num_workers=0,
        batch_size = 64,
        **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.data_config.batch_size = batch_size
        self.handler: BaseDataHandler | None = data_handlers.get(type_of_data)
        self.num_workers = num_workers

        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, num_workers = self.num_workers 
        )

        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, num_workers=self.num_workers
        )
        val_loader = DataLoader(
            val, batch_size=self.data_config.batch_size, shuffle=False, num_workers=self.num_workers
        )

        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 temp_dir/pytorch/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 temp_dir/pytorch/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 temp_dir/pytorch/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 temp_dir/pytorch/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(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:
        """
        Converts predicted class indices into one-hot probability vectors matching OpenML class indices.
        """
        if not isinstance(classes, list):
            raise ValueError("Please convert model classes to list prior to calling this function.")

        n_samples = len(y)
        n_classes = len(classes)

        # Ensure y is an integer array
        y = np.asarray(y, dtype=np.int64)

        result = np.zeros((n_samples, n_classes), dtype=np.float32)
        result[np.arange(n_samples), y] = 1.0  # vectorized one-hot assignment
        return result

    def run_evaluation(self, task, data, model_classes):
        if not isinstance(task, OpenMLSupervisedTask):
            raise ValueError(task)

        self.model.eval()
        with torch.no_grad():

            # Run inference once
            logits = self.pred_test(task, self.model, data.test_dl, lambda x, _: x)

            # Get predictions
            pred_y = self.config.predict(torch.from_numpy(logits), task).numpy()

            proba_y = None
            if isinstance(task, OpenMLClassificationTask):
                try:
                    proba_y = self.config.predict_proba(torch.from_numpy(logits), task).numpy()
                except AttributeError:
                    if task.class_labels is None:
                        raise ValueError("The task has no class labels")
                    proba_y = self._prediction_to_probabilities(pred_y, list(task.class_labels))

                # Adjust shape if needed
                if task.class_labels is not None and proba_y.shape[1] != len(task.class_labels):
                    self.logger.warning("Mismatch in predicted probabilities and number of class labels.")
                    proba_y_new = np.zeros((proba_y.shape[0], len(task.class_labels)))
                    for idx, model_class in enumerate(model_classes):
                        if model_class < proba_y_new.shape[1]:
                            proba_y_new[:, model_class] = proba_y[:, idx]
                    proba_y = proba_y_new

                    if proba_y.shape[1] != len(task.class_labels):
                        message = f"Estimator only predicted for {proba_y.shape[1]}/{len(task.class_labels)} classes!"
                        warnings.warn(message)
                        self.logger.warning(message)

            elif isinstance(task, OpenMLRegressionTask):
                proba_y = None
            else:
                raise TypeError(type(task))

        print("Evaluation done")
        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 != torch.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, test_loader, predict_func):
        model.eval()
        device = self.config.device
        sanitize = self.config.sanitize

        all_outputs = []

        with torch.no_grad():
            for batch in test_loader:
                inputs = sanitize(batch).to(device)
                outputs = model(inputs)
                outputs = predict_func(outputs, task)
                all_outputs.append(outputs.detach().cpu())  # keep as tensor

        return torch.cat(all_outputs, dim=0).numpy()  # only convert once

add_callbacks()

Adds the user-defined callbacks to the list of callbacks

Source code in temp_dir/pytorch/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 temp_dir/pytorch/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 temp_dir/pytorch/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 temp_dir/pytorch/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