import json import pickle import numpy as np 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, rep_dim: int, eta: float = 1.0): """Inits DeepSAD with hyperparameter eta.""" self.rep_dim = rep_dim # representation dimension 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, "indices": None, "file_ids": None, "frame_ids": None, "file_names": None, # mapping of file_ids to file names }, "test": { "time": None, "indices": None, "file_ids": None, "frame_ids": None, "file_names": None, # mapping of file_ids to file names "exp_based": { "auc": None, "roc": None, "prc": None, "ap": None, "scores": None, "indices": None, "labels": None, "valid_mask": None, }, "manual_based": { "auc": None, "roc": None, "prc": None, "ap": None, "scores": None, "indices": None, "labels": None, "valid_mask": None, }, }, "inference": { "time": None, "indices": None, "file_ids": None, "frame_ids": None, "file_names": None, # mapping of file_ids to file names }, } 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, self.rep_dim) 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.c = self.trainer.c.cpu().data.numpy().tolist() # get as list # Store training results including indices self.results["train"]["time"] = self.trainer.train_time self.results["train"]["indices"] = self.trainer.train_indices self.results["train"]["file_ids"] = self.trainer.train_file_ids self.results["train"]["frame_ids"] = self.trainer.train_frame_ids # Get file names mapping for training data self.results["train"]["file_names"] = { file_id: dataset.get_file_name_from_idx(file_id) for file_id in np.unique(self.trainer.train_file_ids) } 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 ) scores, outputs = self.trainer.infer(dataset, self.net) # Store inference indices and mappings self.results["inference"]["time"] = self.trainer.inference_time self.results["inference"]["indices"] = self.trainer.inference_indices self.results["inference"]["file_ids"] = self.trainer.inference_file_ids self.results["inference"]["frame_ids"] = self.trainer.inference_frame_ids # Get file names mapping for inference data self.results["inference"]["file_names"] = { file_id: dataset.get_file_name_from_idx(file_id) for file_id in np.unique(self.trainer.inference_file_ids) } return scores, outputs 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) # Store all test indices and mappings self.results["test"]["time"] = self.trainer.test_time self.results["test"]["indices"] = self.trainer.test_indices self.results["test"]["file_ids"] = self.trainer.test_file_ids self.results["test"]["frame_ids"] = self.trainer.test_frame_ids # Get file names mapping for test data self.results["test"]["file_names"] = { file_id: dataset.get_file_name_from_idx(file_id) for file_id in np.unique(self.trainer.test_file_ids) } # Store experiment-based results self.results["test"]["exp_based"]["auc"] = self.trainer.test_auc_exp_based self.results["test"]["exp_based"]["roc"] = self.trainer.test_roc_exp_based self.results["test"]["exp_based"]["prc"] = self.trainer.test_prc_exp_based self.results["test"]["exp_based"]["ap"] = self.trainer.test_ap_exp_based self.results["test"]["exp_based"]["scores"] = self.trainer.test_scores_exp_based self.results["test"]["exp_based"]["indices"] = self.trainer.test_index_mapping[ "exp_based" ]["indices"] self.results["test"]["exp_based"]["labels"] = self.trainer.test_index_mapping[ "exp_based" ]["labels"] self.results["test"]["exp_based"]["valid_mask"] = ( self.trainer.test_index_mapping["exp_based"]["valid_mask"] ) # Store manual-based results self.results["test"]["manual_based"]["auc"] = self.trainer.test_auc_manual_based self.results["test"]["manual_based"]["roc"] = self.trainer.test_roc_manual_based self.results["test"]["manual_based"]["prc"] = self.trainer.test_prc_manual_based self.results["test"]["manual_based"]["ap"] = self.trainer.test_ap_manual_based self.results["test"]["manual_based"]["scores"] = ( self.trainer.test_scores_manual_based ) self.results["test"]["manual_based"]["indices"] = ( self.trainer.test_index_mapping["manual_based"]["indices"] ) self.results["test"]["manual_based"]["labels"] = ( self.trainer.test_index_mapping["manual_based"]["labels"] ) self.results["test"]["manual_based"]["valid_mask"] = ( self.trainer.test_index_mapping["manual_based"]["valid_mask"] ) 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, self.rep_dim) # 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) # Test self.ae_trainer.test(dataset, self.ae_net, k_fold_idx=k_fold_idx) # Get train results self.ae_results = { "train": { "time": self.ae_trainer.train_time, "indices": self.ae_trainer.train_indices, "labels_exp_based": self.ae_trainer.train_labels_exp_based, "labels_manual_based": self.ae_trainer.train_labels_manual_based, "semi_targets": self.ae_trainer.train_semi_targets, "file_ids": self.ae_trainer.train_file_ids, "frame_ids": self.ae_trainer.train_frame_ids, "scores": self.ae_trainer.train_scores, "loss": self.ae_trainer.train_loss, "file_names": { file_id: dataset.get_file_name_from_idx(file_id) for file_id in np.unique(self.ae_trainer.train_file_ids) }, }, "test": { "time": self.ae_trainer.test_time, "indices": self.ae_trainer.test_indices, "labels_exp_based": self.ae_trainer.test_labels_exp_based, "labels_manual_based": self.ae_trainer.test_labels_manual_based, "semi_targets": self.ae_trainer.test_semi_targets, "file_ids": self.ae_trainer.test_file_ids, "frame_ids": self.ae_trainer.test_frame_ids, "scores": self.ae_trainer.test_scores, "loss": self.ae_trainer.test_loss, "file_names": { file_id: dataset.get_file_name_from_idx(file_id) for file_id in np.unique(self.ae_trainer.test_file_ids) }, }, } # 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.rep_dim) 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_pkl): """Save autoencoder results dict to a JSON-file.""" with open(export_pkl, "wb") as fp: pickle.dump(self.ae_results, fp)