diff --git a/abl/models/basic_model.py b/abl/models/basic_model.py index 9050365..5a20901 100644 --- a/abl/models/basic_model.py +++ b/abl/models/basic_model.py @@ -15,21 +15,51 @@ import sys sys.path.append("..") import torch -from torch.utils.data import Dataset +import numpy +from torch.utils.data import Dataset, DataLoader import os from multiprocessing import Pool +from typing import List, Any, T, Tuple, Optional, Callable class BasicDataset(Dataset): - def __init__(self, X, Y): + def __init__(self, X: List[Any], Y: List[Any]): + """Initialize a basic dataset. + + Parameters + ---------- + X : List[Any] + A list of objects representing the input data. + Y : List[Any] + A list of objects representing the output data. + """ self.X = X self.Y = Y def __len__(self): + """Return the length of the dataset. + + Returns + ------- + int + The length of the dataset. + """ return len(self.X) - def __getitem__(self, index): + def __getitem__(self, index: int) -> Tuple(Any, Any): + """Get an item from the dataset. + + Parameters + ---------- + index : int + The index of the item to retrieve. + + Returns + ------- + Tuple(Any, Any) + A tuple containing the input and output data at the specified index. + """ assert index < len(self), "index range error" img = self.X[index] @@ -39,17 +69,50 @@ class BasicDataset(Dataset): class XYDataset(Dataset): - def __init__(self, X, Y, transform=None): + def __init__(self, X: List[Any], Y: List[int], transform: Callable[...] = None): + """ + Initialize the dataset used for classification task. + + Parameters + ---------- + X : List[Any] + The input data. + Y : List[int] + The target data. + transform : callable, optional + A function/transform that takes in an object and returns a transformed version. Defaults to None. + """ self.X = X self.Y = torch.LongTensor(Y) self.n_sample = len(X) self.transform = transform - def __len__(self): + def __len__(self) -> int: + """ + Return the length of the dataset. + + Returns + ------- + int + The length of the dataset. + """ return len(self.X) - def __getitem__(self, index): + def __getitem__(self, index: int) -> Tuple[Any, torch.Tensor]: + """ + Get the item at the given index. + + Parameters + ---------- + index : int + The index of the item to get. + + Returns + ------- + Tuple[Any, torch.Tensor] + A tuple containing the object and its label. + """ assert index < len(self), "index range error" img = self.X[index] @@ -70,20 +133,102 @@ class FakeRecorder: class BasicModel: + """ + Wrap NN models into the form of an sklearn estimator + + Parameters + ---------- + model : torch.nn.Module + The PyTorch model to be trained or used for prediction. + criterion : torch.nn.Module + The loss function used for training. + optimizer : torch.nn.Module + The optimizer used for training. + device : torch.device + The device on which the model will be trained or used for prediction. + batch_size : int, optional + The batch size used for training, by default 1. + num_epochs : int, optional + The number of epochs used for training, by default 1. + stop_loss : Optional[float], optional + The loss value at which to stop training, by default 0.01. + num_workers : int, optional + The number of workers used for loading data, by default 0. + save_interval : Optional[int], optional + The interval at which to save the model during training, by default None. + save_dir : Optional[str], optional + The directory in which to save the model during training, by default None. + transform : Callable[..., Any], optional + The transformation function used for data augmentation, by default None. + collate_fn : Callable[[List[T]], Any], optional + The function used to collate data, by default None. + recorder : Any, optional + The recorder used to record training progress, by default None. + + Attributes + ---------- + model : torch.nn.Module + The PyTorch model to be trained or used for prediction. + batch_size : int + The batch size used for training. + num_epochs : int + The number of epochs used for training. + stop_loss : Optional[float] + The loss value at which to stop training. + num_workers : int + The number of workers used for loading data. + criterion : torch.nn.Module + The loss function used for training. + optimizer : torch.nn.Module + The optimizer used for training. + transform : Callable[..., Any] + The transformation function used for data augmentation. + device : torch.device + The device on which the model will be trained or used for prediction. + recorder : Any + The recorder used to record training progress. + save_interval : Optional[int] + The interval at which to save the model during training. + save_dir : Optional[str] + The directory in which to save the model during training. + collate_fn : Callable[[List[T]], Any] + The function used to collate data. + + Methods + ------- + fit(data_loader=None, X=None, y=None) + Train the model. + train_epoch(data_loader) + Train the model for one epoch. + predict(data_loader=None, X=None, print_prefix="") + Predict the class of the input data. + predict_proba(data_loader=None, X=None, print_prefix="") + Predict the probability of each class for the input data. + val(data_loader=None, X=None, y=None, print_prefix="") + Validate the model. + score(data_loader=None, X=None, y=None, print_prefix="") + Score the model. + _data_loader(X, y=None) + Load data. + save(epoch_id, save_dir="") + Save the model. + load(epoch_id, load_dir="") + Load the model. + """ def __init__( self, - model, - criterion, - optimizer, - device, - batch_size=1, - num_epochs=1, - stop_loss=0.01, - num_workers=0, - save_interval=None, - save_dir=None, - transform=None, - collate_fn=None, + model: torch.nn.Module, + criterion: torch.nn.Module, + optimizer: torch.nn.Module, + device: torch.device, + batch_size: int = 1, + num_epochs: int = 1, + stop_loss: Optional[float] = 0.01, + num_workers: int = 0, + save_interval: Optional[int] = None, + save_dir: Optional[str] = None, + transform: Callable[...] = None, + collate_fn: Callable[[List[T]], Any] = None, recorder=None, ): @@ -106,7 +251,6 @@ class BasicModel: self.save_interval = save_interval self.save_dir = save_dir self.collate_fn = collate_fn - pass def _fit(self, data_loader, n_epoch, stop_loss): recorder = self.recorder @@ -126,12 +270,44 @@ class BasicModel: recorder.print("Model fitted, minimal loss is ", min_loss) return loss_value - def fit(self, data_loader=None, X=None, y=None): + def fit( + self, data_loader: DataLoader = None, X: List[Any] = None, y: List[int] = None + ) -> float: + """ + Train the model. + + Parameters + ---------- + data_loader : DataLoader, optional + The data loader used for training, by default None + X : List[Any], optional + The input data, by default None + y : List[int], optional + The target data, by default None + + Returns + ------- + float + The loss value of the trained model. + """ if data_loader is None: data_loader = self._data_loader(X, y) return self._fit(data_loader, self.num_epochs, self.stop_loss) - def train_epoch(self, data_loader): + def train_epoch(self, data_loader: DataLoader): + """ + Train the model for one epoch. + + Parameters + ---------- + data_loader : DataLoader + The data loader used for training. + + Returns + ------- + float + The loss value of the trained model. + """ model = self.model criterion = self.criterion optimizer = self.optimizer @@ -169,7 +345,29 @@ class BasicModel: return torch.cat(results, axis=0) - def predict(self, data_loader=None, X=None, print_prefix=""): + def predict( + self, + data_loader: DataLoader = None, + X: List[Any] = None, + print_prefix: str = "", + ) -> numpy.ndarray: + """ + Predict the class of the input data. + + Parameters + ---------- + data_loader : DataLoader, optional + The data loader used for prediction, by default None + X : List[Any], optional + The input data, by default None + print_prefix : str, optional + The prefix used for printing, by default "" + + Returns + ------- + numpy.ndarray + The predicted class of the input data. + """ recorder = self.recorder recorder.print("Start Predict Class ", print_prefix) @@ -177,9 +375,31 @@ class BasicModel: data_loader = self._data_loader(X) return self._predict(data_loader).argmax(axis=1).cpu().numpy() - def predict_proba(self, data_loader=None, X=None, print_prefix=""): + def predict_proba( + self, + data_loader: DataLoader = None, + X: List[Any] = None, + print_prefix: str = "", + ) -> numpy.ndarray: + """ + Predict the probability of each class for the input data. + + Parameters + ---------- + data_loader : DataLoader, optional + The data loader used for prediction, by default None + X : List[Any], optional + The input data, by default None + print_prefix : str, optional + The prefix used for printing, by default "" + + Returns + ------- + numpy.ndarray + The predicted probability of each class for the input data. + """ recorder = self.recorder - # recorder.print('Start Predict Probability ', print_prefix) + recorder.print("Start Predict Probability ", print_prefix) if data_loader is None: data_loader = self._data_loader(X) @@ -215,7 +435,32 @@ class BasicModel: return mean_loss, accuracy - def val(self, data_loader=None, X=None, y=None, print_prefix=""): + def val( + self, + data_loader: DataLoader = None, + X: List[Any] = None, + y: List[int] = None, + print_prefix: str = "", + ) -> float: + """ + Validate the model. + + Parameters + ---------- + data_loader : DataLoader, optional + The data loader used for validation, by default None + X : List[Any], optional + The input data, by default None + y : List[int], optional + The target data, by default None + print_prefix : str, optional + The prefix used for printing, by default "" + + Returns + ------- + float + The accuracy of the model. + """ recorder = self.recorder recorder.print("Start val ", print_prefix) @@ -227,10 +472,54 @@ class BasicModel: ) return accuracy - def score(self, data_loader=None, X=None, y=None, print_prefix=""): + def score( + self, + data_loader: DataLoader = None, + X: List[Any] = None, + y: List[int] = None, + print_prefix: str = "", + ) -> float: + """ + Score the model. + + Parameters + ---------- + data_loader : DataLoader, optional + The data loader used for scoring, by default None + X : List[Any], optional + The input data, by default None + y : List[int], optional + The target data, by default None + print_prefix : str, optional + The prefix used for printing, by default "" + + Returns + ------- + float + The accuracy of the model. + """ return self.val(data_loader, X, y, print_prefix) - def _data_loader(self, X, y=None): + def _data_loader( + self, + X: List[Any], + y: List[int] = None, + ) -> DataLoader: + """ + Generate data_loader for user provided data. + + Parameters + ---------- + X : List[Any] + The input data. + y : List[int], optional + The target data, by default None + + Returns + ------- + DataLoader + The data loader. + """ collate_fn = self.collate_fn transform = self.transform @@ -238,7 +527,7 @@ class BasicModel: y = [0] * len(X) dataset = XYDataset(X, y, transform=transform) sampler = None - data_loader = torch.utils.data.DataLoader( + data_loader = DataLoader( dataset, batch_size=self.batch_size, shuffle=False, @@ -248,7 +537,17 @@ class BasicModel: ) return data_loader - def save(self, epoch_id, save_dir): + def save(self, epoch_id: int, save_dir: str = ""): + """ + Save the model and the optimizer. + + Parameters + ---------- + epoch_id : int + The epoch id. + save_dir : str, optional + The directory to save the model, by default "" + """ recorder = self.recorder if not os.path.exists(save_dir): os.makedirs(save_dir) @@ -259,7 +558,17 @@ class BasicModel: save_path = os.path.join(save_dir, str(epoch_id) + "_opt.pth") torch.save(self.optimizer.state_dict(), save_path) - def load(self, epoch_id, load_dir): + def load(self, epoch_id: int, load_dir: str = ""): + """ + Load the model and the optimizer. + + Parameters + ---------- + epoch_id : int + The epoch id. + load_dir : str, optional + The directory to load the model, by default "" + """ recorder = self.recorder recorder.print("Loading model and opter") load_path = os.path.join(load_dir, str(epoch_id) + "_net.pth")