From 8a7e43422aa90399a188edb8e504bd31f0dbb0ed Mon Sep 17 00:00:00 2001 From: Peng Tan Date: Fri, 10 Nov 2023 21:12:47 +0800 Subject: [PATCH] [DOC] add docs for hetero_reuser --- learnware/reuse/feature_augment_reuser.py | 8 +- learnware/reuse/hetero_reuser/__init__.py | 80 +++++- .../reuse/hetero_reuser/feature_alignment.py | 227 ++++++++++++++++-- 3 files changed, 278 insertions(+), 37 deletions(-) diff --git a/learnware/reuse/feature_augment_reuser.py b/learnware/reuse/feature_augment_reuser.py index af98a0a..eaf6c43 100644 --- a/learnware/reuse/feature_augment_reuser.py +++ b/learnware/reuse/feature_augment_reuser.py @@ -15,7 +15,7 @@ class FeatureAugmentReuser(BaseReuser): def __init__(self, learnware: Learnware = None, mode: str = None): """ - Initializes the FeatureAugmentReuser with a learnware model and a mode. + Initialize the FeatureAugmentReuser with a learnware model and a mode. Parameters ---------- @@ -30,7 +30,7 @@ class FeatureAugmentReuser(BaseReuser): def predict(self, user_data: np.ndarray) -> np.ndarray: """ - Predicts the output for user data using the trained output aligner model. + Predict the output for user data using the trained output aligner model. Parameters ---------- @@ -50,7 +50,7 @@ class FeatureAugmentReuser(BaseReuser): def fit(self, x_train: np.ndarray, y_train: np.ndarray): """ - Trains the output aligner model using the training data augmented with predictions from the learnware model. + Train the output aligner model using the training data augmented with predictions from the learnware model. Parameters ---------- @@ -73,7 +73,7 @@ class FeatureAugmentReuser(BaseReuser): def _fill_data(self, X: np.ndarray): """ - Fills missing data (NaN, Inf) in the input array with the mean of the column. + Fill missing data (NaN, Inf) in the input array with the mean of the column. Parameters ---------- diff --git a/learnware/reuse/hetero_reuser/__init__.py b/learnware/reuse/hetero_reuser/__init__.py index 69d24a1..91edd7a 100644 --- a/learnware/reuse/hetero_reuser/__init__.py +++ b/learnware/reuse/hetero_reuser/__init__.py @@ -5,22 +5,82 @@ from ..feature_augment_reuser import FeatureAugmentReuser class HeteroMapTableReuser(BaseReuser): + """ + HeteroMapTableReuser is a class designed for reusing learnware models with feature alignment and augmentation. + It can handle both classification and regression tasks and supports fine-tuning on additional training data. + + Attributes + ---------- + learnware : Learnware + The learnware model to be reused. + mode : str + The mode of operation, either "classification" or "regression". + cuda_idx : int + Index of the CUDA device to be used for computations. + align_arguments : dict + Additional arguments for feature alignment. + """ def __init__(self, learnware: Learnware = None, mode: str = None, cuda_idx=0, **align_arguments): - self.learnware=learnware - assert mode in ["classification", "regression"] - self.mode=mode - self.cuda_idx=cuda_idx - self.align_arguments=align_arguments + """ + Initialize the HeteroMapTableReuser with a learnware model, mode, CUDA device index, and alignment arguments. + + Parameters + ---------- + learnware : Learnware + A learnware model used for initial predictions. + mode : str + The mode of operation, either "regression" or "classification". + cuda_idx : int + The index of the CUDA device for computations. + align_arguments : dict + Additional arguments to be passed to the feature alignment process. + """ + self.learnware = learnware + assert mode in ["classification", "regression"], "Mode must be either 'classification' or 'regression'" + self.mode = mode + self.cuda_idx = cuda_idx + self.align_arguments = align_arguments def fit(self, user_rkme): - self.feature_aligner=FeatureAligner(learnware=self.learnware, mode=self.mode, cuda_idx=self.cuda_idx, **self.align_arguments) + """ + Fit the feature aligner using the user RKME (Relative Knowledge Model Embeddings) specification. + + Parameters + ---------- + user_rkme : RKMETableSpecification + The RKME specification from the user dataset. + """ + self.feature_aligner = FeatureAligner(learnware=self.learnware, mode=self.mode, cuda_idx=self.cuda_idx, **self.align_arguments) self.feature_aligner.fit(user_rkme) - self.reuser=self.feature_aligner + self.reuser = self.feature_aligner - def finetune(self, x_train,y_train): - self.reuser=FeatureAugmentReuser(learnware=self.feature_aligner, mode=self.mode) + def finetune(self, x_train, y_train): + """ + Fine-tune the feature aligner using additional training data. + + Parameters + ---------- + x_train : ndarray + Training data features. + y_train : ndarray + Training data labels. + """ + self.reuser = FeatureAugmentReuser(learnware=self.feature_aligner, mode=self.mode) self.reuser.fit(x_train, y_train) def predict(self, user_data): - return self.reuser.predict(user_data) \ No newline at end of file + """ + Predict the output for user data using the feature aligner or the fine-tuned model. + + Parameters + ---------- + user_data : ndarray + Input data for making predictions. + + Returns + ------- + ndarray + Predicted output from the model. + """ + return self.reuser.predict(user_data) diff --git a/learnware/reuse/hetero_reuser/feature_alignment.py b/learnware/reuse/hetero_reuser/feature_alignment.py index 2bb3f71..def7764 100644 --- a/learnware/reuse/hetero_reuser/feature_alignment.py +++ b/learnware/reuse/hetero_reuser/feature_alignment.py @@ -16,28 +16,98 @@ from ..base import BaseReuser class FeatureAligner(BaseReuser): + """ + FeatureAligner is a class for aligning features from a user dataset with a target dataset using a learnware model. + It supports both classification and regression tasks and uses a feature alignment trainer for alignment. + + Attributes + ---------- + learnware : Learnware + The learnware model used for final prediction. + mode : str + Operation mode, either "classification" or "regression". + align_arguments : dict + Additional arguments for the feature alignment trainer. + cuda_idx : int + Index of the CUDA device to be used for computations. + device : torch.device + The device (CPU or CUDA) on which computations will be performed. + """ def __init__(self, learnware: Learnware = None, mode: str = None, cuda_idx=0, **align_arguments): - self.learnware=learnware - assert mode in ["classification", "regression"] - self.mode=mode - self.align_arguments=align_arguments - self.cuda_idx=cuda_idx + """ + Initialize the FeatureAligner with a learnware model, mode, CUDA device index, and alignment arguments. + + Parameters + ---------- + learnware : Learnware + A learnware model used for initial predictions. + mode : str + The mode of operation, either "regression" or "classification". + cuda_idx : int + The index of the CUDA device for computations. + align_arguments : dict + Additional arguments to be passed to the feature alignment trainer. + """ + self.learnware = learnware + assert mode in ["classification", "regression"], "Mode must be either 'classification' or 'regression'" + self.mode = mode + self.align_arguments = align_arguments + self.cuda_idx = cuda_idx self.device = choose_device(cuda_idx=cuda_idx) - def fit(self, user_rkme): - target_rkme=self.learnware.specification.get_stat_spec()["RKMETableSpecification"] - trainer=FeatureAlignmentTrainer(target_rkme=target_rkme, user_rkme=user_rkme, cuda_idx=self.cuda_idx, **self.align_arguments) - self.align_model=trainer.model + def fit(self, user_rkme: RKMETableSpecification): + """ + Fit the align model using the RKME (Relative Knowledge Model Embeddings) specifications from the learnware model. + + Parameters + ---------- + user_rkme : RKMETableSpecification + The RKME specification from the user dataset. + """ + target_rkme = self.learnware.specification.get_stat_spec()["RKMETableSpecification"] + trainer = FeatureAlignmentTrainer(target_rkme=target_rkme, user_rkme=user_rkme, cuda_idx=self.cuda_idx, **self.align_arguments) + self.align_model = trainer.model self.align_model.eval() def predict(self, user_data: ndarray) -> ndarray: - user_data=self._fill_data(user_data) - transformed_user_data=self.align_model(torch.tensor(user_data, device=self.device).float()).detach().cpu().numpy() - y_pred=self.learnware.predict(transformed_user_data) + """ + Predict the output for user data using the aligned model and learnware model. + + Parameters + ---------- + user_data : ndarray + Input data for making predictions. + + Returns + ------- + ndarray + Predicted output from the learnware model after alignment. + """ + user_data = self._fill_data(user_data) + transformed_user_data = self.align_model(torch.tensor(user_data, device=self.device).float()).detach().cpu().numpy() + y_pred = self.learnware.predict(transformed_user_data) return y_pred - + def _fill_data(self, X: np.ndarray): + """ + Fill missing data (NaN, Inf) in the input array with the mean of the column. + + Parameters + ---------- + X : np.ndarray + Input data array that may contain missing values. + + Returns + ------- + np.ndarray + Data array with missing values filled. + + Raises + ------ + ValueError + If a column in X contains only exceptional values (NaN, Inf). + """ X[np.isinf(X) | np.isneginf(X) | np.isposinf(X) | np.isneginf(X)] = np.nan if np.any(np.isnan(X)): for col in range(X.shape[1]): @@ -45,15 +115,38 @@ class FeatureAligner(BaseReuser): if np.any(is_nan): if np.all(is_nan): raise ValueError(f"All values in column {col} are exceptional, e.g., NaN and Inf.") - # Fill np.nan with np.nanmean col_mean = np.nanmean(X[:, col]) X[:, col] = np.where(is_nan, col_mean, X[:, col]) return X + class FeatureAlignmentModel(nn.Module): + """ + FeatureAlignmentModel is a neural network module designed for feature alignment tasks. + It consists of multiple fully connected (dense) layers, optional dropout and batch normalization layers, + and supports different activation functions. + """ - def __init__(self, input_dim, output_dim, hidden_dims=[1024], activation="relu", dropout_ratio=0, use_bn=False): + def __init__(self, input_dim: int, output_dim: int, hidden_dims: list = [1024], activation: str = "relu", dropout_ratio: float = 0, use_bn: bool = False): + """ + Initialize the FeatureAlignmentModel. + + Parameters + ---------- + input_dim : int + The dimensionality of the input features. + output_dim : int + The dimensionality of the output features. + hidden_dims : List[int], optional + A list specifying the number of units in each hidden layer. + activation : str, optional + The activation function to use. Supported options are "relu", "gelu", "selu", and "leakyrelu". + dropout_ratio : float, optional + The dropout ratio applied to each layer (0 means no dropout). + use_bn : bool, optional + Whether to use batch normalization after each fully connected layer. + """ super().__init__() dims = [input_dim] + hidden_dims + [output_dim] self.fc_list = nn.ModuleList() @@ -78,16 +171,66 @@ class FeatureAlignmentModel(nn.Module): else: self.activation = F.relu - def forward(self, x): + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass of the model. + + Parameters + ---------- + x : torch.Tensor + Input tensor. + + Returns + ------- + torch.Tensor + Output tensor after passing through the model. + """ if len(self.fc_list) > 0: for fc, drop in zip(self.fc_list, self.drop_list): - x = fc(x) - x = self.activation(x) - x = drop(x) - return self.final_fc(x) + x = fc(x) # Apply fully connected layer + x = self.activation(x) # Apply activation function + x = drop(x) # Apply dropout + return self.final_fc(x) # Return output from final fully connected layer class FeatureAlignmentTrainer(): + """ + FeatureAlignmentTrainer is a class designed to train a neural network for aligning features from a user dataset + to a target dataset. It utilizes Maximum Mean Discrepancy (MMD) as the loss function for training. + + Attributes + ---------- + target_rkme : RKMETableSpecification + The RKME (Relative Knowledge Model Embeddings) specification of the target dataset. + user_rkme : RKMETableSpecification + The RKME specification of the user dataset. + extra_labeled_data : Any, optional + Additional labeled data for training, if available. + target_learnware : Learnware, optional + The learnware model used for the target dataset. + num_epoch : int + The number of training epochs. + lr : float + Learning rate for the optimizer. + gamma : float + The gamma parameter for the Gaussian kernel in MMD computation. + network_type : str + Type of the neural network used for feature alignment. + optimizer_type : str + Type of optimizer to be used in training ('Adam' or 'SGD'). + hidden_dims : List[int] + A list specifying the number of units in each hidden layer. + activation : str + The activation function to use in the network. + dropout_ratio : float + The dropout ratio applied to each layer. + use_bn : bool + Whether to use batch normalization after each fully connected layer. + const : float + A constant value used in training. + cuda_idx : int + Index of the CUDA device to be used for computations. + """ def __init__( self, @@ -107,7 +250,8 @@ class FeatureAlignmentTrainer(): const: float = 1e1, cuda_idx: int = 0 ): - """Training the base mapping network + """ + Initialize the FeatureAlignmentTrainer with the specified parameters. """ self.target_rkme = target_rkme self.user_rkme = user_rkme @@ -129,19 +273,56 @@ class FeatureAlignmentTrainer(): else: self.train() - def gaussian_kernel(self, x1, x2): + def gaussian_kernel(self, x1: torch.Tensor, x2: torch.Tensor): + """ + Compute the Gaussian kernel between two sets of samples. + + Parameters + ---------- + x1 : torch.Tensor + First set of samples. + x2 : torch.Tensor + Second set of samples. + + Returns + ------- + torch.Tensor + The computed Gaussian kernel matrix. + """ x1 = x1.double() x2 = x2.double() X12norm = torch.sum(x1**2, 1, keepdim=True) - 2 * x1 @ x2.T + torch.sum(x2**2, 1, keepdim=True).T return torch.exp(-X12norm * self.args["gamma"]) - def compute_mmd(self, user_X, user_weight, target_X, target_weight): + def compute_mmd(self, user_X: torch.Tensor, user_weight: torch.Tensor, target_X: torch.Tensor, target_weight: torch.Tensor) -> torch.Tensor: + """ + Compute the Maximum Mean Discrepancy (MMD) between the user and target datasets. + + Parameters + ---------- + user_X : torch.Tensor + Transformed user data. + user_weight : torch.Tensor + Weights of the user data. + target_X : torch.Tensor + Target data. + target_weight : torch.Tensor + Weights of the target data. + + Returns + ------- + torch.Tensor + The computed MMD loss. + """ term1 = torch.sum(self.gaussian_kernel(user_X, user_X) * (user_weight.T @ user_weight)) term2 = torch.sum(self.gaussian_kernel(user_X, target_X) * (user_weight.T @ target_weight)) term3 = torch.sum(self.gaussian_kernel(target_X, target_X) * (target_weight.T @ target_weight)) return term1 - 2 * term2 + term3 def train(self): + """ + Train the feature alignment model using MMD as the loss function. + """ args = self.args input_dim = self.user_rkme.get_z().shape[1] output_dim = self.target_rkme.get_z().shape[1]