diff --git a/learnware/market/heterogeneous/organizer/__init__.py b/learnware/market/heterogeneous/organizer/__init__.py index fdc73fe..51cd12a 100644 --- a/learnware/market/heterogeneous/organizer/__init__.py +++ b/learnware/market/heterogeneous/organizer/__init__.py @@ -15,7 +15,7 @@ import torch.multiprocessing as mp from ....learnware import Learnware, get_learnware_from_dirpath from ....logger import get_module_logger -from ....specification.system import HeteroSpecification +from ....specification.system import HeteroMapTableSpecification from ...base import BaseChecker, BaseUserInfo from ...easy import EasyOrganizer from ...easy.database_ops import DatabaseOperations @@ -71,15 +71,15 @@ class HeteroMapTableOrganizer(EasyOrganizer): if not rebuild: if os.path.exists(self.hetero_mappings_path): for hetero_json_path in os.listdir(self.hetero_mappings_path): - idx = hetero_json_path.split('.')[0] - hetero_spec = HeteroSpecification() + idx = hetero_json_path.split(".")[0] + hetero_spec = HeteroMapTableSpecification() hetero_spec.load(os.path.join(self.hetero_mappings_path, f"{idx}.json")) try: - self.learnware_list[idx].update_stat_spec("HeteroSpecification", hetero_spec) + self.learnware_list[idx].update_stat_spec("HeteroMapTableSpecification", hetero_spec) except: logger.warning(f"Learnware ID {idx} NOT Found!") else: - logger.info("No HeteroSpecifications to reload. Use loaded market mapping to regenerate.") + logger.info("No HeteroMapTableSpecification to reload. Use loaded market mapping to regenerate.") self._update_learnware_by_ids(self.learnware_list.keys()) else: logger.warning(f"No market mapping to reload!!") @@ -90,7 +90,8 @@ class HeteroMapTableOrganizer(EasyOrganizer): self.auto_update = auto_update self.market_id = market_id self.training_args = kwargs - if auto_update_limit is not None: self.auto_update_limit = auto_update_limit + if auto_update_limit is not None: + self.auto_update_limit = auto_update_limit def add_learnware( self, zip_path: str, semantic_spec: dict, check_status: int, learnware_id: str = None @@ -98,7 +99,7 @@ class HeteroMapTableOrganizer(EasyOrganizer): if check_status == BaseChecker.INVALID_LEARNWARE: logger.warning("Learnware is invalid!") return None, BaseChecker.INVALID_LEARNWARE - + semantic_spec = copy.deepcopy(semantic_spec) logger.info("Get new learnware from %s" % (zip_path)) @@ -123,7 +124,7 @@ class HeteroMapTableOrganizer(EasyOrganizer): except: pass return None, BaseChecker.INVALID_LEARNWARE - + if new_learnware is None: return None, BaseChecker.INVALID_LEARNWARE @@ -143,7 +144,7 @@ class HeteroMapTableOrganizer(EasyOrganizer): self.use_flags[learnware_id] = learnwere_status self._update_learnware_by_ids([learnware_id]) self.count += 1 - self.training_count += ([learnware_id] == self._get_table_type_learnware_ids([learnware_id])) + self.training_count += [learnware_id] == self._get_table_type_learnware_ids([learnware_id]) if self.auto_update and self.training_count - self.last_training_count == self.auto_update_limit + 1: training_learnware_ids = self._get_table_type_learnware_ids(self.get_learnware_ids()) @@ -151,16 +152,16 @@ class HeteroMapTableOrganizer(EasyOrganizer): logger.warning(f"Leanwares for training: {training_learnware_ids}") updated_market_mapping = self.train( - learnware_list=training_learnwares, - save_dir=self.market_store_path, - **self.training_args + learnware_list=training_learnwares, save_dir=self.market_store_path, **self.training_args + ) + + logger.warning( + f"Market mapping train completed. Now update HeteroMapTableSpecification for {training_learnware_ids}" ) - - logger.warning(f"Market mapping train completed. Now update HeteroSpecification for {training_learnware_ids}") self.market_mapping = updated_market_mapping self._update_learnware_by_ids(training_learnware_ids) self.last_training_count = len(training_learnware_ids) - + return learnware_id, learnwere_status @staticmethod @@ -178,7 +179,7 @@ class HeteroMapTableOrganizer(EasyOrganizer): market_mapping_trainer.save_model(output_dir=save_dir) return market_mapping - + def _update_learnware_by_ids(self, ids: List[str]): ids = self._get_table_type_learnware_ids(ids) for id in ids: @@ -187,14 +188,14 @@ class HeteroMapTableOrganizer(EasyOrganizer): semantic_spec, stat_spec = spec.get_semantic_spec(), spec.get_stat_spec()["RKMETableSpecification"] features = semantic_spec["Input"]["Description"].values() hetero_spec = self.market_mapping.hetero_mapping(stat_spec, features) - self.learnware_list[id].update_stat_spec("HeteroSpecification", hetero_spec) - + self.learnware_list[id].update_stat_spec("HeteroMapTableSpecification", hetero_spec) + save_path = os.path.join(self.hetero_mappings_path, f"{id}.json") hetero_spec.save(save_path) except Exception as err: - logger.warning(f"Learnware {id} generate HeteroSpecification failed! Due to {err}") - - def generate_hetero_map_spec(self, user_info: BaseUserInfo) -> HeteroSpecification: + logger.warning(f"Learnware {id} generate HeteroMapTableSpecification failed! Due to {err}") + + def generate_hetero_map_spec(self, user_info: BaseUserInfo) -> HeteroMapTableSpecification: user_stat_spec = user_info.stat_info["RKMETableSpecification"] user_features = user_info.get_semantic_spec()["Input"]["Description"].values() @@ -210,13 +211,13 @@ class HeteroMapTableOrganizer(EasyOrganizer): features = spec.get_semantic_spec()["Input"]["Description"] learnware_df = pd.DataFrame(data=stat_spec.get_z(), columns=features.values()) learnware_df_dict[tuple(sorted(features))].append(learnware_df) - + return [pd.concat(dfs) for dfs in learnware_df_dict.values()] - + def _get_table_type_learnware_ids(self, ids: List[str]) -> List[str]: ret = [] for id in ids: semantic_spec = self.learnware_list[id].get_specification().get_semantic_spec() if semantic_spec["Data"]["Values"][0] == "Table": ret.append(id) - return ret \ No newline at end of file + return ret diff --git a/learnware/market/heterogeneous/organizer/hetero_mapping/__init__.py b/learnware/market/heterogeneous/organizer/hetero_mapping/__init__.py index 69b7b99..231dfff 100644 --- a/learnware/market/heterogeneous/organizer/hetero_mapping/__init__.py +++ b/learnware/market/heterogeneous/organizer/hetero_mapping/__init__.py @@ -7,7 +7,7 @@ import torch import torch.nn.functional as F from torch import Tensor, nn -from .....specification import HeteroSpecification, RKMETableSpecification +from .....specification import HeteroMapTableSpecification, RKMETableSpecification from .feature_extractor import * from .trainer import Trainer, TransTabCollatorForCL @@ -147,8 +147,8 @@ class HeteroMapping(nn.Module): loss = self._self_supervised_contrastive_loss(feat_x_multiview) return loss - def hetero_mapping(self, rkme_spec: RKMETableSpecification, cols: List[str]) -> HeteroSpecification: - hetero_spec = HeteroSpecification() + def hetero_mapping(self, rkme_spec: RKMETableSpecification, cols: List[str]) -> HeteroMapTableSpecification: + hetero_spec = HeteroMapTableSpecification() hetero_input_df = pd.DataFrame(data=rkme_spec.get_z(), columns=cols) hetero_embedding = self._extract_batch_features(hetero_input_df) hetero_spec.generate_stat_spec_from_system(hetero_embedding, rkme_spec) diff --git a/learnware/market/heterogeneous/organizer/hetero_mapping/feature_extractor.py b/learnware/market/heterogeneous/organizer/hetero_mapping/feature_extractor.py index 1c23587..66eea79 100644 --- a/learnware/market/heterogeneous/organizer/hetero_mapping/feature_extractor.py +++ b/learnware/market/heterogeneous/organizer/hetero_mapping/feature_extractor.py @@ -6,7 +6,6 @@ from typing import Dict import numpy as np import torch import torch.nn.init as nn_init -from loguru import logger from torch import Tensor, nn from transformers import BertTokenizerFast diff --git a/learnware/market/heterogeneous/organizer/hetero_mapping/trainer.py b/learnware/market/heterogeneous/organizer/hetero_mapping/trainer.py index 3667f59..e9083a0 100644 --- a/learnware/market/heterogeneous/organizer/hetero_mapping/trainer.py +++ b/learnware/market/heterogeneous/organizer/hetero_mapping/trainer.py @@ -6,12 +6,14 @@ import time import numpy as np import pandas as pd import torch -from loguru import logger from torch import nn from torch.utils.data import DataLoader, Dataset from tqdm.autonotebook import trange from .feature_extractor import FeatureTokenizer +from .....logger import get_module_logger + +logger = get_module_logger("hetero_mapping_trainer") class Trainer: diff --git a/learnware/market/heterogeneous/searcher.py b/learnware/market/heterogeneous/searcher.py index 9782cf7..c30ad1e 100644 --- a/learnware/market/heterogeneous/searcher.py +++ b/learnware/market/heterogeneous/searcher.py @@ -4,7 +4,7 @@ import numpy as np from ...learnware import Learnware from ...logger import get_module_logger -from ...specification import HeteroSpecification +from ...specification import HeteroMapTableSpecification from ..base import BaseSearcher, BaseUserInfo from ..easy import EasySearcher from ..utils import parse_specification_type @@ -34,28 +34,28 @@ class HeteroMapTableSearcher(EasySearcher): return [(max_dist - dist) / (max_dist - dist_epsilon) for dist in dist_list] def _search_by_hetero_spec_single( - self, - learnware_list: List[Learnware], - user_hetero_spec: HeteroSpecification + self, learnware_list: List[Learnware], user_hetero_spec: HeteroMapTableSpecification ) -> Tuple[List[float], List[Learnware]]: - hetero_spec_list = [learnware.specification.get_stat_spec_by_name("HeteroSpecification") for learnware in learnware_list] + hetero_spec_list = [ + learnware.specification.get_stat_spec_by_name("HeteroMapTableSpecification") for learnware in learnware_list + ] mmd_dist_list = [] for idx, hetero_spec in enumerate(hetero_spec_list): mmd_dist = hetero_spec.dist(user_hetero_spec) mmd_dist_list.append(mmd_dist) - + sorted_idx_list = sorted(range(len(learnware_list)), key=lambda k: mmd_dist_list[k]) sorted_dist_list = [mmd_dist_list[idx] for idx in sorted_idx_list] sorted_learnware_list = [learnware_list[idx] for idx in sorted_idx_list] return sorted_dist_list, sorted_learnware_list - + def _filter_by_hetero_spec_single( self, sorted_score_list: List[float], learnware_list: List[Learnware], filter_score: float = 0.5, - min_num: int = 5 + min_num: int = 5, ) -> Tuple[List[float], List[Learnware]]: idx = min(min_num, len(learnware_list)) while idx < len(learnware_list): @@ -64,11 +64,10 @@ class HeteroMapTableSearcher(EasySearcher): idx += 1 return sorted_score_list[:idx], learnware_list[:idx] - def __call__( - self, - learnware_list: List[Learnware], - user_info: BaseUserInfo, + self, + learnware_list: List[Learnware], + user_info: BaseUserInfo, ) -> Tuple[List[float], List[Learnware], float, List[Learnware]]: # todo: use specially assigned search_gamma for calculating mmd dist user_hetero_spec = self.learnware_oganizer.generate_hetero_map_spec(user_info) @@ -88,6 +87,7 @@ class HeteroMapTableSearcher(EasySearcher): def reset(self, organizer): self.learnware_oganizer = organizer + class HeteroSearcher(EasySearcher): def __init__(self, organizer: HeteroMapTableOrganizer = None): super(HeteroSearcher, self).__init__(organizer) @@ -96,7 +96,7 @@ class HeteroSearcher(EasySearcher): def reset(self, organizer): super().reset(organizer) self.hetero_stat_searcher.reset(organizer) - + @staticmethod def check_user_info(user_info: BaseUserInfo): try: @@ -105,7 +105,9 @@ class HeteroSearcher(EasySearcher): user_task_type = user_info.get_semantic_spec()["Task"]["Values"] if user_task_type not in [["Classification"], ["Regression"]]: - logger.warning("User doesn't provide correct task type, it must be either Classification or Regression.") + logger.warning( + "User doesn't provide correct task type, it must be either Classification or Regression." + ) return False user_input_description = user_info.get_semantic_spec()["Input"] @@ -115,10 +117,12 @@ class HeteroSearcher(EasySearcher): if user_input_shape != user_description_dim or user_input_shape != user_description_feature_num: logger.warning("User data feature dimensions mismatch with semantic specification.") return False - + return True except Exception as e: - logger.info(f"Invalid heterogeneous search information provided. Use homogeneous search instead. Error: {e}") + logger.info( + f"Invalid heterogeneous search information provided. Use homogeneous search instead. Error: {e}" + ) return False def __call__( @@ -136,4 +140,4 @@ class HeteroSearcher(EasySearcher): else: return self.stat_searcher(learnware_list, user_info, max_search_num, search_method) else: - return None, learnware_list, 0.0, None \ No newline at end of file + return None, learnware_list, 0.0, None diff --git a/learnware/reuse/hetero_reuser/feature_alignment.py b/learnware/reuse/hetero_reuser/feature_alignment.py index def7764..7f81187 100644 --- a/learnware/reuse/hetero_reuser/feature_alignment.py +++ b/learnware/reuse/hetero_reuser/feature_alignment.py @@ -6,13 +6,15 @@ import torch.nn.functional as F import torch import time from tqdm import trange -from loguru import logger from learnware.learnware import Learnware from learnware.specification import RKMETableSpecification from learnware.specification.regular.table.rkme import choose_device from ..base import BaseReuser +from ...logger import get_module_logger + +logger = get_module_logger("hetero_feature_alignment") class FeatureAligner(BaseReuser): @@ -66,7 +68,9 @@ class FeatureAligner(BaseReuser): 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) + 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() @@ -85,7 +89,9 @@ class FeatureAligner(BaseReuser): 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() + 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 @@ -120,7 +126,6 @@ class FeatureAligner(BaseReuser): return X - class FeatureAlignmentModel(nn.Module): """ FeatureAlignmentModel is a neural network module designed for feature alignment tasks. @@ -128,7 +133,15 @@ class FeatureAlignmentModel(nn.Module): and supports different activation functions. """ - def __init__(self, input_dim: int, output_dim: int, hidden_dims: list = [1024], activation: str = "relu", dropout_ratio: float = 0, use_bn: bool = 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. @@ -187,13 +200,13 @@ class FeatureAlignmentModel(nn.Module): """ if len(self.fc_list) > 0: for fc, drop in zip(self.fc_list, self.drop_list): - x = fc(x) # Apply fully connected layer + x = fc(x) # Apply fully connected layer x = self.activation(x) # Apply activation function - x = drop(x) # Apply dropout + x = drop(x) # Apply dropout return self.final_fc(x) # Return output from final fully connected layer - -class FeatureAlignmentTrainer(): + +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. @@ -248,7 +261,7 @@ class FeatureAlignmentTrainer(): dropout_ratio: float = 0, use_bn: bool = False, const: float = 1e1, - cuda_idx: int = 0 + cuda_idx: int = 0, ): """ Initialize the FeatureAlignmentTrainer with the specified parameters. @@ -266,7 +279,7 @@ class FeatureAlignmentTrainer(): } self.network_type = network_type self.optimizer_type = optimizer_type - self.const=const + self.const = const self.device = choose_device(cuda_idx=cuda_idx) if extra_labeled_data is not None and target_learnware is not None: self.train_with_labeled_data(extra_labeled_data[0], extra_labeled_data[1], target_learnware) @@ -294,7 +307,9 @@ class FeatureAlignmentTrainer(): 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: torch.Tensor, user_weight: torch.Tensor, target_X: torch.Tensor, target_weight: torch.Tensor) -> torch.Tensor: + 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. @@ -327,7 +342,9 @@ class FeatureAlignmentTrainer(): input_dim = self.user_rkme.get_z().shape[1] output_dim = self.target_rkme.get_z().shape[1] - user_model=FeatureAlignmentModel(input_dim, output_dim, args["hidden_dims"], args["activation"], args["dropout_ratio"], args["use_bn"]) + user_model = FeatureAlignmentModel( + input_dim, output_dim, args["hidden_dims"], args["activation"], args["dropout_ratio"], args["use_bn"] + ) # device = torch.device("cuda" if torch.cuda.is_available() else "cpu") user_model.to(self.device) @@ -355,4 +372,4 @@ class FeatureAlignmentTrainer(): ) self.model = user_model - logger.info("training complete, cost {:.1f} secs.".format(time.time() - start_time)) \ No newline at end of file + logger.info("training complete, cost {:.1f} secs.".format(time.time() - start_time)) diff --git a/learnware/specification/__init__.py b/learnware/specification/__init__.py index fae0c7c..0332890 100644 --- a/learnware/specification/__init__.py +++ b/learnware/specification/__init__.py @@ -1,13 +1,13 @@ from .base import Specification, BaseStatSpecification from .regular import ( - RegularStatsSpecification, + RegularStatSpecification, RKMEStatSpecification, RKMETableSpecification, RKMEImageSpecification, RKMETextSpecification, ) -from .system import HeteroSpecification +from .system import HeteroMapTableSpecification from ..utils import is_torch_avaliable diff --git a/learnware/specification/regular/__init__.py b/learnware/specification/regular/__init__.py index 9d46114..bd01aa9 100644 --- a/learnware/specification/regular/__init__.py +++ b/learnware/specification/regular/__init__.py @@ -1,4 +1,4 @@ -from .base import RegularStatsSpecification +from .base import RegularStatSpecification from ...utils import is_torch_avaliable from .text import RKMETextSpecification diff --git a/learnware/specification/regular/base.py b/learnware/specification/regular/base.py index 6916177..1960f0d 100644 --- a/learnware/specification/regular/base.py +++ b/learnware/specification/regular/base.py @@ -3,7 +3,7 @@ from __future__ import annotations from ..base import BaseStatSpecification -class RegularStatsSpecification(BaseStatSpecification): +class RegularStatSpecification(BaseStatSpecification): def generate_stat_spec(self, **kwargs): self.generate_stat_spec_from_data(**kwargs) diff --git a/learnware/specification/regular/image/rkme.py b/learnware/specification/regular/image/rkme.py index 4421f91..e935afa 100644 --- a/learnware/specification/regular/image/rkme.py +++ b/learnware/specification/regular/image/rkme.py @@ -17,11 +17,11 @@ from torchvision.transforms import Resize from tqdm import tqdm from . import cnn_gp -from ..base import RegularStatsSpecification +from ..base import RegularStatSpecification from ..table.rkme import solve_qp, choose_device, setup_seed -class RKMEImageSpecification(RegularStatsSpecification): +class RKMEImageSpecification(RegularStatSpecification): # INNER_PRODUCT_COUNT = 0 IMAGE_WIDTH = 32 diff --git a/learnware/specification/regular/table/rkme.py b/learnware/specification/regular/table/rkme.py index 17aedc1..d0f113b 100644 --- a/learnware/specification/regular/table/rkme.py +++ b/learnware/specification/regular/table/rkme.py @@ -20,7 +20,7 @@ try: except ImportError: _FAISS_INSTALLED = False -from ..base import RegularStatsSpecification +from ..base import RegularStatSpecification from ....logger import get_module_logger logger = get_module_logger("rkme") @@ -31,7 +31,7 @@ if not _FAISS_INSTALLED: ) -class RKMETableSpecification(RegularStatsSpecification): +class RKMETableSpecification(RegularStatSpecification): """Reduced Kernel Mean Embedding (RKME) Specification""" def __init__(self, gamma: float = 0.1, cuda_idx: int = -1): diff --git a/learnware/specification/system/__init__.py b/learnware/specification/system/__init__.py index 1a8b6ca..e57a155 100644 --- a/learnware/specification/system/__init__.py +++ b/learnware/specification/system/__init__.py @@ -1 +1 @@ -from .heter_table import HeteroSpecification +from .heter_table import HeteroMapTableSpecification diff --git a/learnware/specification/system/base.py b/learnware/specification/system/base.py index 12a7226..281f682 100644 --- a/learnware/specification/system/base.py +++ b/learnware/specification/system/base.py @@ -1,11 +1,7 @@ -from __future__ import annotations - -from loguru import logger - from ..base import BaseStatSpecification -class SystemStatsSpecification(BaseStatSpecification): +class SystemStatSpecification(BaseStatSpecification): def generate_stat_spec(self, **kwargs): self.generate_stat_spec_from_system(**kwargs) diff --git a/learnware/specification/system/heter_table.py b/learnware/specification/system/heter_table.py index a574daf..5329a06 100644 --- a/learnware/specification/system/heter_table.py +++ b/learnware/specification/system/heter_table.py @@ -10,10 +10,10 @@ import torch from ..regular import RKMETableSpecification from ..regular.table.rkme import choose_device, setup_seed, torch_rbf_kernel -from .base import SystemStatsSpecification +from .base import SystemStatSpecification -class HeteroSpecification(SystemStatsSpecification): +class HeteroMapTableSpecification(SystemStatSpecification): """Heterogeneous Embedding Specification""" def __init__(self, gamma: float = 0.1, cuda_idx: int = -1): @@ -26,7 +26,7 @@ class HeteroSpecification(SystemStatsSpecification): torch.cuda.empty_cache() self.device = choose_device(cuda_idx=cuda_idx) setup_seed(0) - super(HeteroSpecification, self).__init__(type=self.__class__.__name__) + super(HeteroMapTableSpecification, self).__init__(type=self.__class__.__name__) def get_z(self) -> np.ndarray: return self.z.detach().cpu().numpy() @@ -38,7 +38,7 @@ class HeteroSpecification(SystemStatsSpecification): self.beta = rkme_spec.beta.to(self.device) self.z = torch.from_numpy(heter_embedding).double().to(self.device) - def inner_prod(self, Embed2: HeteroSpecification) -> float: + def inner_prod(self, Embed2: HeteroMapTableSpecification) -> float: beta_1 = self.beta.reshape(1, -1).double().to(self.device) beta_2 = Embed2.beta.reshape(1, -1).double().to(self.device) Z1 = self.z.double().reshape(self.z.shape[0], -1).to(self.device) @@ -47,7 +47,7 @@ class HeteroSpecification(SystemStatsSpecification): return float(v) - def dist(self, Embed2: HeteroSpecification, omit_term1: bool = False) -> float: + def dist(self, Embed2: HeteroMapTableSpecification, omit_term1: bool = False) -> float: term1 = 0 if omit_term1 else self.inner_prod(self) term2 = self.inner_prod(Embed2) term3 = Embed2.inner_prod(Embed2)