mt/Deep-SAD-PyTorch/src/DeepSAD.py

import json
import pickle

import torch

from base.base_dataset import BaseADDataset
from networks.main import build_autoencoder, build_network
from optim.ae_trainer import AETrainer
from optim.DeepSAD_trainer import DeepSADTrainer


class DeepSAD(object):
    """A class for the Deep SAD method.

    Attributes:
        eta: Deep SAD hyperparameter eta (must be 0 < eta).
        c: Hypersphere center c.
        net_name: A string indicating the name of the neural network to use.
        net: The neural network phi.
        trainer: DeepSADTrainer to train a Deep SAD model.
        optimizer_name: A string indicating the optimizer to use for training the Deep SAD network.
        ae_net: The autoencoder network corresponding to phi for network weights pretraining.
        ae_trainer: AETrainer to train an autoencoder in pretraining.
        ae_optimizer_name: A string indicating the optimizer to use for pretraining the autoencoder.
        results: A dictionary to save the results.
        ae_results: A dictionary to save the autoencoder results.
    """

    def __init__(self, eta: float = 1.0):
        """Inits DeepSAD with hyperparameter eta."""

        self.eta = eta
        self.c = None  # hypersphere center c

        self.net_name = None
        self.net = None  # neural network phi

        self.trainer = None
        self.optimizer_name = None

        self.ae_net = None  # autoencoder network for pretraining
        self.ae_trainer = None
        self.ae_optimizer_name = None

        self.results = {
            "train_time": None,
            "test_auc": None,
            "test_time": None,
            "test_scores": None,
        }

        self.ae_results = {"train_time": None, "test_auc": None, "test_time": None}

    def set_network(self, net_name):
        """Builds the neural network phi."""
        self.net_name = net_name
        self.net = build_network(net_name)

    def train(
        self,
        dataset: BaseADDataset,
        optimizer_name: str = "adam",
        lr: float = 0.001,
        n_epochs: int = 50,
        lr_milestones: tuple = (),
        batch_size: int = 128,
        weight_decay: float = 1e-6,
        device: str = "cuda",
        n_jobs_dataloader: int = 0,
        k_fold_idx: int = None,
    ):
        """Trains the Deep SAD model on the training data."""

        self.optimizer_name = optimizer_name
        self.trainer = DeepSADTrainer(
            self.c,
            self.eta,
            optimizer_name=optimizer_name,
            lr=lr,
            n_epochs=n_epochs,
            lr_milestones=lr_milestones,
            batch_size=batch_size,
            weight_decay=weight_decay,
            device=device,
            n_jobs_dataloader=n_jobs_dataloader,
        )
        # Get the model
        self.net = self.trainer.train(dataset, self.net, k_fold_idx=k_fold_idx)
        self.results["train_time"] = self.trainer.train_time
        self.c = self.trainer.c.cpu().data.numpy().tolist()  # get as list

    def inference(
        self, dataset: BaseADDataset, device: str = "cuda", n_jobs_dataloader: int = 0
    ):
        """Tests the Deep SAD model on the test data."""

        if self.trainer is None:
            self.trainer = DeepSADTrainer(
                self.c, self.eta, device=device, n_jobs_dataloader=n_jobs_dataloader
            )

        return self.trainer.infer(dataset, self.net)

    def test(
        self,
        dataset: BaseADDataset,
        device: str = "cuda",
        n_jobs_dataloader: int = 0,
        k_fold_idx: int = None,
    ):
        """Tests the Deep SAD model on the test data."""

        if self.trainer is None:
            self.trainer = DeepSADTrainer(
                self.c, self.eta, device=device, n_jobs_dataloader=n_jobs_dataloader
            )

        self.trainer.test(dataset, self.net, k_fold_idx=k_fold_idx)

        # Get results
        self.results["test_auc"] = self.trainer.test_auc
        self.results["test_roc"] = self.trainer.test_roc
        self.results["test_prc"] = self.trainer.test_prc
        self.results["test_ap"] = self.trainer.test_ap
        self.results["test_time"] = self.trainer.test_time
        self.results["test_scores"] = self.trainer.test_scores

    def pretrain(
        self,
        dataset: BaseADDataset,
        optimizer_name: str = "adam",
        lr: float = 0.001,
        n_epochs: int = 100,
        lr_milestones: tuple = (),
        batch_size: int = 128,
        weight_decay: float = 1e-6,
        device: str = "cuda",
        n_jobs_dataloader: int = 0,
        k_fold_idx: int = None,
    ):
        """Pretrains the weights for the Deep SAD network phi via autoencoder."""

        # Set autoencoder network
        self.ae_net = build_autoencoder(self.net_name)

        # Train
        self.ae_optimizer_name = optimizer_name
        self.ae_trainer = AETrainer(
            optimizer_name,
            lr=lr,
            n_epochs=n_epochs,
            lr_milestones=lr_milestones,
            batch_size=batch_size,
            weight_decay=weight_decay,
            device=device,
            n_jobs_dataloader=n_jobs_dataloader,
        )
        self.ae_net = self.ae_trainer.train(dataset, self.ae_net, k_fold_idx=k_fold_idx)

        # Get train results
        self.ae_results["train_time"] = self.ae_trainer.train_time

        # Test
        self.ae_trainer.test(dataset, self.ae_net, k_fold_idx=k_fold_idx)

        # Get test results
        self.ae_results["test_auc"] = self.ae_trainer.test_auc
        self.ae_results["test_time"] = self.ae_trainer.test_time

        # Initialize Deep SAD network weights from pre-trained encoder
        self.init_network_weights_from_pretraining()

    def init_network_weights_from_pretraining(self):
        """Initialize the Deep SAD network weights from the encoder weights of the pretraining autoencoder."""

        net_dict = self.net.state_dict()
        ae_net_dict = self.ae_net.state_dict()

        # Filter out decoder network keys
        ae_net_dict = {k: v for k, v in ae_net_dict.items() if k in net_dict}
        # Overwrite values in the existing state_dict
        net_dict.update(ae_net_dict)
        # Load the new state_dict
        self.net.load_state_dict(net_dict)

    def save_model(self, export_model, save_ae=True):
        """Save Deep SAD model to export_model."""

        net_dict = self.net.state_dict()
        ae_net_dict = self.ae_net.state_dict() if save_ae else None

        torch.save(
            {"c": self.c, "net_dict": net_dict, "ae_net_dict": ae_net_dict},
            export_model,
        )

    def load_model(self, model_path, load_ae=False, map_location="cpu"):
        """Load Deep SAD model from model_path."""

        model_dict = torch.load(model_path, map_location=map_location)

        self.c = model_dict["c"]
        self.net.load_state_dict(model_dict["net_dict"])

        # load autoencoder parameters if specified
        if load_ae:
            if self.ae_net is None:
                self.ae_net = build_autoencoder(self.net_name)
            self.ae_net.load_state_dict(model_dict["ae_net_dict"])

    def save_results(self, export_pkl):
        """Save results dict to a JSON-file."""
        with open(export_pkl, "wb") as fp:
            # json.dump(self.results, fp)
            pickle.dump(self.results, fp)

    def save_ae_results(self, export_json):
        """Save autoencoder results dict to a JSON-file."""
        with open(export_json, "w") as fp:
            json.dump(self.ae_results, fp)