pax_global_header00006660000000000000000000000064142774154230014523gustar00rootroot0000000000000052 comment=001ccd80def887f555fbf4a860e375993d7ff908 fedlab-nlp/000077500000000000000000000000001427741542300130735ustar00rootroot00000000000000fedlab-nlp/.gitignore000066400000000000000000000027701427741542300150710ustar00rootroot00000000000000# ---> JetBrains # Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio and WebStorm # Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839 # User-specific stuff .idea/**/workspace.xml .idea/**/tasks.xml .idea/**/usage.statistics.xml .idea/**/dictionaries .idea/**/shelf # Generated files .idea/**/contentModel.xml # Sensitive or high-churn files .idea/**/dataSources/ .idea/**/dataSources.ids .idea/**/dataSources.local.xml .idea/**/sqlDataSources.xml .idea/**/dynamic.xml .idea/**/uiDesigner.xml .idea/**/dbnavigator.xml # Gradle .idea/**/gradle.xml .idea/**/libraries # Gradle and Maven with auto-import # When using Gradle or Maven with auto-import, you should exclude module files, # since they will be recreated, and may cause churn. Uncomment if using # auto-import. # .idea/artifacts # .idea/compiler.xml # .idea/jarRepositories.xml # .idea/modules.xml # .idea/*.iml # .idea/modules # *.iml # *.ipr # CMake cmake-build-*/ # Mongo Explorer plugin .idea/**/mongoSettings.xml # File-based project format *.iws # IntelliJ out/ # mpeltonen/sbt-idea plugin .idea_modules/ # JIRA plugin atlassian-ide-plugin.xml # Cursive Clojure plugin .idea/replstate.xml # Crashlytics plugin (for Android Studio and IntelliJ) com_crashlytics_export_strings.xml crashlytics.properties crashlytics-build.properties fabric.properties # Editor-based Rest Client .idea/httpRequests # Android studio 3.1+ serialized cache file .idea/caches/build_file_checksums.ser fedlab-nlp/LICENSE000066400000000000000000000241471427741542300141100ustar00rootroot00000000000000Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document. "Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License. "Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity. "You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License. "Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files. "Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types. "Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below). "Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof. "Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution." "Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work. 2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form. 3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed. 4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions: (a) You must give any other recipients of the Work or Derivative Works a copy of this License; and (b) You must cause any modified files to carry prominent notices stating that You changed the files; and (c) You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and (d) If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License. You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License. 5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions. 6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file. 7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License. 8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages. 9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability. END OF TERMS AND CONDITIONS APPENDIX: How to apply the Apache License to your work. To apply the Apache License to your work, attach the following boilerplate notice, with the fields enclosed by brackets "[]" replaced with your own identifying information. (Don't include the brackets!) The text should be enclosed in the appropriate comment syntax for the file format. We also recommend that a file or class name and description of purpose be included on the same "printed page" as the copyright notice for easier identification within third-party archives. Copyright [yyyy] [name of copyright owner] Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. fedlab-nlp/README.md000066400000000000000000000046311427741542300143560ustar00rootroot00000000000000

FedLab-NLP

SMILELAB Licence Build

为联邦自然语言处理发展添瓦加砖

1️⃣ 本项目来自 [SMILELab-FL](https://github.com/SMILELab-FL/FedLab-NLP) , 我们主要致力于联邦自然语言处理的发展； 2️⃣ `FedLab-NLP`是一款轻量级、可自由定制联邦学习过程并面向于自然语言处理的代码框架； 3️⃣ `FedLab-NLP`主要实现基于 [FedLab](https://github.com/SMILELab-FL/FedLab) 和 [HuggingFace](https://github.com/huggingface/transformers)，能完成自然语言预训练大模型在联邦学习场景下工作； 4️⃣ 任何使用问题可以联系 :e-mail: iezhuo17@gmail.com ## Installation ### 目录组织 我们建议按照如下的文件结构来使用`FedLab-NLP`: data: 存放数据集，注意该数据集需要已经经过 iid/niid 划分的联邦数据集; pretrained: 在nlp文件夹里存储从huggingface中下载的预训练模型，如`bert-base-uncased`; code: 在该文件夹下面使用`FedLab-NLP`; output: 存储模型输出的记录等。 目录组织如下： ```grapha ├── workspace │ └── data | | ├── fedglue | | └── fedner │ ├── pretrained │ │ ├── nlp │ │ └── cv │ ├── output │ └── code │ └── FedLab-NLP ``` 运行路径生成： ```bash mkdir workspace cd workspace mkdir data mkdir code mkdir pretrained cd pretrained mkdir nlp cd .. cd code ``` ### 环境安装 建议运行环境的`python`版本为`3.7+`，我们建议使用`pytorch`版本为`1.10+` ```bash git clone git@git.openi.org.cn:Trustworthy-DL/FedLab-NLP.git cd FedLab-NLP pip install -r resquirements.txt ``` ## Usage 支持 `fedavg` 联邦学习算法 ```bash bash fed_run.sh {your_file_path}/workspace {task_name} fedavg 10001 {server_gpu_id} {client1_gpu_id} {client2_gpu_id} ``` 支持 `centralized` 集中训练学习算法 ```bash bash cen_run.sh {your_file_path}/workspace {task_name} centralized 10001 {server_gpu_id} ``` ## Contributing fedlab-nlp/cen_run.sh000066400000000000000000000015111427741542300150560ustar00rootroot00000000000000#!/bin/bash # shellcheck disable=SC2068 # 读取参数 idx=0 for i in $@ do args[${idx}]=$i let "idx=${idx}+1" done # 分离参数 run_dirs=${args[0]} task_name=${args[1]} fl_algorithm=${args[2]} port=${args[3]} idx=0 for(( i=4;i<${#args[@]};i++ )) do device[${idx}]=${args[i]} let "idx=${idx}+1" done world_size=${#device[@]} #echo "word_size is ${word_size}" if [ ${task_name} == "qnli" ]; then max_seq=256 else max_seq=128 fi echo "${task_name}'s max_seq is ${max_seq}" CUDA_VISIBLE_DEVICES=${device[0]} python main.py \ --model_name_or_path ${run_dirs}/pretrain/nlp/bert-base-uncased/ \ --output_dir ${run_dirs}/output/fedglue \ --task_name ${task_name} \ --fl_algorithm ${fl_algorithm} \ --raw_dataset_path ${run_dirs}/data/fedglue \ --partition_dataset_path ${run_dirs}/data/fedglue \ --max_seq_length ${max_seq} waitfedlab-nlp/configs/000077500000000000000000000000001427741542300145235ustar00rootroot00000000000000fedlab-nlp/configs/__init__.py000066400000000000000000000005501427741542300166340ustar00rootroot00000000000000from configs.models import ModelArguments from configs.trainers import TrainingArguments, TrainArguments from configs.datasets import DataTrainingArguments from configs.federated import FederatedTrainingArguments __all__ = [ "ModelArguments", "TrainingArguments", "TrainArguments", "DataTrainingArguments", "FederatedTrainingArguments" ] fedlab-nlp/configs/datasets.py000066400000000000000000000076461427741542300167220ustar00rootroot00000000000000import os from typing import Optional from dataclasses import dataclass, field @dataclass class DataTrainingArguments: task_name: Optional[str] = field( default=None, metadata={"help": "The name of the task"}, ) dataset_name: Optional[str] = field( default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."} ) raw_dataset_path: Optional[str] = field( default=None, metadata={"help": "The raw data dir. if None, it builds from globalhost.py"} ) partition_dataset_path: Optional[str] = field( default=None, metadata={"help": "The partition data dir. if None, it builds from globalhost.py"} ) dataset_config_name: Optional[str] = field( default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."} ) max_seq_length: int = field( default=128, metadata={ "help": ( "The maximum total input sequence length after tokenization. Sequences longer " "than this will be truncated, sequences shorter will be padded." ) }, ) cache_dir: Optional[str] = field( default=None, metadata={"help": "The cached data dir. if None, it builds from globalhost.py"} ) overwrite_cache: bool = field( default=False, metadata={"help": "Overwrite the cached preprocessed datasets or not."} ) pad_to_max_length: bool = field( default=True, metadata={ "help": ( "Whether to pad all samples to `max_seq_length`. " "If False, will pad the samples dynamically when batching to the maximum length in the batch." ) }, ) max_train_samples: Optional[int] = field( default=None, metadata={ "help": ( "For debugging purposes or quicker training, truncate the number of training examples to this " "value if set." ) }, ) max_eval_samples: Optional[int] = field( default=None, metadata={ "help": ( "For debugging purposes or quicker training, truncate the number of evaluation examples to this " "value if set." ) }, ) max_predict_samples: Optional[int] = field( default=None, metadata={ "help": ( "For debugging purposes or quicker training, truncate the number of prediction examples to this " "value if set." ) }, ) train_file: Optional[str] = field( default=None, metadata={"help": "A csv or a json file containing the training data."} ) validation_file: Optional[str] = field( default=None, metadata={"help": "A csv or a json file containing the validation data."} ) test_file: Optional[str] = field(default=None, metadata={"help": "A csv or a json file containing the test data."}) def __post_init__(self): if self.task_name is None: raise ValueError(f"The task_name must be set, but {self.task_name} found") else: self.task_name = self.task_name.lower() if self.raw_dataset_path is None: raise ValueError(f"The raw_dataset_path must be set, but {self.raw_dataset_path} found") else: if not self.raw_dataset_path.endswith(".pkl"): self.raw_dataset_path = os.path.join( self.raw_dataset_path, f"{self.task_name}_data.pkl" ) if self.partition_dataset_path is None: raise ValueError(f"The raw_dataset_path must be set, but {self.raw_dataset_path} found") else: if not self.partition_dataset_path.endswith(".pkl"): self.partition_dataset_path = os.path.join( self.partition_dataset_path, f"{self.task_name}_partition.pkl" ) fedlab-nlp/configs/federated.py000066400000000000000000000053371427741542300170300ustar00rootroot00000000000000"""federated configs for FedLab-NLP""" from typing import Optional from dataclasses import dataclass, field @dataclass class FederatedTrainingArguments: fl_algorithm: str = field( default="fedavg", metadata={"help": "The name of the federated learning algorithm"}, ) clients_num: int = field( default=10, metadata={"help": "The number of participant clients"}, ) alpha: Optional[float] = field( default=1.0, metadata={"help": "Non-IID shift and None denotes IID"}, ) partition_method: str = field( default=None, metadata={"help": "The partition methods"}, ) do_mimic: bool = field( default=True, metadata={"help": "Important! we only process once data processing in server if True"} ) ip: str = field( default="127.0.0.1" ) port: str = field( default="10001", metadata={"help": "The torch.distribute find this port to transmit stuff."} ) rank: int = field( default=0, metadata={"help": "-1: centralized, 0: server, >0: client"} ) world_size: int = field( default=None, metadata={"help": "The number of sub-server"} ) ethernet: Optional[str] = field( default=None, metadata={"help": "not set"} ) rounds: int = field( default=100, metadata={"help": "The number of training round"} ) sample: float = field( default=1, metadata={"help": "The participant ratio in each training round"} ) pson: bool = field( default=False, metadata={"help": "Whether to use personalized test(local) metric"} ) _clients_num_per_sub_server: int = field( init=False, metadata={"help": "The number of clients in different works"} ) def __post_init__(self): if self.alpha is None: # IID self.alpha = "inf" if not self.do_mimic: print("Please check whether federated device has its own data") @property def clients_num_per_sub_server(self): return int(self.clients_num / (self.world_size-1)) @property def clients_id_list(self): if self.rank == -1: return [1] elif self.rank == 0: return [i for i in range(self.clients_num)] else: client_id_end = min(self.clients_num, self.rank * self.clients_num_per_sub_server) client_id_list = [ i for i in range((self.rank - 1) * self.clients_num_per_sub_server, client_id_end) ] return client_id_list @property def clients_num_in_total(self): if self.rank == -1: # centralized return 1 else: # federated return self.clients_num fedlab-nlp/configs/models.py000066400000000000000000000043451427741542300163660ustar00rootroot00000000000000from typing import Optional, List from dataclasses import dataclass, field @dataclass class ModelArguments: """ Arguments pertaining to which model/configs/tokenizer we are going to fine-tune from. """ model_name_or_path: str = field( metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"} ) model_type: Optional[str] = field( default=None, metadata={"help": "model_name"} ) model_output_mode: Optional[str] = field( default=None, metadata={"help": "model_output_type: {seq_classification, seq_tagging, language_model}"} ) config_name: Optional[str] = field( default=None, metadata={"help": "Pretrained configs name or path if not the same as model_name"} ) tokenizer_name: Optional[str] = field( default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"} ) use_fast_tokenizer: bool = field( default=True, metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."}, ) model_revision: str = field( default="main", metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."}, ) use_auth_token: bool = field( default=False, metadata={ "help": ( "Will use the token generated when running `transformers-cli login` (necessary to use this script " "with private models)." ) }, ) ignore_mismatched_sizes: bool = field( default=False, metadata={"help": "Will enable to load a pretrained model whose head dimensions are different."}, ) # Personalized Model Config permutation_layers: bool = field( default=False, metadata={"help": "The flag controls model layers."}, ) client_model_layers: List[int] = field( default_factory=list, metadata={"help": "The client model's size, default=[0,1,2]"}, ) server_model_layers: List[int] = field( default_factory=list, metadata={"help": "The client model's size, default=[i for i in range(12)]"}, ) def __post_init__(self): ... fedlab-nlp/configs/trainers.py000066400000000000000000000010641427741542300167250ustar00rootroot00000000000000from typing import Optional from dataclasses import dataclass, field from transformers import TrainingArguments @dataclass class TrainArguments(TrainingArguments): do_reuse: bool = field( default=False, metadata={"help": "whether to load last checkpoint"} ) metric_name: str = field( default="glue", metadata={"help": "whether to load last checkpoint"} ) loss_name: str = field( default="xent", metadata={"help": "{xent: cross_entropy}"} ) is_decreased_valid_metric: bool = field( default=False ) fedlab-nlp/data/000077500000000000000000000000001427741542300140045ustar00rootroot00000000000000fedlab-nlp/data/__init__.py000066400000000000000000000002251427741542300161140ustar00rootroot00000000000000from data.base_dataloader import BaseDataLoader from data.dataloader import GlueDataLoader __all__ = [ "BaseDataLoader", "GlueDataLoader" ] fedlab-nlp/data/base_dataloader.py000066400000000000000000000047251427741542300174600ustar00rootroot00000000000000"""Base DataLoader for FedLab-NLP""" import os from abc import ABC from utils import registry, pickle_read from transformers import AutoTokenizer class BaseDataLoader(ABC): def __init__(self): config = registry.get("config") self.model_config = config.model_config self.data_config = config.data_config self.training_config = config.training_config self.federated_config = config.federated_config self.partition_name = self.federated_config.partition_method self.clients_list = self.federated_config.clients_id_list self._load_attributes() self._build_tokenizer() self._build_registry() self.logger = registry.get("logger") def _load_data(self): if self.federated_config.rank == -1: self._load_centralized_data() elif self.federated_config.rank == 0: self._load_federated_data_on_server() else: self._load_federated_data_on_client() def _load_federated_data_on_client(self): raise NotImplementedError def _load_federated_data_on_server(self): raise NotImplementedError def _load_centralized_data(self): raise NotImplementedError def _convert_examples_to_features(self): raise NotImplementedError def _load_attributes(self): partition_data = pickle_read(self.data_config.partition_dataset_path) self.attribute = partition_data[self.partition_name]["attribute"] def _build_registry(self): if self.model_config.model_output_mode == "seq_classification": registry.register("num_labels", len(self.attribute["label_list"])) def _build_tokenizer(self): self.tokenizer = AutoTokenizer.from_pretrained(self.model_config.model_name_or_path) @property def cached_data_file(self): if self.federated_config.rank != -1: cached_file = os.path.join( self.data_config.cache_dir, f"models={self.model_config.model_type}_" f"seq={self.data_config.max_seq_length}_" f"clients={self.federated_config.clients_num}_" f"alpha={self.federated_config.alpha}" ) else: cached_file = os.path.join( self.data_config.cache_dir, f"models={self.model_config.model_type}_" f"seq={self.data_config.max_seq_length}_" f"centralized" ) return cached_file fedlab-nlp/data/dataloader.py000066400000000000000000000206171427741542300164640ustar00rootroot00000000000000"""Customer Dataloader for FedLab-NLP""" import os import numpy as np from data import BaseDataLoader from utils import registry, pickle_write, pickle_read, check_cached_data import torch from torch.utils.data import (DataLoader, RandomSampler, SequentialSampler, TensorDataset) from transformers import glue_convert_examples_to_features @registry.register_data("glue") class GlueDataLoader(BaseDataLoader): def __init__(self): super().__init__() self.output_mode = self.attribute["output_mode"] self.label_list = self.attribute["label_list"] self._load_data() def _load_federated_data_on_server(self): if os.path.isfile(self.cached_data_file): self.logger.info(f"loading cached data ...") train_features_dict, valid_features_dict, valid_fedtures_all, test_fedtures_all, \ train_examples_num_dict, valid_examples_num_dict, train_num, valid_num, test_num \ = pickle_read(self.cached_data_file) # server doesn't use each client's train & test dataset del train_features_dict, valid_features_dict else: self.logger.info(f"generating cached data ...") train_features_dict, valid_features_dict, valid_fedtures_all, test_fedtures_all, \ train_examples_num_dict, valid_examples_num_dict = self._convert_examples_to_features() if self.federated_config.do_mimic and self.federated_config.rank == 0: with open(os.path.join(self.data_config.cache_dir, "server_write.flag"), "w") as file: file.write("BaseServer wrote OK\n") self.valid_dataloader = self.build_dataloader(valid_fedtures_all, "valid") self.test_dataloader = self.build_dataloader(test_fedtures_all, "test") self.train_examples_num_dict = train_examples_num_dict self.valid_examples_num_dict = valid_examples_num_dict def _load_federated_data_on_client(self): train_dataloader_dict, valid_dataloader_dict = {}, {} if self.federated_config.do_mimic: self.logger.info(f"local rank {self.federated_config.rank} is waiting for processed features") while not check_cached_data(self.data_config.cache_dir): ... self.logger.info(f"local rank {self.federated_config.rank} builds dataloader") train_features_dict, valid_features_dict, valid_fedtures_all, test_fedtures_all, \ train_examples_num_dict, valid_examples_num_dict, train_num, valid_num, test_num \ = pickle_read(self.cached_data_file) del valid_fedtures_all, test_fedtures_all for idx in self.clients_list: train_dataloader_dict[idx] = self.build_dataloader(train_features_dict[idx], "train") valid_dataloader_dict[idx] = self.build_dataloader(valid_features_dict[idx], "valid") else: # Local data loading self.logger.info("Sorry, the current glue_dataloader doesn't support local loading") raise NotImplementedError self.train_dataloader_dict = train_dataloader_dict self.valid_dataloader_dict = valid_dataloader_dict self.train_examples_num_dict = train_examples_num_dict self.valid_examples_num_dict = valid_examples_num_dict self.train_num, self.valid_num, self.test_num = train_num, valid_num, test_num def _load_centralized_data(self): train_dataloader_dict, valid_dataloader_dict = {}, {} if os.path.isfile(self.cached_data_file): self.logger.info(f"loading cached data ...") train_features_dict, valid_features_dict, valid_fedtures_all, test_fedtures_all, \ train_examples_num_dict, valid_examples_num_dict, train_num, valid_num, test_num \ = pickle_read(self.cached_data_file) else: self.logger.info(f"generating cached data ...") train_features_dict, valid_features_dict, valid_fedtures_all, test_fedtures_all, \ train_examples_num_dict, valid_examples_num_dict = self._convert_examples_to_features() train_features_all = [] for idx, train_features in train_features_dict.items(): train_features_all += list(train_features) train_dataloader_dict[-1] = self.build_dataloader(train_features_all, "train") valid_dataloader_dict[-1] = self.build_dataloader(valid_fedtures_all, "valid") self.train_dataloader_dict = train_dataloader_dict self.valid_dataloader_dict = valid_dataloader_dict self.test_dataloader = self.build_dataloader(test_fedtures_all, "test") def build_dataloader(self, features, mode="train"): # Convert to Tensors and build dataset all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long) all_attention_mask = torch.tensor([f.attention_mask for f in features], dtype=torch.long) all_token_type_ids = torch.tensor([f.token_type_ids for f in features], dtype=torch.long) if self.output_mode == "classification": all_labels = torch.tensor([f.label for f in features], dtype=torch.long) elif self.output_mode == "regression": all_labels = torch.tensor([f.label for f in features], dtype=torch.float) dataset = TensorDataset(all_input_ids, all_attention_mask, all_token_type_ids, all_labels) sampler = RandomSampler(dataset) if mode == "train" else SequentialSampler(dataset) dataloader = DataLoader(dataset, sampler=sampler, batch_size=self.training_config.train_batch_size) return dataloader def _convert_examples_to_features(self): raw_data = pickle_read(self.data_config.raw_dataset_path) partition_data = pickle_read(self.data_config.partition_dataset_path) train_examples_num_dict, valid_examples_num_dict = {}, {} train_features_dict, valid_features_dict = {}, {} clients_partition_data = partition_data[self.partition_name] n_clients = self.attribute["clients_num"] if n_clients != self.federated_config.clients_num: raise ValueError(f"partition data have {n_clients} clients " f"that mismatches your input {self.federated_config.clients_num} clients") self.logger.info("convert train examples into features ...") train_features_all = np.array(glue_convert_examples_to_features( examples=raw_data["train"], tokenizer=self.tokenizer, max_length=self.data_config.max_seq_length, label_list=self.label_list, output_mode=self.output_mode )) self.logger.info("convert valid examples into features ...") valid_fedtures_all = np.array(glue_convert_examples_to_features( examples=raw_data["valid"], tokenizer=self.tokenizer, max_length=self.data_config.max_seq_length, label_list=self.label_list, output_mode=self.output_mode )) self.logger.info("convert test examples into features ...") test_fedtures_all = np.array(glue_convert_examples_to_features( examples=raw_data["test"], tokenizer=self.tokenizer, max_length=self.data_config.max_seq_length, label_list=self.label_list, output_mode=self.output_mode )) self.logger.info("build clients train & valid features ...") for idx in range(n_clients): client_train_list = clients_partition_data["train"][idx] train_examples_num_dict[idx] = len(client_train_list) train_features_dict[idx] = train_features_all[client_train_list] client_valid_list = clients_partition_data["valid"][idx] valid_examples_num_dict[idx] = len(client_valid_list) valid_features_dict[idx] = valid_fedtures_all[client_valid_list] self.train_num, self.valid_num, self.test_num = \ len(train_features_all), len(valid_fedtures_all), len(test_fedtures_all) federated_data = ( train_features_dict, valid_features_dict, valid_fedtures_all, test_fedtures_all, train_examples_num_dict, valid_examples_num_dict, self.train_num, self.valid_num, self.test_num ) self.logger.info("saving processed features ...") pickle_write(federated_data, self.cached_data_file) return train_features_dict, valid_features_dict, valid_fedtures_all, test_fedtures_all, \ train_examples_num_dict, valid_examples_num_dict, fedlab-nlp/data/utils.py000066400000000000000000000000001427741542300155040ustar00rootroot00000000000000fedlab-nlp/fed_run.sh000066400000000000000000000027011427741542300150510ustar00rootroot00000000000000#!/bin/bash # shellcheck disable=SC2068 # 读取参数 idx=0 for i in $@ do args[${idx}]=$i let "idx=${idx}+1" done # 分离参数 run_dirs=${args[0]} task_name=${args[1]} fl_algorithm=${args[2]} port=${args[3]} # 读取 GPU 参数 idx=0 for(( i=4;i<${#args[@]};i++ )) do device[${idx}]=${args[i]} let "idx=${idx}+1" done world_size=${#device[@]} echo "world_size is ${world_size}" if [ ${task_name} == "qnli" ]; then max_seq=256 else max_seq=128 fi echo "${task_name}'s max_seq is ${max_seq}" CUDA_VISIBLE_DEVICES=${device[0]} python main.py \ --model_name_or_path ${run_dirs}/pretrain/nlp/bert-base-uncased/ \ --output_dir ${run_dirs}/output/fedglue \ --rank 0 \ --task_name ${task_name} \ --fl_algorithm ${fl_algorithm} \ --raw_dataset_path ${run_dirs}/data/fedglue \ --partition_dataset_path ${run_dirs}/data/fedglue \ --max_seq_length ${max_seq} \ --world_size ${world_size} \ --port ${port} & sleep 2s for(( i=1;i<${world_size};i++)) do { echo "client ${i} started" CUDA_VISIBLE_DEVICES=${device[i]} python main.py \ --model_name_or_path ${run_dirs}/pretrain/nlp/bert-base-uncased/ \ --output_dir ${run_dirs}/output/fedglue \ --rank ${i} \ --task_name ${task_name} \ --fl_algorithm ${fl_algorithm} \ --raw_dataset_path ${run_dirs}/data/fedglue \ --partition_dataset_path ${run_dirs}/data/fedglue \ --max_seq_length ${max_seq} \ --world_size ${world_size} \ --port ${port} & sleep 2s } done wait fedlab-nlp/fedlab/000077500000000000000000000000001427741542300143105ustar00rootroot00000000000000fedlab-nlp/fedlab/__init__.py000066400000000000000000000012331427741542300164200ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. __version__ = "1.2.1" fedlab-nlp/fedlab/core/000077500000000000000000000000001427741542300152405ustar00rootroot00000000000000fedlab-nlp/fedlab/core/__init__.py000066400000000000000000000013751427741542300173570ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from . import client, server, communicator from .network import DistNetwork from .network_manager import NetworkManager fedlab-nlp/fedlab/core/client/000077500000000000000000000000001427741542300165165ustar00rootroot00000000000000fedlab-nlp/fedlab/core/client/__init__.py000066400000000000000000000015571427741542300206370ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. ORDINARY_TRAINER = 0 SERIAL_TRAINER = 1 from .manager import ClientManager, ActiveClientManager, PassiveClientManager from .trainer import ClientTrainer, SGDClientTrainer from .serial_trainer import SerialTrainer, SubsetSerialTrainerfedlab-nlp/fedlab/core/client/manager.py000066400000000000000000000142131427741542300205030ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import torch from . import ORDINARY_TRAINER, SERIAL_TRAINER from ..network_manager import NetworkManager from ...utils import Logger, MessageCode class ClientManager(NetworkManager): """Base class for ClientManager. :class:`ClientManager` defines client activation in different communication stages. Args: network (DistNetwork): Network configuration and interfaces. trainer (BaseClient): Subclass of :class:`BaseClient`. Provides :meth:`local_process` and :attr:`uplink_package`. Define local client trainers procedure. """ def __init__(self, network, trainer): super().__init__(network) self._trainer = trainer def setup(self): """Initialization stage. :class:`ClientManager` reports number of clients simulated by current client process. """ super().setup() tensor = torch.Tensor([self._trainer.client_num]).int() self._network.send(content=tensor, message_code=MessageCode.SetUp, dst=0) class PassiveClientManager(ClientManager): """Passive communication :class:`NetworkManager` for client in synchronous FL pattern. Args: network (DistNetwork): Network configuration and interfaces. trainer (BaseClient): Subclass of :class:`BaseClient`. Provides :meth:`local_process` and :attr:`uplink_package`. Define local client trainers procedure. logger (Logger, optional): Object of :class:`Logger`. """ def __init__(self, network, trainer, logger=None): super().__init__(network, trainer) self._LOGGER = Logger() if logger is None else logger def main_loop(self): """Actions to perform when receiving a new message, including local trainers. Main procedure of each client: 1. client waits for data from server (PASSIVELY). 2. after receiving data, client start local model trainers procedure. 3. client synchronizes with server actively. """ while True: sender_rank, message_code, payload = self._network.recv(src=0) if message_code == MessageCode.Exit: # client exit feedback if self._network.rank == self._network.world_size - 1: self._network.send(message_code=MessageCode.Exit, dst=0) break elif message_code == MessageCode.ParameterUpdate: id_list, payload = payload[0].to( torch.int32).tolist(), payload[1:] # check the trainer type if self._trainer.type == SERIAL_TRAINER: self._trainer.local_process(id_list=id_list, payload=payload) elif self._trainer.type == ORDINARY_TRAINER: assert len(id_list) == 1 self._trainer.local_process(payload=payload) self.synchronize() else: raise ValueError( "Invalid MessageCode {}. Please check MessageCode list.". format(message_code)) def synchronize(self): """Synchronize with server""" self._LOGGER.info("Uploading information to server.") self._network.send(content=self._trainer.uplink_package, message_code=MessageCode.ParameterUpdate, dst=0) class ActiveClientManager(ClientManager): """Active communication :class:`NetworkManager` for client in asynchronous FL pattern. Args: network (DistNetwork): Network configuration and interfaces. trainer (BaseClient): Subclass of :class:`BaseClient`. Provides :meth:`local_process` and :attr:`uplink_package`. Define local client trainers procedure. logger (Logger, optional): Object of :class:`Logger`. """ def __init__(self, network, trainer, logger=None): super().__init__(network, trainer) self._LOGGER = Logger() if logger is None else logger def main_loop(self): """Actions to perform on receiving new message, including local trainers 1. client requests data from server (ACTIVELY). 2. after receiving data, client will train local model. 3. client will synchronize with server actively. """ while True: # request model actively self.request() # waits for data from server _, message_code, payload = self._network.recv(src=0) if message_code == MessageCode.Exit: # client exit feedback if self._network.rank == self._network.world_size - 1: self._network.send(message_code=MessageCode.Exit, dst=0) break elif message_code == MessageCode.ParameterUpdate: self._trainer.local_process(payload) self.synchronize() else: raise ValueError( "Invalid MessageCode {}. Please check MessageCode Enum.". format(message_code)) def request(self): """BaseClient request""" self._LOGGER.info("request parameter procedure.") self._network.send(message_code=MessageCode.ParameterRequest, dst=0) def synchronize(self): """Synchronize with server""" self._LOGGER.info("Uploading information to server.") self._network.send(content=self._trainer.uplink_package, message_code=MessageCode.ParameterUpdate, dst=0) fedlab-nlp/fedlab/core/client/serial_trainer.py000066400000000000000000000144331427741542300221000ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import torch from utils import registry from .trainer import ClientTrainer from ..client import SERIAL_TRAINER from ...utils import Logger, SerializationTool from ...utils.dataset.sampler import SubsetSampler class SerialTrainer(ClientTrainer): """Base class. Simulate multiple clients in sequence in a single process. Args: model (torch.nn.Module): Model used in this federation. client_num (int): Number of clients in current trainer. cuda (bool): Use GPUs or not. Default: ``False``. """ def __init__(self, model, client_num, cuda=False, logger=None): super().__init__(model, cuda) self.client_num = client_num self.type = SERIAL_TRAINER # represent serial trainer self._LOGGER = Logger() if logger is None else logger self.param_list = [] @property def uplink_package(self): return self.param_list def _train_alone(self, model_parameters, train_loader): """Train local model with :attr:`model_parameters` on :attr:`train_loader`. Args: model_parameters (torch.Tensor): Serialized model parameters of one model. train_loader (torch.utils.data.DataLoader): :class:`torch.utils.data.DataLoader` for this client. """ raise NotImplementedError def _get_dataloader(self, client_id): """Get :class:`DataLoader` for ``client_id``.""" raise NotImplementedError def local_process(self, id_list, payload): """Train local model with different dataset according to client id in ``id_list``. Args: id_list (list[int]): BaseClient id in this trainers serial. payload (list[torch.Tensor]): communication payload from server. """ self.param_list = [] model_parameters = payload[0] self._LOGGER.info( "Local trainers with client id list: {}".format(id_list)) for idx in id_list: self._LOGGER.info( "Starting trainers procedure of client [{}]".format(idx)) data_loader = self._get_dataloader(client_id=idx) self._train_alone(model_parameters=model_parameters, train_loader=data_loader) self.param_list.append(self.model_parameters) return self.param_list class SubsetSerialTrainer(SerialTrainer): """Train multiple clients in a single process. Customize :meth:`_get_dataloader` or :meth:`_train_alone` for specific algorithm design in clients. Args: model (torch.nn.Module): Model used in this federation. dataset (torch.utils.data.Dataset): Local dataset for this group of clients. data_slices (list[list]): Subset of indices of dataset. logger (Logger, optional): Object of :class:`Logger`. cuda (bool): Use GPUs or not. Default: ``False``. args (dict): Uncertain variables. Default: ``{"epochs": 5, "batch_size": 100, "lr": 0.1}`` .. note:: ``len(data_slices) == client_num``, that is, each sub-index of :attr:`dataset` corresponds to a client's local dataset one-by-one. """ def __init__(self, model, dataset, data_slices, logger=None, cuda=False, args={ "epochs": 5, "batch_size": 100, "lr": 0.1 }) -> None: super(SubsetSerialTrainer, self).__init__(model=model, client_num=len(data_slices), cuda=cuda, logger=logger) self.dataset = dataset self.data_slices = data_slices # [0, client_num) self.args = args def _get_dataloader(self, client_id): """Return a trainers dataloader used in :meth:`train` for client with :attr:`id` Args: client_id (int): :attr:`client_id` of client to generate dataloader Note: :attr:`client_id` here is not equal to ``client_id`` in global FL setting. It is the index of client in current :class:`SerialTrainer`. Returns: :class:`DataLoader` for specific client's sub-dataset """ batch_size = self.args["batch_size"] train_loader = torch.utils.data.DataLoader( self.dataset, sampler=SubsetSampler(indices=self.data_slices[client_id], shuffle=True), batch_size=batch_size) return train_loader def _train_alone(self, model_parameters, train_loader): """Single round of local trainers for one client. Note: Overwrite this method to customize the PyTorch trainers pipeline. Args: model_parameters (torch.Tensor): serialized model parameters. train_loader (torch.utils.data.DataLoader): :class:`torch.utils.data.DataLoader` for this client. """ epochs, lr = self.args["epochs"], self.args["lr"] SerializationTool.deserialize_model(self._model, model_parameters) criterion = torch.nn.CrossEntropyLoss() optimizer = torch.optim.SGD(self._model.parameters(), lr=lr) self._model.train() for _ in range(epochs): for data, target in train_loader: if self.cuda: data = data.cuda(self.gpu) target = target.cuda(self.gpu) output = self.model(data) loss = criterion(output, target) optimizer.zero_grad() loss.backward() optimizer.step() return self.model_parameters fedlab-nlp/fedlab/core/client/trainer.py000066400000000000000000000114751427741542300205440ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from abc import abstractclassmethod, abstractproperty from typing import List import torch from ..client import ORDINARY_TRAINER from ..model_maintainer import ModelMaintainer from ...utils import Logger, SerializationTool class ClientTrainer(ModelMaintainer): """An abstract class representing a client trainer. In FedLab, we define the backend of client trainer show manage its local model. It should have a function to update its model called :meth:`local_process`. If you use our framework to define the activities of client, please make sure that your self-defined class should subclass it. All subclasses should overwrite :meth:`local_process` and property `uplink_package`. Args: model (torch.nn.Module): PyTorch model. cuda (bool): Use GPUs or not. """ def __init__(self, model, cuda): super().__init__(model, cuda) self.client_num = 1 # default is 1. self.type = ORDINARY_TRAINER @abstractproperty def uplink_package(self) -> List[torch.Tensor]: """Return a tensor list for uploading to server. This attribute will be called by client manager. Customize it for new algorithms. """ raise NotImplementedError @abstractclassmethod def local_process(self, payload) -> bool: """Manager of the upper layer will call this function with accepted payload In synchronous mode, return True to end current FL round. """ raise NotImplementedError def train(self): """Override this method to define the algorithm of trainers your model. This function should manipulate :attr:`self._model`""" raise NotImplementedError def evaluate(self): """Evaluate quality of local model.""" raise NotImplementedError class SGDClientTrainer(ClientTrainer): """BaseClient backend handler, this class provides data process method to upper layer. Args: model (torch.nn.Module): PyTorch model. data_loader (torch.utils.data.DataLoader): :class:`torch.utils.data.DataLoader` for this client. epochs (int): the number of local epoch. optimizer (torch.optim.Optimizer): optimizer for this client's model. criterion (torch.nn.Loss): loss function used in local trainers process. cuda (bool, optional): use GPUs or not. Default: ``False``. logger (Logger, optional): :object of :class:`Logger`. """ def __init__(self, model, data_loader, epochs, optimizer, criterion, cuda=False, logger=None): super(SGDClientTrainer, self).__init__(model, cuda) self._data_loader = data_loader self.epochs = epochs self.optimizer = optimizer self.criterion = criterion self._LOGGER = Logger() if logger is None else logger self.model_time = 0 @property def uplink_package(self): """Return a tensor list for uploading to server. This attribute will be called by client manager. Customize it for new algorithms. """ return [self.model_parameters] def local_process(self, payload): model_parameters = payload[0] self.train(model_parameters) def train(self, model_parameters) -> None: """BaseClient trains its local model on local dataset. Args: model_parameters (torch.Tensor): Serialized model parameters. """ SerializationTool.deserialize_model( self._model, model_parameters) # load parameters self._LOGGER.info("Local train procedure is running") for ep in range(self.epochs): self._model.train() for inputs, labels in self._data_loader: if self.cuda: inputs, labels = inputs.cuda(self.gpu), labels.cuda( self.gpu) outputs = self._model(inputs) loss = self.criterion(outputs, labels) self.optimizer.zero_grad() loss.backward() self.optimizer.step() self._LOGGER.info("Local train procedure is finished") fedlab-nlp/fedlab/core/communicator/000077500000000000000000000000001427741542300177405ustar00rootroot00000000000000fedlab-nlp/fedlab/core/communicator/__init__.py000066400000000000000000000023711427741542300220540ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """FedLab communication API""" HEADER_SENDER_RANK_IDX = 0 HEADER_RECEIVER_RANK_IDX = 1 HEADER_SLICE_SIZE_IDX = 2 HEADER_MESSAGE_CODE_IDX = 3 HEADER_DATA_TYPE_IDX = 4 DEFAULT_RECEIVER_RANK = -1 DEFAULT_SLICE_SIZE = 0 DEFAULT_MESSAGE_CODE_VALUE = 0 HEADER_SIZE = 5 # DATA TYPE CONSTANT INT8 = 0 INT16 = 1 INT32 = 2 INT64 = 3 FLOAT16 = 4 FLOAT32 = 5 FLOAT64 = 6 def dtype_torch2flab(torch_type): return supported_torch_dtypes.index(torch_type) def dtype_flab2torch(fedlab_type): return supported_torch_dtypes[fedlab_type] from .package import Package, supported_torch_dtypes from .processor import PackageProcessor fedlab-nlp/fedlab/core/communicator/package.py000066400000000000000000000156651427741542300217220ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import warnings from typing import List from copy import deepcopy import torch import torch.distributed as dist from . import HEADER_SENDER_RANK_IDX, HEADER_RECEIVER_RANK_IDX, HEADER_SLICE_SIZE_IDX, \ HEADER_MESSAGE_CODE_IDX, HEADER_DATA_TYPE_IDX from . import DEFAULT_SLICE_SIZE, DEFAULT_MESSAGE_CODE_VALUE from . import HEADER_SIZE from ...utils.message_code import MessageCode supported_torch_dtypes = [ torch.int8, torch.int16, torch.int32, torch.int64, torch.float16, torch.float32, torch.float64 ] class Package(object): """A basic network package data structure used in FedLab. Everything is Tensor in FedLab. Note: ``slice_size_i = tensor_i.shape[0]``, that is, every element in slices indicates the size of a sub-Tensor in content. :class:`Package` maintains 3 variables: - :attr:`header` : ``torch.Tensor([sender_rank, recv_rank, content_size, message_code, data_type])`` - :attr:`slices` : ``list[slice_size_1, slice_size_2]`` - :attr:`content` : ``torch.Tensor([tensor_1, tensor_2, ...])`` Args: message_code (MessageCode): Message code content (torch.Tensor, optional): Tensors contained in this package. """ def __init__(self, message_code=None, content=None): if message_code is None: message_code = DEFAULT_MESSAGE_CODE_VALUE else: if isinstance(message_code, MessageCode): message_code = message_code.value assert isinstance( message_code, int ), "message_code can only be MessageCode or integer, not {}".format( type(message_code)) # initialize header. The dtype of header is set as torch.int32 as default. self.header = torch.zeros(size=(HEADER_SIZE, ), dtype=torch.int32) if dist.is_initialized(): self.header[HEADER_SENDER_RANK_IDX] = dist.get_rank() else: self.header[HEADER_SENDER_RANK_IDX] = -1 self.header[HEADER_RECEIVER_RANK_IDX] = -1 # assigned by processor self.header[HEADER_MESSAGE_CODE_IDX] = message_code self.header[HEADER_SLICE_SIZE_IDX] = DEFAULT_SLICE_SIZE self.header[HEADER_DATA_TYPE_IDX] = -1 # assigned by processor # initialize content and slices self.slices = [] self.content = None self.dtype = None if isinstance(content, torch.Tensor): self.append_tensor(content) if isinstance(content, List): self.append_tensor_list(content) def append_tensor(self, tensor): """Append new tensor to :attr:`Package.content` Args: tensor (torch.Tensor): Tensor to append in content. """ if not isinstance(tensor, torch.Tensor): raise ValueError( "Invalid content type, expecting torch.Tensor but get {}". format(type(tensor))) shape = list(tensor.shape) slice = [tensor.numel(), len(shape)] + shape tensor = tensor.view(-1) if self.content is None: self.content = deepcopy(tensor) self.dtype = tensor.dtype else: if tensor.dtype is not self.dtype: warnings.warn( "The dtype of current tensor is {}. But package dtype is {}. The current data type will be casted to {} and fedlab do not guarantee lossless conversion." .format(tensor.dtype, self.dtype, self.dtype)) tensor = tensor.to(self.dtype) self.content = torch.cat((self.content, tensor)) self.slices += slice self.header[HEADER_SLICE_SIZE_IDX] = len(self.slices) def append_tensor_list(self, tensor_list): """Append a list of tensors to :attr:`Package.content`. Args: tensor_list (list[torch.Tensor]): A list of tensors to append to :attr:`Package.content`. """ for tensor in tensor_list: self.append_tensor(tensor) def to(self, dtype): if dtype in supported_torch_dtypes: self.dtype = dtype self.content.to(self.dtype) else: warnings.warn( "FedLab only supports following data types: torch.int8, torch.int16, torch.int32, torch.int64, torch.float16, torch.float32, torch.float64." ) @staticmethod def parse_content(slices, content): """Parse package content into a list of tensors Args: slices (list[int]): A list containing number of elements of each tensor. Each number is used as offset in parsing process. content (torch.Tensor): :attr:`Package.content`, a 1-D tensor composed of several 1-D tensors and their corresponding offsets. For more details about :class:`Package`. Returns: list[torch.Tensor]: A list of 1-D tensors parsed from ``content`` """ index = 0 # parse variable for content iter = 0 # parse variable for slices parse_result = [] while iter < len(slices): offset = slices[iter] # offset of content shape_len = slices[iter + 1] # offset of shape tuple shape = tuple(slices[iter + 2:iter + 2 + shape_len]) # obtain shape tuple seg_tensor = content[index:index + offset] reshape_tensor = seg_tensor.view(size=shape) # reshape parse_result.append(reshape_tensor) index += offset iter += shape_len + 2 return parse_result @staticmethod def parse_header(header): """Parse header to get information of current package. Args: header (torch.Tensor): :attr:`Package.header`, a 1-D tensor composed of 4 elements: ``torch.Tensor([sender_rank, recv_rank, slice_size, message_code, data_type])``. For more details about :class:`Package`. Returns: tuple: A tuple containing 5 elements: ``(sender_rank, recv_rank, slice_size, message_code, data_type)``. """ sender_rank = int(header[HEADER_SENDER_RANK_IDX]) receiver_rank = int(header[HEADER_RECEIVER_RANK_IDX]) slice_size = int(header[HEADER_SLICE_SIZE_IDX]) message_code = MessageCode(int(header[HEADER_MESSAGE_CODE_IDX])) data_type = int(header[HEADER_DATA_TYPE_IDX]) return sender_rank, receiver_rank, slice_size, message_code, data_type fedlab-nlp/fedlab/core/communicator/processor.py000066400000000000000000000113361427741542300223350ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import numpy as np import torch import torch.distributed as dist from .package import Package from . import HEADER_DATA_TYPE_IDX, HEADER_SIZE, HEADER_RECEIVER_RANK_IDX, HEADER_SLICE_SIZE_IDX, dtype_flab2torch, dtype_torch2flab class PackageProcessor(object): """Provide more flexible distributed tensor communication functions based on :func:`torch.distributed.send` and :func:`torch.distributed.recv`. :class:`PackageProcessor` defines the details of point-to-point package communication. EVERYTHING is :class:`torch.Tensor` in FedLab. """ @staticmethod def send_package(package, dst): """Three-segment tensor communication pattern based on ``torch.distributed`` Pattern is shown as follows: 1.1 sender: send a header tensor containing ``slice_size`` to receiver 1.2 receiver: receive the header, and get the value of ``slice_size`` and create a buffer for incoming slices of content 2.1 sender: send a list of slices indicating the size of every content size. 2.2 receiver: receive the slices list. 3.1 sender: send a content tensor composed of a list of tensors. 3.2 receiver: receive the content tensor, and parse it to obtain slices list using parser function """ def send_header(header, dst): header[HEADER_RECEIVER_RANK_IDX] = dst dist.send(header, dst=dst) def send_slices(slices, dst): np_slices = np.array(slices, dtype=np.int32) tensor_slices = torch.from_numpy(np_slices) dist.send(tensor_slices, dst=dst) def send_content(content, dst): dist.send(content, dst=dst) # body if package.dtype is not None: package.header[HEADER_DATA_TYPE_IDX] = dtype_torch2flab( package.dtype) # sender header firstly send_header(header=package.header, dst=dst) # if package got content, then send remain parts if package.header[HEADER_SLICE_SIZE_IDX] > 0: send_slices(slices=package.slices, dst=dst) send_content(content=package.content, dst=dst) @staticmethod def recv_package(src=None): """Three-segment tensor communication pattern based on ``torch.distributed`` Pattern is shown as follows: 1.1 sender: send a header tensor containing ``slice_size`` to receiver 1.2 receiver: receive the header, and get the value of ``slice_size`` and create a buffer for incoming slices of content 2.1 sender: send a list of slices indicating the size of every content size. 2.2 receiver: receive the slices list. 3.1 sender: send a content tensor composed of a list of tensors. 3.2 receiver: receive the content tensor, and parse it to obtain slices list using parser function """ def recv_header(src=src, parse=True): buffer = torch.zeros(size=(HEADER_SIZE, ), dtype=torch.int32) dist.recv(buffer, src=src) if parse is True: return Package.parse_header(buffer) else: return buffer def recv_slices(slices_size, src): buffer_slices = torch.zeros(size=(slices_size, ), dtype=torch.int32) dist.recv(buffer_slices, src=src) slices = [slc.item() for slc in buffer_slices] return slices def recv_content(slices, data_type, src): content_size = sum(slices) dtype = dtype_flab2torch(data_type) buffer = torch.zeros(size=(content_size, ), dtype=dtype) dist.recv(buffer, src=src) return Package.parse_content(slices, buffer) # body sender_rank, _, slices_size, message_code, data_type = recv_header( src=src) if slices_size > 0: slices = recv_slices(slices_size=slices_size, src=sender_rank) content = recv_content(slices, data_type, src=sender_rank) else: content = None return sender_rank, message_code, content fedlab-nlp/fedlab/core/coordinator.py000066400000000000000000000066101427741542300201400ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. class Coordinator(object): """Deal with the mapping relation between client id and process rank in FL system. Note BaseServer Manager creates a Coordinator following: 1. init network connection. 2. client send local group info (the number of client simulating in local) to server. 4. server receive all info and init a server Coordinator. Args: setup_dict (dict): A dict like {rank:client_num ...}, representing the map relation between process rank and client id. mode (str, optional): “GLOBAL” and "LOCAL". Coordinator will map client id to (rank, global id) or (rank, local id) according to mode. For example, client id 51 is in a machine which has 1 manager and serial trainer simulating 10 clients. LOCAL id means the index of its 10 clients. Therefore, global id 51 will be mapped into local id 1 (depending on setting). """ def __init__(self, setup_dict, mode='LOCAL') -> None: self.map = setup_dict self.mode = mode def map_id(self, id): """a map function from client id to (rank,local id) Args: id (int): client id Returns: rank, id : rank in distributed group and local id. """ m_id = id for rank, num in self.map.items(): if m_id >= num: m_id -= num else: local_id = m_id global_id = id ret_id = local_id if self.mode == 'LOCAL' else global_id return rank, ret_id def map_id_list(self, id_list): """a map function from id_list to dict{rank:local id} This can be very useful in Scale modules. Args: id_list (list(int)): a list of client id. Returns: map_dict (dict): contains process rank and its relative local client ids. """ map_dict = {} for id in id_list: rank, id = self.map_id(id) if rank in map_dict.keys(): map_dict[rank].append(id) else: map_dict[rank] = [id] return map_dict def switch(self): if self.mode == 'GLOBAL': self.mode = 'LOCAL' elif self.mode == 'LOCAL': self.mode = 'GLOBAL' else: raise ValueError("Invalid Map Mode {}".format(self.mode)) @property def total(self): return int(sum(self.map.values())) def __str__(self) -> str: return "Coordinator map information: {} \nMap mode: {} \nTotal: {}".format( self.map, self.mode, self.total) def __call__(self, info, *args, **kwds): if isinstance(info, int): return self.map_id(info) if isinstance(info, list): return self.map_id_list(info)fedlab-nlp/fedlab/core/model_maintainer.py000066400000000000000000000037341427741542300211300ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from ..utils.serialization import SerializationTool from ..utils.functional import get_best_gpu class ModelMaintainer(object): """Maintain PyTorch model. Provide necessary attributes and operation methods. More features with local or global model will be implemented here. Args: model (torch.nn.Module): PyTorch model. cuda (bool): use GPUs or not. """ def __init__(self, model, cuda) -> None: self.cuda = cuda if cuda: # dynamic gpu acquire. self.gpu = get_best_gpu() self._model = model.cuda(self.gpu) else: self._model = model.cpu() @property def model(self): """Return :class:`torch.nn.module`""" return self._model @property def model_parameters(self): """Return serialized model parameters.""" return SerializationTool.serialize_model(self._model) @property def model_gradients(self): """Return serialized model gradients.""" return SerializationTool.serialize_model_gradients(self._model) @property def shape_list(self): """Return shape of model parameters. Currently, this attributes used in tensor compression. """ shape_list = [param.shape for param in self._model.parameters()] return shape_list fedlab-nlp/fedlab/core/network.py000066400000000000000000000112661427741542300173110ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import os import torch import torch.distributed as dist import datetime from .communicator.processor import Package, PackageProcessor from ..utils import Logger type2byte = { torch.int8: 1, torch.int16: 2, torch.int32: 4, torch.int64: 8, torch.float16: 2, torch.float32: 4, torch.float64: 8 } class DistNetwork(object): """Manage ``torch.distributed`` network. Args: address (tuple): Address of this server in form of ``(SERVER_ADDR, SERVER_IP)`` world_size (int): the size of this distributed group (including server). rank (int): the rank of process in distributed group. ethernet (str) dist_backend (str or torch.distributed.Backend): :attr:`backend` of ``torch.distributed``. Valid values include ``mpi``, ``gloo``, and ``nccl``. Default: ``"gloo"``. """ def __init__(self, address, world_size, rank, ethernet=None, dist_backend="gloo"): super(DistNetwork, self).__init__() self.address = address self.rank = rank self.world_size = world_size self.dist_backend = dist_backend self.ethernet = ethernet self._LOGGER = Logger(log_name="network {}".format(self.rank)) self.send_volume_intotal = 0 # byte self.recv_volume_intotal = 0 # byte def init_network_connection(self): """Initialize ``torch.distributed`` communication group""" self._LOGGER.info(self.__str__()) if self.ethernet is not None: os.environ["GLOO_SOCKET_IFNAME"] = self.ethernet dist.init_process_group( backend=self.dist_backend, init_method="tcp://{}:{}".format(self.address[0], self.address[1]), rank=self.rank, world_size=self.world_size, timeout=datetime.timedelta(seconds=18000) ) def close_network_connection(self): """Destroy current ``torch.distributed`` process group""" self._LOGGER.info( "Overall communication volume: sent {} bytes, received {} bytes.". format(self.send_volume_intotal, self.recv_volume_intotal)) if dist.is_initialized(): dist.destroy_process_group() def send(self, content=None, message_code=None, dst=0, count=True): """Send tensor to process rank=dst""" pack = Package(message_code=message_code, content=content) PackageProcessor.send_package(pack, dst=dst) if pack.content is not None and count is True: self.send_volume_intotal += pack.content.numel() * type2byte[ pack.dtype] self._LOGGER.info( "Sent package to destination {}, message code {}, content length {}" .format(dst, message_code, 0 if pack.content is None else pack.content.numel())) def recv(self, src=None, count=True): """Receive tensor from process rank=src""" sender_rank, message_code, content = PackageProcessor.recv_package( src=src) if content is not None and count is True: volumn = sum([data.numel() for data in content]) # content from server to client, the first content is id_list. # remove the size of id_list in the count. if self.rank != 0: volumn -= content[0].numel() self.recv_volume_intotal += volumn * type2byte[content[0].dtype] self._LOGGER.info( "Received package from source {}, message code {}, content length {}" .format(sender_rank, message_code, 0 if content is None else volumn)) return sender_rank, message_code, content def __str__(self): info_str = "torch.distributed connection is initializing with ip address {}:{}, rank {}, world size: {}, backend {}, ethernet {}.".format( self.address[0], self.address[1], self.rank, self.world_size, self.dist_backend, self.ethernet, ) return info_str fedlab-nlp/fedlab/core/network_manager.py000066400000000000000000000035421427741542300210010ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from torch.multiprocessing import Process class NetworkManager(Process): """Abstract class Args: network (DistNetwork): object to manage torch.distributed network communication. """ def __init__(self, network): super(NetworkManager, self).__init__() self._network = network def run(self): """ Main Process: 1. Initialization stage. 2. FL communication stage. 3. Shutdown stage. Close network connection. """ self.setup() self.main_loop() self.shutdown() def setup(self, *args, **kwargs): """Initialize network connection and necessary setups. At first, ``self._network.init_network_connection()`` is required to be called. Overwrite this method to implement system setup message communication procedure. """ self._network.init_network_connection() def main_loop(self, *args, **kwargs): """Define the actions of communication stage.""" raise NotImplementedError def shutdown(self, *args, **kwargs): """Shutdown stage. Close the network connection in the end. """ self._network.close_network_connection() fedlab-nlp/fedlab/core/server/000077500000000000000000000000001427741542300165465ustar00rootroot00000000000000fedlab-nlp/fedlab/core/server/__init__.py000066400000000000000000000014341427741542300206610ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from .handler import SyncParameterServerHandler, AsyncParameterServerHandler from .manager import SynchronousServerManager, AsynchronousServerManager fedlab-nlp/fedlab/core/server/handler.py000066400000000000000000000205351427741542300205420ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from abc import abstractproperty from typing import List import random import torch from ..model_maintainer import ModelMaintainer from ...utils import Logger, Aggregators, SerializationTool class ParameterServerBackendHandler(ModelMaintainer): """An abstract class representing handler of parameter server. Please make sure that your self-defined server handler class subclasses this class Example: Read source code of :class:`SyncParameterServerHandler` and :class:`AsyncParameterServerHandler`. """ def __init__(self, model, cuda=False): super().__init__(model, cuda) @abstractproperty def downlink_package(self) -> List[torch.Tensor]: """Property for manager layer. BaseServer manager will call this property when activates clients.""" raise NotImplementedError @abstractproperty def if_stop(self) -> bool: """:class:`NetworkManager` keeps monitoring this attribute, and it will stop all related processes and threads when ``True`` returned.""" return False def _update_global_model(self, payload): """Override this function to define how to update global model (aggregation or optimization).""" raise NotImplementedError class SyncParameterServerHandler(ParameterServerBackendHandler): """Synchronous Parameter BaseServer Handler. Backend of synchronous parameter server: this class is responsible for backend computing in synchronous server. Synchronous parameter server will wait for every client to finish local trainers process before the next FL round. Details in paper: http://proceedings.mlr.press/v54/mcmahan17a.html Args: model (torch.nn.Module): Model used in this federation. global_round (int): stop condition. Shut down FL system when global round is reached. sample_ratio (float): The result of ``sample_ratio * client_num`` is the number of clients for every FL round. cuda (bool): Use GPUs or not. Default: ``False``. logger (Logger, optional): object of :class:`Logger`. """ def __init__(self, model, global_round, sample_ratio, cuda=False, logger=None): super(SyncParameterServerHandler, self).__init__(model, cuda) self._LOGGER = Logger() if logger is None else logger assert sample_ratio >= 0.0 and sample_ratio <= 1.0 # basic setting self.client_num_in_total = 0 self.sample_ratio = sample_ratio # client buffer self.client_buffer_cache = [] # stop condition self.global_round = global_round self.round = 0 @property def downlink_package(self): """Property for manager layer. BaseServer manager will call this property when activates clients.""" return [self.model_parameters] @property def if_stop(self): """:class:`NetworkManager` keeps monitoring this attribute, and it will stop all related processes and threads when ``True`` returned.""" return self.round >= self.global_round @property def client_num_per_round(self): return max(1, int(self.sample_ratio * self.client_num_in_total)) def sample_clients(self): """Return a list of client rank indices selected randomly. The client ID is from ``1`` to ``self.client_num_in_total + 1``.""" selection = random.sample(range(self.client_num_in_total), self.client_num_per_round) return selection def _update_global_model(self, payload): """Update global model with collected parameters from clients. Note: BaseServer handler will call this method when its ``client_buffer_cache`` is full. User can overwrite the strategy of aggregation to apply on :attr:`model_parameters_list`, and use :meth:`SerializationTool.deserialize_model` to load serialized parameters after aggregation into :attr:`self._model`. Args: payload (list[torch.Tensor]): A list of tensors passed by manager layer. """ assert len(payload) > 0 if len(payload) == 1: self.client_buffer_cache.append(payload[0].clone()) else: self.client_buffer_cache += payload # serial trainer assert len(self.client_buffer_cache) <= self.client_num_per_round if len(self.client_buffer_cache) == self.client_num_per_round: model_parameters_list = self.client_buffer_cache # use aggregator serialized_parameters = Aggregators.fedavg_aggregate( model_parameters_list) SerializationTool.deserialize_model(self._model, serialized_parameters) self.round += 1 # reset cache cnt self.client_buffer_cache = [] return True # return True to end this round. else: return False class AsyncParameterServerHandler(ParameterServerBackendHandler): """Asynchronous Parameter BaseServer Handler Update global model immediately after receiving a ParameterUpdate message Paper: https://arxiv.org/abs/1903.03934 Args: model (torch.nn.Module): Global model in server alpha (float): Weight used in async aggregation. total_time (int): Stop condition. Shut down FL system when total_time is reached. strategy (str): Adaptive strategy. ``constant``, ``hinge`` and ``polynomial`` is optional. Default: ``constant``. cuda (bool): Use GPUs or not. logger (Logger, optional): Object of :class:`Logger`. """ def __init__(self, model, alpha, total_time, strategy="constant", cuda=False, logger=None): super(AsyncParameterServerHandler, self).__init__(model, cuda) self._LOGGER = Logger() if logger is None else logger self.client_num_in_total = 0 self.time = 1 self.total_time = total_time # async aggregation params self.alpha = alpha self.strategy = strategy # "constant", "hinge", "polynomial" self.a = 10 self.b = 4 @property def if_stop(self): """:class:`NetworkManager` keeps monitoring this attribute, and it will stop all related processes and threads when ``True`` returned.""" return self.time >= self.total_time @property def downlink_package(self): return [self.model_parameters, torch.Tensor([self.time])] def _update_global_model(self, payload): client_model_parameters, model_time = payload[0], payload[1].item() """ "update global model from client_model_queue""" alpha_T = self._adapt_alpha(model_time) aggregated_params = Aggregators.fedasync_aggregate( self.model_parameters, client_model_parameters, alpha_T) # use aggregator SerializationTool.deserialize_model(self._model, aggregated_params) self.time += 1 def _adapt_alpha(self, receive_model_time): """update the alpha according to staleness""" staleness = self.time - receive_model_time if self.strategy == "constant": return torch.mul(self.alpha, 1) elif self.strategy == "hinge" and self.b is not None and self.a is not None: if staleness <= self.b: return torch.mul(self.alpha, 1) else: return torch.mul(self.alpha, 1 / (self.a * ((staleness - self.b) + 1))) elif self.strategy == "polynomial" and self.a is not None: return torch.mul(self.alpha, (staleness + 1)**(-self.a)) else: raise ValueError("Invalid strategy {}".format(self.strategy)) fedlab-nlp/fedlab/core/server/hierarchical/000077500000000000000000000000001427741542300211645ustar00rootroot00000000000000fedlab-nlp/fedlab/core/server/hierarchical/__init__.py000066400000000000000000000013371427741542300233010ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from .connector import ClientConnector, ServerConnector from .scheduler import Scheduler fedlab-nlp/fedlab/core/server/hierarchical/connector.py000066400000000000000000000203041427741542300235270ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import threading from time import sleep import torch from ...network_manager import NetworkManager from ...communicator.processor import PackageProcessor from ...coordinator import Coordinator from ....utils import MessageCode torch.multiprocessing.set_sharing_strategy("file_system") class Connector(NetworkManager): """Abstract class for basic Connector, which is a sub-module of :class:`Scheduler`. Connector inherits :class:`NetworkManager`, maintaining two Message Queue. One is for sending messages to collaborator, the other is for read messages from others. Note: Connector is a basic component for scheduler, Example code can be seen in ``scheduler.py``. Args: network (DistNetwork): Manage ``torch.distributed`` network communication. write_queue (torch.multiprocessing.Queue): Message queue to write. read_queue (torch.multiprocessing.Queue): Message queue to read. """ def __init__(self, network, write_queue, read_queue): super(Connector, self).__init__(network) # bidirectional message queues self.mq_read = read_queue self.mq_write = write_queue def process_meessage_queue(self): """Define the procedure of dealing with message queue.""" raise NotImplementedError class ServerConnector(Connector): """Connect with server. this process will act like a client. This class is a part of middle server which used in hierarchical structure. Args: network (DistNetwork): Network configuration and interfaces. write_queue (torch.multiprocessing.Queue): Message queue to write. read_queue (torch.multiprocessing.Queue): Message queue to read. logger (Logger, optional): object of :class:`Logger`. """ def __init__(self, network, write_queue, read_queue, logger): super(ServerConnector, self).__init__(network, write_queue, read_queue) self.mq_write = write_queue self.mq_read = read_queue self.group_client_num = 0 self._LOGGER = logger def run(self): """ Main Process: 1. Initialization stage. 2. FL communication stage. 3. Shutdown stage. Close network connection. """ self.setup() self.main_loop() self.shutdown() def setup(self, *args, **kwargs): super().setup() _, message_code, payload = self.mq_read.get() assert message_code == MessageCode.SetUp self.group_client_num = payload[0].item() self._network.send(content=torch.Tensor([self.group_client_num]).int(), message_code=message_code, dst=0) def main_loop(self, *args, **kwargs): # start a thread watching message queue watching_queue = threading.Thread(target=self.process_meessage_queue, daemon=True) watching_queue.start() while True: # server -> client sender, message_code, payload = PackageProcessor.recv_package() self.mq_write.put_nowait((sender, message_code, payload)) if message_code == MessageCode.Exit: # client exit feedback if self._network.rank == self._network.world_size - 1: self._network.send(message_code=MessageCode.Exit, dst=0) self._LOGGER.info("the main loop exit.") watching_queue.join() break def process_meessage_queue(self): """ client -> server directly transport. """ while True: sender, message_code, payload = self.mq_read.get() self._LOGGER.info( "[Queue-Thread: client -> server] recv data from rank {}, message code {}." .format(sender, message_code)) if message_code == MessageCode.Exit: self._LOGGER.info( "[Queue-Thread] process_meessage_queue thread exit.") break self._network.send(content=payload, message_code=message_code, dst=0) class ClientConnector(Connector): """Connect with clients. This class is a part of middle server which used in hierarchical structure. Args: network (DistNetwork): Network configuration and interfaces. write_queue (torch.multiprocessing.Queue): Message queue to write. read_queue (torch.multiprocessing.Queue): Message queue to read. """ def __init__(self, network, write_queue, read_queue, logger): super(ClientConnector, self).__init__(network, write_queue, read_queue) self.mq_read = read_queue self.mq_write = write_queue self.coordinator = None self.group_client_num = 0 self._LOGGER = logger def run(self): """ Main Process: 1. Initialization stage. 2. FL communication stage. 3. Shutdown stage. Close network connection. """ self.setup() self.main_loop() self.shutdown( ) # There is a bug. ClientConnector process will be blocked here. TODO:unsolved. def setup(self, *args, **kwargs): super().setup() rank_client_id_map = {} for rank in range(1, self._network.world_size): _, _, content = self._network.recv(src=rank) rank_client_id_map[rank] = content[0].item() self.coordinator = Coordinator(rank_client_id_map) self.group_client_num = self.coordinator.total self.mq_write.put_nowait((self._network.rank, MessageCode.SetUp, torch.Tensor([self.group_client_num]).int())) def main_loop(self): # start a thread to watch message queue watching_queue = threading.Thread(target=self.process_meessage_queue, daemon=True) watching_queue.start() while True: sender, message_code, payload = self._network.recv( ) # unexpected block. TODO: fix this. if message_code == MessageCode.Exit: self._LOGGER.info("main loop exit.") self.mq_write.put_nowait((None, MessageCode.Exit, None)) watching_queue.join() break self._LOGGER.info( "[client -> server] recv data from rank {}, message code {}.". format(sender, message_code)) self.mq_write.put_nowait((sender, message_code, payload)) def process_meessage_queue(self): """Process message queue Strategy of processing message from server. """ while True: # server -> client sender, message_code, payload = self.mq_read.get() self._LOGGER.info( "[Queue-Thread: server -> client] recv data from rank {}, message code {}." .format(sender, message_code)) # broadcast message id_list, payload = payload[0].to(torch.int32).tolist(), payload[1:] rank_dict = self.coordinator.map_id_list(id_list) for rank, values in rank_dict.items(): id_list = torch.Tensor(values).to(torch.int32) self._network.send(content=[id_list] + payload, message_code=message_code, dst=rank) if message_code == MessageCode.Exit: sleep(5) self._LOGGER.info( "[Queue-Thread] process_meessage_queue thread exit.") break fedlab-nlp/fedlab/core/server/hierarchical/scheduler.py000066400000000000000000000047621427741542300235250ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import torch from torch.multiprocessing import Queue from .connector import ClientConnector, ServerConnector from ....utils import Logger torch.multiprocessing.set_sharing_strategy("file_system") class Scheduler(): """Middle Topology for hierarchical communication pattern. Scheduler uses message queues to decouple connector modules. Args: net_upper (DistNetwork): Distributed network manager of server from upper level. net_lower (DistNetwork): Distributed network manager of clients from lower level. """ def __init__(self, net_upper, net_lower): super(Scheduler, self).__init__() self.__MQs = [Queue(), Queue()] self.net_upper = net_upper self.logger_upper = Logger( log_name="Scheduler{}-ServerConnector".format(self.net_upper.rank)) self.net_lower = net_lower self.logger_lower = Logger( log_name="Scheduler{}-ClientConnector".format(self.net_upper.rank)) def run(self): connect_server = ServerConnector(self.net_upper, write_queue=self.__MQs[1], read_queue=self.__MQs[0], logger=self.logger_upper) connect_client = ClientConnector(self.net_lower, write_queue=self.__MQs[0], read_queue=self.__MQs[1], logger=self.logger_lower) connect_server.start() connect_client.start() # This is a tiny bug. # The process with connect_client.join() is always blocked somehow in the shutdown stage. # You are will come to solve this. (I think it's some process synchronization problems) connect_server.join() # connect_client.join() connect_client.kill() fedlab-nlp/fedlab/core/server/manager.py000066400000000000000000000220731427741542300205360ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import threading import torch from torch.multiprocessing import Queue from ..network_manager import NetworkManager from ..coordinator import Coordinator from ...utils import Logger, MessageCode DEFAULT_SERVER_RANK = 0 class ServerManager(NetworkManager): """Base class of ServerManager. Args: network (DistNetwork): Network configuration and interfaces. handler (ParameterServerBackendHandler): Performe global model update procedure. """ def __init__(self, network, handler): super().__init__(network) self._handler = handler self.coordinator = None # initialized in setup stage. def setup(self): """Initialization Stage. - BaseServer accept local client num report from client manager. - Init a coordinator for client_id -> rank mapping. """ super().setup() rank_client_id_map = {} for rank in range(1, self._network.world_size): _, _, content = self._network.recv(src=rank) rank_client_id_map[rank] = content[0].item() self.coordinator = Coordinator(rank_client_id_map) if self._handler is not None: self._handler.client_num_in_total = self.coordinator.total class SynchronousServerManager(ServerManager): """Synchronous communication This is the top class in our framework which is mainly responsible for network communication of SERVER!. Synchronously communicate with clients following agreements defined in :meth:`main_loop`. Args: network (DistNetwork): Network configuration and interfaces. handler (ParameterServerBackendHandler): Backend calculation handler for parameter server. logger (Logger, optional): Object of :class:`Logger`. """ def __init__(self, network, handler, logger=None): super(SynchronousServerManager, self).__init__(network, handler) self._LOGGER = Logger() if logger is None else logger def setup(self): return super().setup() def main_loop(self): """Actions to perform in server when receiving a package from one client. BaseServer transmits received package to backend computation handler for aggregation or others manipulations. Loop: 1. activate clients for current trainers round. 2. listen for message from clients -> transmit received parameters to server backend. Note: Communication agreements related: user can overwrite this function to customize communication agreements. This method is key component connecting behaviors of :class:`ParameterServerBackendHandler` and :class:`NetworkManager`. Raises: Exception: Unexpected :class:`MessageCode`. """ while self._handler.if_stop is not True: activate = threading.Thread(target=self.activate_clients) activate.start() while True: sender_rank, message_code, payload = self._network.recv() if message_code == MessageCode.ParameterUpdate: if self._handler._update_global_model(payload): break else: raise Exception( "Unexpected message code {}".format(message_code)) def shutdown(self): """Shutdown stage.""" self.shutdown_clients() super().shutdown() def activate_clients(self): """Activate subset of clients to join in one FL round Manager will start a new thread to send activation package to chosen clients' process rank. The id of clients are obtained from :meth:`handler.sample_clients`. And their communication ranks are are obtained via coordinator. """ self._LOGGER.info("BaseClient activation procedure") clients_this_round = self._handler.sample_clients() rank_dict = self.coordinator.map_id_list(clients_this_round) self._LOGGER.info("BaseClient id list: {}".format(clients_this_round)) for rank, values in rank_dict.items(): downlink_package = self._handler.downlink_package id_list = torch.Tensor(values).to(downlink_package[0].dtype) self._network.send(content=[id_list] + downlink_package, message_code=MessageCode.ParameterUpdate, dst=rank) def shutdown_clients(self): """Shutdown all clients. Send package to each client with :attr:`MessageCode.Exit`. Note: Communication agreements related: User can overwrite this function to define package for exiting information. """ client_list = range(self._handler.client_num_in_total) rank_dict = self.coordinator.map_id_list(client_list) for rank, values in rank_dict.items(): downlink_package = self._handler.downlink_package id_list = torch.Tensor(values).to(downlink_package[0].dtype) self._network.send(content=[id_list] + downlink_package, message_code=MessageCode.Exit, dst=rank) # wait for client exit feedback _, message_code, _ = self._network.recv(src=self._network.world_size - 1) assert message_code == MessageCode.Exit class AsynchronousServerManager(ServerManager): """Asynchronous communication network manager for server This is the top class in our framework which is mainly responsible for network communication of SERVER!. Asynchronously communicate with clients following agreements defined in :meth:`mail_loop`. Args: network (DistNetwork): Network configuration and interfaces. handler (ParameterServerBackendHandler): Backend computation handler for parameter server. logger (Logger, optional): Object of :class:`Logger`. """ def __init__(self, network, handler, logger=None): super(AsynchronousServerManager, self).__init__(network, handler) self._LOGGER = Logger() if logger is None else logger self.message_queue = Queue() def setup(self): return super().setup() def main_loop(self): """Communication agreements of asynchronous FL. - BaseServer receive ParameterRequest from client. Send model parameter to client. - BaseServer receive ParameterUpdate from client. Transmit parameters to queue waiting for aggregation. Raises: ValueError: invalid message code. """ updater = threading.Thread(target=self.updater_thread, daemon=True) updater.start() while self._handler.if_stop is not True: sender, message_code, payload = self._network.recv() if message_code == MessageCode.ParameterRequest: self._network.send(content=self._handler.downlink_package, message_code=MessageCode.ParameterUpdate, dst=sender) elif message_code == MessageCode.ParameterUpdate: self.message_queue.put((sender, message_code, payload)) else: raise ValueError( "Unexpected message code {}.".format(message_code)) def shutdown(self): self.shutdown_clients() super().shutdown() def updater_thread(self): """Asynchronous communication maintain a message queue. A new thread will be started to keep monitoring message queue.""" while self._handler.if_stop is not True: _, message_code, payload = self.message_queue.get() self._handler._update_global_model(payload) assert message_code == MessageCode.ParameterUpdate def shutdown_clients(self): """Shutdown all clients. Send package to clients with ``MessageCode.Exit``. """ for rank in range(1, self._network.world_size): _, message_code, _ = self._network.recv(src=rank) # client request if message_code == MessageCode.ParameterUpdate: self._network.recv( src=rank ) # the next package is model request, which is ignored in shutdown stage. self._network.send(message_code=MessageCode.Exit, dst=rank) # wait for client exit feedback _, message_code, _ = self._network.recv(src=self._network.world_size - 1) assert message_code == MessageCode.Exit fedlab-nlp/fedlab/utils/000077500000000000000000000000001427741542300154505ustar00rootroot00000000000000fedlab-nlp/fedlab/utils/__init__.py000066400000000000000000000014271427741542300175650ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from .aggregator import Aggregators from .logger import Logger from .message_code import MessageCode from .serialization import SerializationToolfedlab-nlp/fedlab/utils/aggregator.py000066400000000000000000000040271427741542300201470ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import torch class Aggregators(object): """Define the algorithm of parameters aggregation""" @staticmethod def fedavg_aggregate(serialized_params_list, weights=None): """FedAvg aggregator Paper: http://proceedings.mlr.press/v54/mcmahan17a.html Args: serialized_params_list (list[torch.Tensor])): Merge all tensors following FedAvg. weights (list, numpy.array or torch.Tensor, optional): Weights for each params, the length of weights need to be same as length of ``serialized_params_list`` Returns: torch.Tensor """ if weights is None: weights = torch.ones(len(serialized_params_list)) if not isinstance(weights, torch.Tensor): weights = torch.tensor(weights) weights = weights / torch.sum(weights) assert torch.all(weights >= 0), "weights should be non-negative values" serialized_parameters = torch.sum( torch.stack(serialized_params_list, dim=-1) * weights, dim=-1) return serialized_parameters @staticmethod def fedasync_aggregate(server_param, new_param, alpha): """FedAsync aggregator Paper: https://arxiv.org/abs/1903.03934 """ serialized_parameters = torch.mul(1 - alpha, server_param) + \ torch.mul(alpha, new_param) return serialized_parameters fedlab-nlp/fedlab/utils/dataset/000077500000000000000000000000001427741542300170755ustar00rootroot00000000000000fedlab-nlp/fedlab/utils/dataset/__init__.py000066400000000000000000000020771427741542300212140ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from .sampler import SubsetSampler, RawPartitionSampler, DictFileSampler from .slicing import noniid_slicing, random_slicing from .partition import DataPartitioner, BasicPartitioner, VisionPartitioner from .partition import CIFAR10Partitioner, CIFAR100Partitioner, FMNISTPartitioner, MNISTPartitioner, \ SVHNPartitioner from .partition import FCUBEPartitioner from .partition import AdultPartitioner, RCV1Partitioner, CovtypePartitioner fedlab-nlp/fedlab/utils/dataset/functional.py000066400000000000000000000357041427741542300216220ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import numpy as np import pandas as pd import warnings def split_indices(num_cumsum, rand_perm): """Splice the sample index list given number of each client. Args: num_cumsum (np.ndarray): Cumulative sum of sample number for each client. rand_perm (list): List of random sample index. Returns: dict: ``{ client_id: indices}``. """ client_indices_pairs = [(cid, idxs) for cid, idxs in enumerate(np.split(rand_perm, num_cumsum)[:-1])] client_dict = dict(client_indices_pairs) return client_dict def balance_split(num_clients, num_samples): """Assign same sample sample for each client. Args: num_clients (int): Number of clients for partition. num_samples (int): Total number of samples. Returns: numpy.ndarray: A numpy array consisting ``num_clients`` integer elements, each represents sample number of corresponding clients. """ num_samples_per_client = int(num_samples / num_clients) client_sample_nums = (np.ones(num_clients) * num_samples_per_client).astype( int) return client_sample_nums def lognormal_unbalance_split(num_clients, num_samples, unbalance_sgm): """Assign different sample number for each client using Log-Normal distribution. Sample numbers for clients are drawn from Log-Normal distribution. Args: num_clients (int): Number of clients for partition. num_samples (int): Total number of samples. unbalance_sgm (float): Log-normal variance. When equals to ``0``, the partition is equal to :func:`balance_partition`. Returns: numpy.ndarray: A numpy array consisting ``num_clients`` integer elements, each represents sample number of corresponding clients. """ num_samples_per_client = int(num_samples / num_clients) if unbalance_sgm != 0: client_sample_nums = np.random.lognormal(mean=np.log(num_samples_per_client), sigma=unbalance_sgm, size=num_clients) client_sample_nums = ( client_sample_nums / np.sum(client_sample_nums) * num_samples).astype(int) diff = np.sum(client_sample_nums) - num_samples # diff <= 0 # Add/Subtract the excess number starting from first client if diff != 0: for cid in range(num_clients): if client_sample_nums[cid] > diff: client_sample_nums[cid] -= diff break else: client_sample_nums = (np.ones(num_clients) * num_samples_per_client).astype(int) return client_sample_nums def dirichlet_unbalance_split(num_clients, num_samples, alpha): """Assign different sample number for each client using Log-Normal distribution. Sample numbers for clients are drawn from Log-Normal distribution. Args: num_clients (int): Number of clients for partition. num_samples (int): Total number of samples. alpha (float): Dirichlet concentration parameter Returns: numpy.ndarray: A numpy array consisting ``num_clients`` integer elements, each represents sample number of corresponding clients. """ min_size = 0 while min_size < 10: proportions = np.random.dirichlet(np.repeat(alpha, num_clients)) proportions = proportions / proportions.sum() min_size = np.min(proportions * num_samples) client_sample_nums = (proportions * num_samples).astype(int) return client_sample_nums def homo_partition(client_sample_nums, num_samples): """Partition data indices in IID way given sample numbers for each clients. Args: client_sample_nums (numpy.ndarray): Sample numbers for each clients. num_samples (int): Number of samples. Returns: dict: ``{ client_id: indices}``. """ rand_perm = np.random.permutation(num_samples) num_cumsum = np.cumsum(client_sample_nums).astype(int) client_dict = split_indices(num_cumsum, rand_perm) return client_dict def hetero_dir_partition(targets, num_clients, num_classes, dir_alpha, min_require_size=None): """ Non-iid partition based on Dirichlet distribution. The method is from "hetero-dir" partition of `Bayesian Nonparametric Federated Learning of Neural Networks `_ and `Federated Learning with Matched Averaging `_. This method simulates heterogeneous partition for which number of data points and class proportions are unbalanced. Samples will be partitioned into :math:`J` clients by sampling :math:`p_k \sim \\text{Dir}_{J}({\\alpha})` and allocating a :math:`p_{p,j}` proportion of the samples of class :math:`k` to local client :math:`j`. Sample number for each client is decided in this function. Args: targets (list or numpy.ndarray): Sample targets. Unshuffled preferred. num_clients (int): Number of clients for partition. num_classes (int): Number of classes in samples. dir_alpha (float): Parameter alpha for Dirichlet distribution. min_require_size (int, optional): Minimum required sample number for each client. If set to ``None``, then equals to ``num_classes``. Returns: dict: ``{ client_id: indices}``. """ if min_require_size is None: min_require_size = num_classes if not isinstance(targets, np.ndarray): targets = np.array(targets) num_samples = targets.shape[0] min_size = 0 while min_size < min_require_size: idx_batch = [[] for _ in range(num_clients)] # for each class in the dataset for k in range(num_classes): idx_k = np.where(targets == k)[0] np.random.shuffle(idx_k) proportions = np.random.dirichlet( np.repeat(dir_alpha, num_clients)) # Balance proportions = np.array( [p * (len(idx_j) < num_samples / num_clients) for p, idx_j in zip(proportions, idx_batch)]) proportions = proportions / proportions.sum() proportions = (np.cumsum(proportions) * len(idx_k)).astype(int)[:-1] idx_batch = [idx_j + idx.tolist() for idx_j, idx in zip(idx_batch, np.split(idx_k, proportions))] min_size = min([len(idx_j) for idx_j in idx_batch]) client_dict = dict() for cid in range(num_clients): np.random.shuffle(idx_batch[cid]) client_dict[cid] = np.array(idx_batch[cid]) return client_dict def shards_partition(targets, num_clients, num_shards): """Non-iid partition used in FedAvg `paper `_. Args: targets (list or numpy.ndarray): Sample targets. Unshuffled preferred. num_clients (int): Number of clients for partition. num_shards (int): Number of shards in partition. Returns: dict: ``{ client_id: indices}``. """ if not isinstance(targets, np.ndarray): targets = np.array(targets) num_samples = targets.shape[0] size_shard = int(num_samples / num_shards) if num_samples % num_shards != 0: warnings.warn("warning: length of dataset isn't divided exactly by num_shards. " "Some samples will be dropped.") shards_per_client = int(num_shards / num_clients) if num_shards % num_clients != 0: warnings.warn("warning: num_shards isn't divided exactly by num_clients. " "Some shards will be dropped.") indices = np.arange(num_samples) # sort sample indices according to labels indices_targets = np.vstack((indices, targets)) indices_targets = indices_targets[:, indices_targets[1, :].argsort()] # corresponding labels after sorting are [0, .., 0, 1, ..., 1, ...] sorted_indices = indices_targets[0, :] # permute shards idx, and slice shards_per_client shards for each client rand_perm = np.random.permutation(num_shards) num_client_shards = np.ones(num_clients) * shards_per_client # sample index must be int num_cumsum = np.cumsum(num_client_shards).astype(int) # shard indices for each client client_shards_dict = split_indices(num_cumsum, rand_perm) # map shard idx to sample idx for each client client_dict = dict() for cid in range(num_clients): shards_set = client_shards_dict[cid] current_indices = [ sorted_indices[shard_id * size_shard: (shard_id + 1) * size_shard] for shard_id in shards_set] client_dict[cid] = np.concatenate(current_indices, axis=0) return client_dict def client_inner_dirichlet_partition(targets, num_clients, num_classes, dir_alpha, client_sample_nums, verbose=True): """Non-iid Dirichlet partition. The method is from The method is from paper `Federated Learning Based on Dynamic Regularization `_. This function can be used by given specific sample number for all clients ``client_sample_nums``. It's different from :func:`hetero_dir_partition`. Args: targets (list or numpy.ndarray): Sample targets. num_clients (int): Number of clients for partition. num_classes (int): Number of classes in samples. dir_alpha (float): Parameter alpha for Dirichlet distribution. client_sample_nums (numpy.ndarray): A numpy array consisting ``num_clients`` integer elements, each represents sample number of corresponding clients. verbose (bool, optional): Whether to print partition process. Default as ``True``. Returns: dict: ``{ client_id: indices}``. """ if not isinstance(targets, np.ndarray): targets = np.array(targets) rand_perm = np.random.permutation(targets.shape[0]) targets = targets[rand_perm] class_priors = np.random.dirichlet(alpha=[dir_alpha] * num_classes, size=num_clients) prior_cumsum = np.cumsum(class_priors, axis=1) idx_list = [np.where(targets == i)[0] for i in range(num_classes)] class_amount = [len(idx_list[i]) for i in range(num_classes)] client_indices = [np.zeros(client_sample_nums[cid]).astype(np.int64) for cid in range(num_clients)] while np.sum(client_sample_nums) != 0: curr_cid = np.random.randint(num_clients) # If current node is full resample a client if verbose: print('Remaining Data: %d' % np.sum(client_sample_nums)) if client_sample_nums[curr_cid] <= 0: continue client_sample_nums[curr_cid] -= 1 curr_prior = prior_cumsum[curr_cid] while True: curr_class = np.argmax(np.random.uniform() <= curr_prior) # Redraw class label if no rest in current class samples if class_amount[curr_class] <= 0: continue class_amount[curr_class] -= 1 client_indices[curr_cid][client_sample_nums[curr_cid]] = \ idx_list[curr_class][class_amount[curr_class]] break client_dict = {cid: client_indices[cid] for cid in range(num_clients)} return client_dict def label_skew_quantity_based_partition(targets, num_clients, num_classes, major_classes_num): """Label-skew:quantity-based partition. For details, please check `Federated Learning on Non-IID Data Silos: An Experimental Study `_. Args: targets (np.ndarray): Labels od dataset. num_clients (int): Number of clients. num_classes (int): Number of unique classes. major_classes_num (int): Number of classes for each client, should be less then ``num_classes``. Returns: dict: ``{ client_id: indices}``. """ idx_batch = [np.ndarray(0, dtype=np.int64) for _ in range(num_clients)] # only for major_classes_num < num_classes. # if major_classes_num = num_classes, it equals to IID partition times = [0 for _ in range(num_classes)] contain = [] for cid in range(num_clients): current = [cid % num_classes] times[cid % num_classes] += 1 j = 1 while j < major_classes_num: ind = np.random.randint(num_classes) if ind not in current: j += 1 current.append(ind) times[ind] += 1 contain.append(current) for k in range(num_classes): idx_k = np.where(targets == k)[0] np.random.shuffle(idx_k) split = np.array_split(idx_k, times[k]) ids = 0 for cid in range(num_clients): if k in contain[cid]: idx_batch[cid] = np.append(idx_batch[cid], split[ids]) ids += 1 client_dict = {cid: idx_batch[cid] for cid in range(num_clients)} return client_dict def fcube_synthetic_partition(data): """Feature-distribution-skew:synthetic partition. Synthetic partition for FCUBE dataset. This partition is from `Federated Learning on Non-IID Data Silos: An Experimental Study `_. Args: data (np.ndarray): Data of dataset :class:`FCUBE`. Returns: dict: ``{ client_id: indices}``. """ num_clients = 4 client_indices = [[] for _ in range(num_clients)] for idx, sample in enumerate(data): p1, p2, p3 = sample if (p1 > 0 and p2 > 0 and p3 > 0) or (p1 < 0 and p2 < 0 and p3 < 0): client_indices[0].append(idx) elif (p1 > 0 and p2 > 0 and p3 < 0) or (p1 < 0 and p2 < 0 and p3 > 0): client_indices[1].append(idx) elif (p1 > 0 and p2 < 0 and p3 > 0) or (p1 < 0 and p2 > 0 and p3 < 0): client_indices[2].append(idx) else: client_indices[3].append(idx) client_dict = {cid: np.array(client_indices[cid]).astype(int) for cid in range(num_clients)} return client_dict def samples_num_count(client_dict, num_clients): """Return sample count for all clients in ``client_dict``. Args: client_dict (dict): Data partition result for different clients. num_clients (int): Total number of clients. Returns: pandas.DataFrame """ client_samples_nums = [[cid, client_dict[cid].shape[0]] for cid in range(num_clients)] client_sample_count = pd.DataFrame(data=client_samples_nums, columns=['client', 'num_samples']).set_index('client') return client_sample_count fedlab-nlp/fedlab/utils/dataset/partition.py000066400000000000000000000454011427741542300214640ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from abc import ABC, abstractmethod import numpy as np from . import functional as F class DataPartitioner(ABC): """Base class for data partition in federated learning. Examples of :class:`DataPartitioner`: :class:`BasicPartitioner`, :class:`CIFAR10Partitioner`. Details and tutorials of different data partition and datasets, please check `Federated Dataset and DataPartitioner `_. """ @abstractmethod def _perform_partition(self): raise NotImplementedError @abstractmethod def __getitem__(self, index): raise NotImplementedError @abstractmethod def __len__(self): raise NotImplementedError class CIFAR10Partitioner(DataPartitioner): """CIFAR10 data partitioner. Partition CIFAR10 given specific client number. Currently 6 supported partition schemes can be achieved by passing different combination of parameters in initialization: - ``balance=None`` - ``partition="dirichlet"``: non-iid partition used in `Bayesian Nonparametric Federated Learning of Neural Networks `_ and `Federated Learning with Matched Averaging `_. Refer to :func:`fedlab.utils.dataset.functional.hetero_dir_partition` for more information. - ``partition="shards"``: non-iid method used in FedAvg `paper `_. Refer to :func:`fedlab.utils.dataset.functional.shards_partition` for more information. - ``balance=True``: "Balance" refers to FL scenario that sample numbers for different clients are the same. Refer to :func:`fedlab.utils.dataset.functional.balance_partition` for more information. - ``partition="iid"``: Random select samples from complete dataset given sample number for each client. - ``partition="dirichlet"``: Refer to :func:`fedlab.utils.dataset.functional.client_inner_dirichlet_partition` for more information. - ``balance=False``: "Unbalance" refers to FL scenario that sample numbers for different clients are different. For unbalance method, sample number for each client is drown from Log-Normal distribution with variance ``unbalanced_sgm``. When ``unbalanced_sgm=0``, partition is balanced. Refer to :func:`fedlab.utils.dataset.functional.lognormal_unbalance_partition` for more information. The method is from paper `Federated Learning Based on Dynamic Regularization `_. - ``partition="iid"``: Random select samples from complete dataset given sample number for each client. - ``partition="dirichlet"``: Refer to :func:`fedlab.utils.dataset.functional.client_inner_dirichlet_partition` for more information. For detail usage, please check `Federated Dataset and DataPartitioner `_. Args: targets (list or numpy.ndarray): Targets of dataset for partition. Each element is in range of [0, 1, ..., 9]. num_clients (int): Number of clients for data partition. balance (bool, optional): Balanced partition over all clients or not. Default as ``True``. partition (str, optional): Partition type, only ``"iid"``, ``shards``, ``"dirichlet"`` are supported. Default as ``"iid"``. unbalance_sgm (float, optional): Log-normal distribution variance for unbalanced data partition over clients. Default as ``0`` for balanced partition. num_shards (int, optional): Number of shards in non-iid ``"shards"`` partition. Only works if ``partition="shards"``. Default as ``None``. dir_alpha (float, optional): Dirichlet distribution parameter for non-iid partition. Only works if ``partition="dirichlet"``. Default as ``None``. verbose (bool, optional): Whether to print partition process. Default as ``True``. seed (int, optional): Random seed. Default as ``None``. Returns: dict: ``{ client_id: indices}``. """ num_classes = 10 def __init__(self, targets, num_clients, balance=True, partition="iid", unbalance_sgm=0, num_shards=None, dir_alpha=None, verbose=True, seed=None): self.targets = np.array(targets) # with shape (num_samples,) self.num_samples = self.targets.shape[0] self.num_clients = num_clients self.client_dict = dict() self.partition = partition self.balance = balance self.dir_alpha = dir_alpha self.num_shards = num_shards self.unbalance_sgm = unbalance_sgm self.verbose = verbose # self.rng = np.random.default_rng(seed) # rng currently not supports randint np.random.seed(seed) # partition scheme check if balance is None: assert partition in ["dirichlet", "shards"], f"When balance=None, 'partition' only " \ f"accepts 'dirichlet' and 'shards'." elif isinstance(balance, bool): assert partition in ["iid", "dirichlet"], f"When balance is bool, 'partition' only " \ f"accepts 'dirichlet' and 'iid'." else: raise ValueError(f"'balance' can only be NoneType or bool, not {type(balance)}.") # perform partition according to setting self.client_dict = self._perform_partition() # get sample number count for each client self.client_sample_count = F.samples_num_count(self.client_dict, self.num_clients) def _perform_partition(self): if self.balance is None: if self.partition == "dirichlet": client_dict = F.hetero_dir_partition(self.targets, self.num_clients, self.num_classes, self.dir_alpha, min_require_size=10) else: # partition is 'shards' client_dict = F.shards_partition(self.targets, self.num_clients, self.num_shards) else: # if balance is True or False # perform sample number balance/unbalance partition over all clients if self.balance is True: client_sample_nums = F.balance_split(self.num_clients, self.num_samples) else: client_sample_nums = F.lognormal_unbalance_split(self.num_clients, self.num_samples, self.unbalance_sgm) # perform iid/dirichlet partition for each client if self.partition == "iid": client_dict = F.homo_partition(client_sample_nums, self.num_samples) else: # for dirichlet client_dict = F.client_inner_dirichlet_partition(self.targets, self.num_clients, self.num_classes, self.dir_alpha, client_sample_nums, self.verbose) return client_dict def __getitem__(self, index): """Obtain sample indices for client ``index``. Args: index (int): BaseClient ID. Returns: list: List of sample indices for client ID ``index``. """ return self.client_dict[index] def __len__(self): """Usually equals to number of clients.""" return len(self.client_dict) class CIFAR100Partitioner(CIFAR10Partitioner): """CIFAR100 data partitioner. This is a subclass of the :class:`CIFAR10Partitioner`. For details, please check `Federated Dataset and DataPartitioner `_. """ num_classes = 100 class BasicPartitioner(DataPartitioner): """Basic data partitioner. Basic data partitioner, supported partition: - label-distribution-skew:quantity-based - label-distribution-skew:distributed-based (Dirichlet) - quantity-skew (Dirichlet) - IID For more details, please check `Federated Learning on Non-IID Data Silos: An Experimental Study `_ and `Federated Dataset and DataPartitioner `_. Args: targets (list or numpy.ndarray): Sample targets. Unshuffled preferred. num_clients (int): Number of clients for partition. partition (str): Partition name. Only supports ``"noniid-#label"``, ``"noniid-labeldir"``, ``"unbalance"`` and ``"iid"`` partition schemes. dir_alpha (float): Parameter alpha for Dirichlet distribution. Only works if ``partition="noniid-labeldir"``. major_classes_num (int): Number of major class for each clients. Only works if ``partition="noniid-#label"``. verbose (bool): Whether output intermediate information. Default as ``True``. seed (int): Random seed. Default as ``None``. Returns: dict: ``{ client_id: indices}``. """ num_classes = 2 def __init__(self, targets, num_clients, partition='iid', dir_alpha=None, major_classes_num=1, verbose=True, seed=None): self.targets = np.array(targets) # with shape (num_samples,) self.num_samples = self.targets.shape[0] self.num_clients = num_clients self.client_dict = dict() self.partition = partition self.dir_alpha = dir_alpha self.verbose = verbose # self.rng = np.random.default_rng(seed) # rng currently not supports randint np.random.seed(seed) if partition == "noniid-#label": # label-distribution-skew:quantity-based assert isinstance(major_classes_num, int), f"'major_classes_num' should be integer, " \ f"not {type(major_classes_num)}." assert major_classes_num > 0, f"'major_classes_num' should be positive." assert major_classes_num < self.num_classes, f"'major_classes_num' for each client " \ f"should be less than number of total " \ f"classes {self.num_classes}." self.major_classes_num = major_classes_num elif partition in ["noniid-labeldir", "unbalance"]: # label-distribution-skew:distributed-based (Dirichlet) and quantity-skew (Dirichlet) assert dir_alpha > 0, f"Parameter 'dir_alpha' for Dirichlet distribution should be " \ f"positive." elif partition == "iid": # IID pass else: raise ValueError( f"tabular data partition only supports 'noniid-#label', 'noniid-labeldir', " f"'unbalance', 'iid'. {partition} is not supported.") self.client_dict = self._perform_partition() # get sample number count for each client self.client_sample_count = F.samples_num_count(self.client_dict, self.num_clients) def _perform_partition(self): if self.partition == "noniid-#label": # label-distribution-skew:quantity-based client_dict = F.label_skew_quantity_based_partition(self.targets, self.num_clients, self.num_classes, self.major_classes_num) elif self.partition == "noniid-labeldir": # label-distribution-skew:distributed-based (Dirichlet) client_dict = F.hetero_dir_partition(self.targets, self.num_clients, self.num_classes, self.dir_alpha, min_require_size=10) elif self.partition == "unbalance": # quantity-skew (Dirichlet) client_sample_nums = F.dirichlet_unbalance_split(self.num_clients, self.num_samples, self.dir_alpha) client_dict = F.homo_partition(client_sample_nums, self.num_samples) else: # IID client_sample_nums = F.balance_split(self.num_clients, self.num_samples) client_dict = F.homo_partition(client_sample_nums, self.num_samples) return client_dict def __getitem__(self, index): return self.client_dict[index] def __len__(self): return len(self.client_dict) class VisionPartitioner(BasicPartitioner): """Data partitioner for vision data. Supported partition for vision data: - label-distribution-skew:quantity-based - label-distribution-skew:distributed-based (Dirichlet) - quantity-skew (Dirichlet) - IID For more details, please check `Federated Learning on Non-IID Data Silos: An Experimental Study `_. Args: targets (list or numpy.ndarray): Sample targets. Unshuffled preferred. num_clients (int): Number of clients for partition. partition (str): Partition name. Only supports ``"noniid-#label"``, ``"noniid-labeldir"``, ``"unbalance"`` and ``"iid"`` partition schemes. dir_alpha (float): Parameter alpha for Dirichlet distribution. Only works if ``partition="noniid-labeldir"``. major_classes_num (int): Number of major class for each clients. Only works if ``partition="noniid-#label"``. verbose (bool): Whether output intermediate information. Default as ``True``. seed (int): Random seed. Default as ``None``. Returns: dict: ``{ client_id: indices}``. """ num_classes = 10 def __init__(self, targets, num_clients, partition='iid', dir_alpha=None, major_classes_num=None, verbose=True, seed=None): super(VisionPartitioner, self).__init__(targets=targets, num_clients=num_clients, partition=partition, dir_alpha=dir_alpha, major_classes_num=major_classes_num, verbose=verbose, seed=seed) class MNISTPartitioner(VisionPartitioner): """Data partitioner for MNIST. For details, please check :class:`VisionPartitioner` and `Federated Dataset and DataPartitioner `_. """ num_features = 784 class FMNISTPartitioner(VisionPartitioner): """Data partitioner for FashionMNIST. For details, please check :class:`VisionPartitioner` and `Federated Dataset and DataPartitioner `_ """ num_features = 784 class SVHNPartitioner(VisionPartitioner): """Data partitioner for SVHN. For details, please check :class:`VisionPartitioner` and `Federated Dataset and DataPartitioner `_ """ num_features = 1024 class FCUBEPartitioner(DataPartitioner): """FCUBE data partitioner. FCUBE is a synthetic dataset for research in non-IID scenario with feature imbalance. This dataset and its partition methods are proposed in `Federated Learning on Non-IID Data Silos: An Experimental Study `_. Supported partition methods for FCUBE: - feature-distribution-skew:synthetic - IID For more details, please refer to Section (IV-B-b) of original paper. For detailed usage, please check `Federated Dataset and DataPartitioner `_. Args: data (numpy.ndarray): Data of dataset :class:`FCUBE`. partition (str): Partition type. Only supports `'synthetic'` and `'iid'`. """ num_classes = 2 num_clients = 4 # only accept partition for 4 clients def __init__(self, data, partition): if partition not in ['synthetic', 'iid']: raise ValueError( f"FCUBE only supports 'synthetic' and 'iid' partition, not {partition}.") self.partition = partition self.data = data if isinstance(data, np.ndarray): self.num_samples = data.shape[0] else: self.num_samples = len(data) self.client_dict = self._perform_partition() def _perform_partition(self): if self.partition == 'synthetic': # feature-distribution-skew:synthetic client_dict = F.fcube_synthetic_partition(self.data) else: # IID partition client_sample_nums = F.balance_split(self.num_clients, self.num_samples) client_dict = F.homo_partition(client_sample_nums, self.num_samples) return client_dict def __getitem__(self, index): return self.client_dict[index] def __len__(self): return self.num_clients class AdultPartitioner(BasicPartitioner): """Data partitioner for Adult. For details, please check :class:`BasicPartitioner` and `Federated Dataset and DataPartitioner `_ """ num_features = 123 num_classes = 2 class RCV1Partitioner(BasicPartitioner): """Data partitioner for RCV1. For details, please check :class:`BasicPartitioner` and `Federated Dataset and DataPartitioner `_ """ num_features = 47236 num_classes = 2 class CovtypePartitioner(BasicPartitioner): """Data partitioner for Covtype. For details, please check :class:`BasicPartitioner` and `Federated Dataset and DataPartitioner `_ """ num_features = 54 num_classes = 2 fedlab-nlp/fedlab/utils/dataset/sampler.py000066400000000000000000000061371427741542300211210ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import random import torch import torch.distributed as dist from ..functional import load_dict class SubsetSampler(torch.utils.data.Sampler): """Samples elements from a given list of indices, always in the same order once initialized. It is a :class:`Sampler` used in :class:`Dataloader`, that each partition will be fixed once initialized. Args: indices (list[int]): Indices in the whole set selected for subset shuffle (bool): shuffle the indices or not. """ def __init__(self, indices, shuffle=False) -> None: self.indices = indices if shuffle is True: random.shuffle(self.indices) def __iter__(self): return iter(self.indices) def __len__(self): return len(self.indices) class RawPartitionSampler(torch.utils.data.Sampler): """Partition dataset according to ``num_replicas``. Every client get a equal shard of dataset. Args: dataset (torch.utils.data.Dataset): client_id (int): num_replicas (int, optional): Number of data replications. Default ``None`` means total number of client processes. """ def __init__(self, dataset, client_id, num_replicas=None): self.indices = [index for index in range(len(dataset))] if num_replicas is None: self.num_replicas = dist.get_world_size() - 1 # world size includes 1 server process else: self.num_replicas = num_replicas assert isinstance(client_id, int), "client_id should be int, not {}".format(type(client_id)) assert client_id <= self.num_replicas, "client_id should not be greater than num_replicas={}".format( self.num_replicas) self.id = client_id self.num_samples = int(len(dataset) / self.num_replicas) def __iter__(self): local_indices = self.indices[(self.id - 1) * self.num_samples:self.id * self.num_samples] assert len(local_indices) == self.num_samples return iter(local_indices) def __len__(self): return self.num_samples class DictFileSampler(torch.utils.data.Sampler): """Get data sample indices given client id from data file with dict.""" def __init__(self, dict_file, client_id): data_indices = load_dict(dict_file) self.indices = data_indices[client_id] def __iter__(self): return iter(self.indices) def __len__(self): return len(self.indices) fedlab-nlp/fedlab/utils/dataset/slicing.py000066400000000000000000000064401427741542300211030ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """functions associated with data and dataset slicing""" import warnings import numpy as np def noniid_slicing(dataset, num_clients, num_shards): """Slice a dataset for non-IID. Args: dataset (torch.utils.data.Dataset): Dataset to slice. num_clients (int): Number of client. num_shards (int): Number of shards. Notes: The size of a shard equals to ``int(len(dataset)/num_shards)``. Each client will get ``int(num_shards/num_clients)`` shards. Returns： dict: ``{ 0: indices of dataset, 1: indices of dataset, ..., k: indices of dataset }`` """ total_sample_nums = len(dataset) size_of_shards = int(total_sample_nums / num_shards) if total_sample_nums % num_shards != 0: warnings.warn( "warning: the length of dataset isn't divided exactly by num_shard.some samples will be dropped." ) # the number of shards that each one of clients can get shard_pc = int(num_shards / num_clients) if num_shards % num_clients != 0: warnings.warn( "warning: num_shard isn't divided exactly by num_clients. some samples will be dropped." ) dict_users = {i: np.array([], dtype='int64') for i in range(num_clients)} labels = np.array(dataset.targets) idxs = np.arange(total_sample_nums) # sort sample indices according to labels idxs_labels = np.vstack((idxs, labels)) idxs_labels = idxs_labels[:, idxs_labels[1, :].argsort()] idxs = idxs_labels[0, :] # corresponding labels after sorting are [0, .., 0, 1, ..., 1, ...] # assign idx_shard = [i for i in range(num_shards)] for i in range(num_clients): rand_set = set(np.random.choice(idx_shard, shard_pc, replace=False)) idx_shard = list(set(idx_shard) - rand_set) for rand in rand_set: dict_users[i] = np.concatenate( (dict_users[i], idxs[rand * size_of_shards:(rand + 1) * size_of_shards]), axis=0) return dict_users def random_slicing(dataset, num_clients): """Slice a dataset randomly and equally for IID. Args： dataset (torch.utils.data.Dataset): a dataset for slicing. num_clients (int): the number of client. Returns： dict: ``{ 0: indices of dataset, 1: indices of dataset, ..., k: indices of dataset }`` """ num_items = int(len(dataset) / num_clients) dict_users, all_idxs = {}, [i for i in range(len(dataset))] for i in range(num_clients): dict_users[i] = list( np.random.choice(all_idxs, num_items, replace=False)) all_idxs = list(set(all_idxs) - set(dict_users[i])) return dict_users fedlab-nlp/fedlab/utils/functional.py000066400000000000000000000203371427741542300201710ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import torch import os import json import pynvml import numpy as np import pickle from collections import Counter class AverageMeter(object): """Record metrics information""" def __init__(self): self.reset() def reset(self): self.val = 0.0 self.avg = 0.0 self.sum = 0.0 self.count = 0.0 def update(self, val, n=1): self.val = val self.sum += val self.count += n self.avg = self.sum / self.count def evaluate(model, criterion, test_loader): """Evaluate classify task model accuracy.""" model.eval() gpu = next(model.parameters()).device loss_ = AverageMeter() acc_ = AverageMeter() with torch.no_grad(): for inputs, labels in test_loader: inputs = inputs.to(gpu) labels = labels.to(gpu) outputs = model(inputs) loss = criterion(outputs, labels) _, predicted = torch.max(outputs, 1) loss_.update(loss.item()) acc_.update(torch.sum(predicted.eq(labels)).item(), len(labels)) return loss_.sum, acc_.avg def read_config_from_json(json_file: str, user_name: str): """Read configs from `json_file` to get configs for `user_name` Args: json_file (str): path for json_file user_name (str): read configs for this user, it can be 'server' or 'client_id' Returns: a tuple with ip, port, world_size, rank about user with `user_name` Examples: read_config_from_json('../../../tests/data/configs.json', 'server') Notes: configs.json example as follows { "server": { "ip" : "127.0.0.1", "port": "3002", "world_size": 3, "rank": 0 }, "client_0": { "ip": "127.0.0.1", "port": "3002", "world_size": 3, "rank": 1 }, "client_1": { "ip": "127.0.0.1", "port": "3002", "world_size": 3, "rank": 2 } } """ with open(json_file) as f: config = json.load(f) config_info = config[user_name] return (config_info["ip"], config_info["port"], config_info["world_size"], config_info["rank"]) def get_best_gpu(): """Return gpu (:class:`torch.device`) with largest free memory.""" assert torch.cuda.is_available() pynvml.nvmlInit() deviceCount = pynvml.nvmlDeviceGetCount() if "CUDA_VISIBLE_DEVICES" in os.environ.keys() is not None: cuda_devices = [ int(device) for device in os.environ["CUDA_VISIBLE_DEVICES"].split(',') ] else: cuda_devices = range(deviceCount) assert max(cuda_devices) < deviceCount deviceMemory = [] for i in cuda_devices: handle = pynvml.nvmlDeviceGetHandleByIndex(i) mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle) deviceMemory.append(mem_info.free) deviceMemory = np.array(deviceMemory, dtype=np.int64) best_device_index = np.argmax(deviceMemory) return torch.device("cuda:%d" % (best_device_index)) def save_dict(dict, path): with open(path, 'wb') as f: pickle.dump(dict, f) def load_dict(path): with open(path, 'rb') as f: return pickle.load(f) def partition_report(targets, data_indices, class_num=None, verbose=True, file=None): """Generate data partition report for clients in ``data_indices``. Generate data partition report for each client according to ``data_indices``, including ratio of each class and dataset size in current client. Report can be printed in screen or into file. The output format is comma-separated values which can be read by :func:`pandas.read_csv` or :func:`csv.reader`. Args: targets (list or numpy.ndarray): Targets for all data samples, with each element is in range of ``0`` to ``class_num-1``. data_indices (dict): Dict of ``client_id: [data indices]``. class_num (int, optional): Total number of classes. If set to ``None``, then ``class_num = max(targets) + 1``. verbose (bool, optional): Whether print data partition report in screen. Default as ``True``. file (str, optional): Output file name of data partition report. If ``None``, then no output in file. Default as ``None``. Examples: First generate synthetic data labels and data partition to obtain ``data_indices`` (``{ client_id: sample indices}``): >>> sample_num = 15 >>> class_num = 4 >>> clients_num = 3 >>> num_per_client = int(sample_num/clients_num) >>> labels = np.random.randint(class_num, size=sample_num) # generate 15 labels, each label is 0 to 3 >>> rand_per = np.random.permutation(sample_num) >>> # partition synthetic data into 3 clients >>> data_indices = {0: rand_per[0:num_per_client], ... 1: rand_per[num_per_client:num_per_client*2], ... 2: rand_per[num_per_client*2:num_per_client*3]} Check ``data_indices`` may look like: >>> data_indices {0: array([8, 6, 5, 7, 2]), 1: array([ 3, 10, 14, 4, 1]), 2: array([13, 9, 12, 11, 0])} Now generate partition report for each client and each class: >>> partition_report(labels, data_indices, class_num=class_num, verbose=True, file=None) Class frequencies: client,class0,class1,class2,class3,Amount BaseClient 0,0.200,0.00,0.200,0.600,5 BaseClient 1,0.400,0.200,0.200,0.200,5 BaseClient 2,0.00,0.400,0.400,0.200,5 """ if not verbose and file is None: print("No partition report generated") return if not isinstance(targets, np.ndarray): targets = np.array(targets) if not class_num: class_num = max(targets) + 1 sorted_cid = sorted( data_indices.keys()) # sort client id in ascending order header_line = "Class frequencies:" col_name = "client," + ','.join([f"class{i}" for i in range(class_num)]) + ",Amount" if verbose: print(header_line) print(col_name) if file is not None: reports = [header_line, col_name] else: reports = None for client_id in sorted_cid: indices = data_indices[client_id] client_targets = targets[indices] client_sample_num = len( indices) # total number of samples of current client client_target_cnt = Counter( client_targets) # { cls1: num1, cls2: num2, ... } report_line = f"BaseClient {client_id:3d}," + \ ','.join([ f"{client_target_cnt[cls] / client_sample_num:.3f}" if cls in client_target_cnt else "0.00" for cls in range(class_num)]) + \ f",{client_sample_num}" if verbose: print(report_line) if file is not None: reports.append(report_line) if file is not None: fh = open(file, "w") fh.write("\n".join(reports)) fh.close() def accuracy(output, target, topk=(1, )): """Computes the top-k accuracy for the specified values of k, in range of [0, 1]""" maxk = max(topk) batch_size = target.size(0) _, pred = output.topk(maxk, 1, True, True) pred = pred.t() correct = pred.eq(target.view(1, -1).expand_as(pred)) res = [] for k in topk: correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True) res.append(correct_k.mul_(1.0 / batch_size)) return res fedlab-nlp/fedlab/utils/logger.py000066400000000000000000000036521427741542300173070ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import logging import sys logging.basicConfig( stream=sys.stdout, level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', datefmt='%H:%M:%S') logging.getLogger().setLevel(logging.INFO) class Logger(object): """record cmd info to file and print it to cmd at the same time Args: log_name (str): log name for output. log_file (str): a file path of log file. """ def __init__(self, log_name=None, log_file=None): if log_name is not None: self.logger = logging.getLogger(log_name) self.name = log_name else: logging.getLogger().setLevel(logging.INFO) self.logger = logging self.name = "root" if log_file is not None: handler = logging.FileHandler(log_file, mode='w') handler.setLevel(level=logging.INFO) formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') handler.setFormatter(formatter) self.logger.addHandler(handler) def info(self, log_str): """Print information to logger""" self.logger.info(log_str) def warning(self, warning_str): """Print warning to logger""" self.logger.warning(warning_str) fedlab-nlp/fedlab/utils/message_code.py000066400000000000000000000017261427741542300204460ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from enum import Enum class MessageCode(Enum): """Different types of messages between client and server that we support go here.""" # BaseServer and BaseClient communication agreements ParameterRequest = 0 GradientUpdate = 1 ParameterUpdate = 2 EvaluateParams = 3 Exit = 4 SetUp = 5 Activation = 6 fedlab-nlp/fedlab/utils/serialization.py000066400000000000000000000056071427741542300207070ustar00rootroot00000000000000# Copyright 2021 Peng Cheng Laboratory (http://www.szpclab.com/) and FedLab Authors (smilelab.group) # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # http://www.apache.org/licenses/LICENSE-2.0 # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import torch class SerializationTool(object): @staticmethod def serialize_model_gradients(model: torch.nn.Module) -> torch.Tensor: gradients = [param.grad.data.view(-1) for param in model.parameters()] m_gradients = torch.cat(gradients) m_gradients = m_gradients.cpu() return m_gradients @staticmethod def serialize_model(model: torch.nn.Module) -> torch.Tensor: """Unfold model parameters Unfold every layer of model, concate all of tensors into one. Return a `torch.Tensor` with shape (size, ). Args: model (torch.nn.Module): model to serialize. """ parameters = [param.data.view(-1) for param in model.parameters()] m_parameters = torch.cat(parameters) m_parameters = m_parameters.cpu() return m_parameters @staticmethod def deserialize_model(model: torch.nn.Module, serialized_parameters: torch.Tensor, mode="copy"): """Assigns serialized parameters to model.parameters. This is done by iterating through ``model.parameters()`` and assigning the relevant params in ``grad_update``. NOTE: this function manipulates ``model.parameters``. Args: model (torch.nn.Module): model to deserialize. serialized_parameters (torch.Tensor): serialized model parameters. mode (str): deserialize mode. "copy" or "add". """ current_index = 0 # keep track of where to read from grad_update for parameter in model.parameters(): numel = parameter.data.numel() size = parameter.data.size() if mode == "copy": parameter.data.copy_( serialized_parameters[current_index:current_index + numel].view(size)) elif mode == "add": parameter.data.add_( serialized_parameters[current_index:current_index + numel].view(size)) else: raise ValueError( "Invalid deserialize mode {}, require \"copy\" or \"add\" " .format(mode)) current_index += numel fedlab-nlp/main.py000066400000000000000000000005601427741542300143720ustar00rootroot00000000000000"""main function for FedLab-NLP""" from utils import registry from utils import build_config from utils import setup_logger, setup_imports def main(): setup_imports() setup_logger() config = build_config() trainer = registry.get_fl_class( config.federated_config.fl_algorithm)() trainer.train() if __name__ == "__main__": main() fedlab-nlp/models/000077500000000000000000000000001427741542300143565ustar00rootroot00000000000000fedlab-nlp/models/__init__.py000066400000000000000000000002271427741542300164700ustar00rootroot00000000000000from models.base_models import BaseModels from models.classification import SeqClassification __all__ = [ "BaseModels", "SeqClassification" ] fedlab-nlp/models/base_models.py000066400000000000000000000013411427741542300172040ustar00rootroot00000000000000 from abc import ABC, abstractmethod from utils import registry import torch.nn as nn class BaseModels(nn.Module): def __init__(self, task_name): super().__init__() self.task_name = task_name config = registry.get("config") self.model_config = config.model_config self.rank = config.federated_config.rank self.logger = registry.get("logger") def _build_config(self): raise NotImplementedError def _build_model(self): raise NotImplementedError def permutate_layers(self, model): raise NotImplementedError def freezed_layers(self, model): raise NotImplementedError def forward(self, inputs): raise NotImplementedError fedlab-nlp/models/classification.py000066400000000000000000000051441427741542300177270ustar00rootroot00000000000000"""SeqClassification Model For FedLab-NLP """ import copy from abc import ABC import torch from utils import registry from models.base_models import BaseModels from transformers import AutoConfig, AutoModelForSequenceClassification @registry.register_model("seq_classification") class SeqClassification(BaseModels, ABC): def __init__(self, task_name): super().__init__(task_name) self.num_labels = registry.get("num_labels") self.auto_config = self._build_config() self.backbone = self._build_model() def _build_config(self): auto_config = AutoConfig.from_pretrained( self.model_config.config_name if self.model_config.config_name else self.model_config.model_name_or_path, num_labels=self.num_labels, finetuning_task=self.task_name if self.task_name else None, # cache_dir=self.model_config.cache_dir, revision=self.model_config.model_revision, use_auth_token=True if self.model_config.use_auth_token else None, ) return auto_config def _build_model(self): backbone = AutoModelForSequenceClassification.from_pretrained( self.model_config.model_name_or_path, from_tf=bool(".ckpt" in self.model_config.model_name_or_path), config=self.auto_config, # cache_dir=self.model_config.cache_dir, revision=self.model_config.model_revision, use_auth_token=True if self.model_config.use_auth_token else None, # ignore_mismatched_sizes=self.model_config.ignore_mismatched_sizes, ) if getattr(self.model_config, "permutation_layers", None): backbone = self.permutate_layers(backbone) return backbone def permutate_layers(self, model): old_modules = model.bert.encoder.layer scrambled_modules = torch.nn.ModuleList() # Now iterate over all layers, # appending to the new module list according to the new order. if self.rank > 0: permutation = self.model_config.client_model_layers else: permutation = self.model_config.server_model_layers self.logger.debug(f"model's layer: {permutation}") for i in permutation: assert i <= 11, permutation scrambled_modules.append(old_modules[i]) # Create a copy of the model, modify it with the new list, and return model_copy = copy.deepcopy(model) model_copy.bert.encoder.layer = scrambled_modules return model_copy def forward(self, inputs): output = self.backbone(**inputs) return output fedlab-nlp/requirement.txt000066400000000000000000000001521427741542300161720ustar00rootroot00000000000000numpy wandb pandas loguru gpustat omegaconf matplotlib scikit-learn tensorboardX transformers torch==1.11 fedlab-nlp/run/000077500000000000000000000000001427741542300136775ustar00rootroot00000000000000fedlab-nlp/run/centralized/000077500000000000000000000000001427741542300162035ustar00rootroot00000000000000fedlab-nlp/run/centralized/__init__.py000066400000000000000000000001321427741542300203100ustar00rootroot00000000000000from run.centralized.client import CenClientTrainer __all__ = [ "CenClientTrainer" ] fedlab-nlp/run/centralized/client.py000066400000000000000000000004521427741542300200340ustar00rootroot00000000000000"""Centralized Training for FedLab-NLP""" from abc import ABC from trainers.BaseClient.base_client import BaseClientTrainer class CenClientTrainer(BaseClientTrainer, ABC): def __init__(self, model, train_dataset, valid_dataset): super().__init__(model, train_dataset, valid_dataset) fedlab-nlp/run/centralized/config.yaml000066400000000000000000000006621427741542300203400ustar00rootroot00000000000000data_config: dataset_name: glue federated_config: clients_num: 100 rounds: 100 alpha: 1.0 sample: 0.1 rank: -1 pson: true model_config: model_type: bert model_output_mode: seq_classification permutation_layers: true client_model_layers: [0,1,2] server_model_layers: [0,1,2] training_config: per_device_train_batch_size: 32 num_train_epochs: 5 learning_rate: 5e-5 metric_name: glue do_predict: truefedlab-nlp/run/centralized/trainer.py000066400000000000000000000015011427741542300202160ustar00rootroot00000000000000"""centralized Trainer""" import time from abc import ABC from trainers import BaseTrainer from utils import registry, cen_metric_save from run.centralized.client import CenClientTrainer @registry.register_fl_algorithm("centralized") class CentralizedTrainer(BaseTrainer, ABC): def __init__(self): super().__init__() self._before_training() def _build_client(self): self.client_trainer = CenClientTrainer( model=self.model, train_dataset=self.data.train_dataloader_dict, valid_dataset=self.data.valid_dataloader_dict, ) def on_client_end(self): if self.training_config.do_predict: self.client_trainer.test_on_client(self.data.test_dataloader) cen_metric_save(self.client_trainer, self.training_config, self.logger) fedlab-nlp/run/dry_run/000077500000000000000000000000001427741542300153615ustar00rootroot00000000000000fedlab-nlp/run/dry_run/config.yaml000066400000000000000000000007071427741542300175160ustar00rootroot00000000000000data_config: dataset_name: glue task_name: qqp federated_config: clients_num: 100 rounds: 100 alpha: 1.0 sample: 0.1 rank: -1 pson: true model_config: model_type: bert model_output_mode: seq_classification permutation_layers: false client_model_layers: [0,1,2] server_model_layers: [0,1,2] training_config: per_device_train_batch_size: 32 num_train_epochs: 5000 learning_rate: 5e-5 metric_name: glue do_predict: truefedlab-nlp/run/dry_run/trainer.py000066400000000000000000000010601427741542300173740ustar00rootroot00000000000000"""Dry running Trainer""" from abc import ABC from utils import registry from trainers import BaseTrainer from run.centralized.client import CenClientTrainer @registry.register_fl_algorithm("dry_run") class DryTrainer(BaseTrainer, ABC): def __init__(self): super().__init__() self._before_training() def _build_client(self): self.client_trainer = CenClientTrainer( model=self.model, train_dataset=self.data.train_dataloader_dict, valid_dataset=self.data.valid_dataloader_dict, ) fedlab-nlp/run/fedavg/000077500000000000000000000000001427741542300151335ustar00rootroot00000000000000fedlab-nlp/run/fedavg/__init__.py000066400000000000000000000005231427741542300172440ustar00rootroot00000000000000from run.fedavg.trainer import FedAvgTrainer from run.fedavg.server import FedAvgSyncServerHandler, FedAvgServerManager from run.fedavg.client import FedAvgClientTrainer, FedAvgClientManager __all__ = [ "FedAvgTrainer", "FedAvgClientTrainer", "FedAvgClientManager", "FedAvgServerManager", "FedAvgSyncServerHandler", ] fedlab-nlp/run/fedavg/client.py000066400000000000000000000006631427741542300167700ustar00rootroot00000000000000"""federated average client""" from abc import ABC from trainers.BaseClient import BaseClientTrainer, BaseClientManager class FedAvgClientTrainer(BaseClientTrainer, ABC): def __init__(self, model, train_dataset, valid_dataset): super().__init__(model, train_dataset, valid_dataset) class FedAvgClientManager(BaseClientManager, ABC): def __init__(self, network, trainer): super().__init__(network, trainer) fedlab-nlp/run/fedavg/config.yaml000066400000000000000000000006111427741542300172620ustar00rootroot00000000000000data_config: dataset_name: glue federated_config: clients_num: 100 rounds: 100 alpha: 1.0 sample: 0.1 model_config: model_type: bert model_output_mode: seq_classification permutation_layers: true client_model_layers: [0,1,2] server_model_layers: [0,1,2] training_config: per_device_train_batch_size: 32 num_train_epochs: 1 learning_rate: 5e-5 metric_name: glue fedlab-nlp/run/fedavg/server.py000066400000000000000000000006631427741542300170200ustar00rootroot00000000000000"""federated average server""" from abc import ABC from trainers.BaseServer import BaseSyncServerHandler, BaseServerManager class FedAvgSyncServerHandler(BaseSyncServerHandler, ABC): def __init__(self, model, valid_data, test_data): super().__init__(model, valid_data, test_data) class FedAvgServerManager(BaseServerManager, ABC): def __init__(self, network, handler): super().__init__(network, handler) fedlab-nlp/run/fedavg/trainer.py000066400000000000000000000022411427741542300171500ustar00rootroot00000000000000"""federated average trainer""" from utils import registry from trainers.FedBaseTrainer import BaseTrainer from run.fedavg.client import FedAvgClientTrainer, FedAvgClientManager from run.fedavg.server import FedAvgSyncServerHandler, FedAvgServerManager @registry.register_fl_algorithm("fedavg") class FedAvgTrainer(BaseTrainer): def __init__(self): super().__init__() self._before_training() def _build_server(self): self.handler = FedAvgSyncServerHandler( self.model, valid_data=self.data.valid_dataloader, test_data=self.data.test_dataloader ) self.server_manger = FedAvgServerManager( network=self.network, handler=self.handler, ) def _build_client(self): self.client_trainer = FedAvgClientTrainer( model=self.model, train_dataset=self.data.train_dataloader_dict, valid_dataset=self.data.valid_dataloader_dict, # data_slices=self.federated_config.clients_id_list, ) self.client_manager = FedAvgClientManager( trainer=self.client_trainer, network=self.network ) fedlab-nlp/test.py000066400000000000000000000000361427741542300144230ustar00rootroot00000000000000"""Test some code snippets""" fedlab-nlp/tools/000077500000000000000000000000001427741542300142335ustar00rootroot00000000000000fedlab-nlp/tools/glue_scripts/000077500000000000000000000000001427741542300167365ustar00rootroot00000000000000fedlab-nlp/tools/glue_scripts/glue.py000066400000000000000000000157471427741542300202620ustar00rootroot00000000000000import os import pickle import argparse from loguru import logger from sklearn.model_selection import train_test_split from transformers import glue_output_modes as output_modes from utils import make_sure_dirs from tools.glue_scripts.partition import GlueDataPartition from tools.glue_scripts.glue_utils import glue_processors as processors def parser_args(): parser = argparse.ArgumentParser() parser.add_argument("--data_dir", default=None, type=str, required=True, help="The input data dir. Should contain the .tsv files (or other data files) for the task.") parser.add_argument("--task", default=None, type=str, required=True, help="Task name") parser.add_argument("--output_dir", default=None, type=str, required=True, help="The output directory to save partition or raw data") parser.add_argument("--clients_num", default=None, type=int, required=True, help="All clients numbers") parser.add_argument("--alpha", default=None, type=float, help="The label skew degree.") parser.add_argument("--overwrite", default=None, type=int, help="overwrite") args = parser.parse_args() return args def load_glue_examples(args): task_name = args.task.lower() processor = processors[task_name]() output_mode = output_modes[task_name] label_list = processor.get_labels() # if task in ['mnli', 'mnli-mm'] and args.model_type in ['roberta']: # HACK(label indices are swapped in RoBERTa pretrained model) # label_list[1], label_list[2] = label_list[2], label_list[1] train_examples = processor.get_train_examples(args.data_dir) valid_examples = processor.get_dev_examples(args.data_dir) test_examples = processor.get_test_examples(args.data_dir) return train_examples, valid_examples, test_examples, output_mode, label_list def get_partition_data(examples, num_classes, num_clients, label_vocab, dir_alpha, partition, ): targets = [example.label for example in examples] clients_partition_data = GlueDataPartition( targets=targets, num_classes=num_classes, num_clients=num_clients, label_vocab=label_vocab, dir_alpha=dir_alpha, partition=partition, verbose=False ) partition_data = {} for idx in range(len(clients_partition_data)): client_idxs = clients_partition_data[idx] partition_data[idx] = client_idxs return partition_data def convert_glue_to_federated_pkl(args): logger.info("reading examples ...") if os.path.isfile(args.output_data_file): logger.info(f"Examples in {args.output_data_file} have existed ...") with open(args.output_data_file, "rb") as file: data = pickle.load(file) train_examples, valid_examples, test_examples = data["train"], data["valid"], data["test"] output_mode, label_list = data["output_mode"], data["label_list"] else: logger.info(f"Generating examples from {args.data_dir} ...") train_examples, original_valid_examples, original_test_examples, output_mode, label_list \ = load_glue_examples(args) # we need to split original valid_examples into new valid and test sets original_valid_examples_idx = [i for i in range(len(original_valid_examples))] original_valid_examples_label = [example.label for example in original_valid_examples] valid_idx, test_idx, valid_y, test_y = train_test_split( original_valid_examples_idx, original_valid_examples_label, test_size=0.5, random_state=42, stratify=original_valid_examples_label ) valid_examples = [original_valid_examples[idx] for idx in valid_idx] test_examples = [original_valid_examples[idx] for idx in test_idx] data = { "train": train_examples, "valid": valid_examples, "test": test_examples, "output_mode": output_mode, "label_list": label_list } with open(args.output_data_file, "wb") as file: pickle.dump(data, file) logger.info("partition data ...") if os.path.isfile(args.output_partition_file): logger.info("loading partition data ...") with open(args.output_partition_file, "rb") as file: partition_data = pickle.load(file) else: partition_data = {} logger.info(f"partition data's keys: {partition_data.keys()}") if f"clients={args.clients_num}_alpha={args.alpha}" in partition_data and not args.overwrite: logger.info(f"Partition method 'clients={args.clients_num}_alpha={args.alpha}' has existed " f"and overwrite={args.overwrite}, then skip") else: lable_mapping = {label: idx for idx, label in enumerate(label_list)} attribute = {"lable_mapping": lable_mapping, "label_list": label_list, "clients_num": args.clients_num, "alpha": args.alpha, "output_mode": output_mode } clients_partition_data = {"train": get_partition_data( examples=train_examples, num_classes=len(label_list), num_clients=args.clients_num, label_vocab=label_list, dir_alpha=args.alpha, partition="dirichlet" ), "valid": get_partition_data( examples=valid_examples, num_classes=len(label_list), num_clients=args.clients_num, label_vocab=label_list, dir_alpha=args.alpha, partition="dirichlet" ), "test": get_partition_data( examples=test_examples, num_classes=len(label_list), num_clients=args.clients_num, label_vocab=label_list, dir_alpha=args.alpha, partition="dirichlet" ), "attribute": attribute} logger.info(f"writing clients={args.clients_num}_alpha={args.alpha} ...") partition_data[f"clients={args.clients_num}_alpha={args.alpha}"] = clients_partition_data with open(args.output_partition_file, "wb") as file: pickle.dump(partition_data, file) logger.info("end") if __name__ == "__main__": logger.info("start...") args = parser_args() data_dir = args.data_dir output_dir = args.output_dir tasks = ["MRPC", "SST-2", "QNLI", "QQP", "MNLI", "CoLA", "RTE"] # tasks = ["MRPC"] client_nums = [100, 10] args.overwrite = True for task in tasks: for client_num in client_nums: args.clients_num = client_num args.task = task args.data_dir = os.path.join(data_dir, args.task) args.output_dir = os.path.join(output_dir, "fedglue") make_sure_dirs(args.output_dir) args.output_data_file = os.path.join(args.output_dir, f"{args.task.lower()}_data.pkl") args.output_partition_file = os.path.join(args.output_dir, f"{args.task.lower()}_partition.pkl") logger.info(f"clients_num: {args.clients_num}") logger.info(f"data_dir: {args.data_dir}") logger.info(f"output_dir: {args.output_dir}") logger.info(f"output_data_file: {args.output_data_file}") logger.info(f"output_partition_file: {args.output_partition_file}") convert_glue_to_federated_pkl(args) fedlab-nlp/tools/glue_scripts/glue_metric.py000066400000000000000000000034441427741542300216140ustar00rootroot00000000000000import numpy as np from scipy.stats import pearsonr, spearmanr from sklearn.metrics import f1_score, matthews_corrcoef def simple_accuracy(preds, labels): return (preds == labels).astype(np.int32).mean() def acc_and_f1(preds, labels): acc = simple_accuracy(preds, labels) f1 = f1_score(y_true=labels, y_pred=preds) return { "acc": acc, "f1": f1, "acc_and_f1": (acc + f1) / 2, } def pearson_and_spearman(preds, labels): pearson_corr = pearsonr(preds, labels)[0] spearman_corr = spearmanr(preds, labels)[0] return { "pearson": pearson_corr, "spearmanr": spearman_corr, "corr": (pearson_corr + spearman_corr) / 2, } def glue_compute_metrics(task_name, preds, labels): assert len(preds) == len(labels), f"Predictions and labels have mismatched lengths {len(preds)} and {len(labels)}" if task_name == "cola": return {"mcc": matthews_corrcoef(labels, preds)} elif task_name == "sst-2": return {"acc": simple_accuracy(preds, labels)} elif task_name == "mrpc": return acc_and_f1(preds, labels) elif task_name == "sts-b": return pearson_and_spearman(preds, labels) elif task_name == "qqp": return acc_and_f1(preds, labels) elif task_name == "mnli": return {"acc": simple_accuracy(preds, labels)} elif task_name == "mnli-mm": return {"acc": simple_accuracy(preds, labels)} elif task_name == "qnli": return {"acc": simple_accuracy(preds, labels)} elif task_name == "rte": return {"acc": simple_accuracy(preds, labels)} elif task_name == "wnli": return {"acc": simple_accuracy(preds, labels)} elif task_name == "hans": return {"acc": simple_accuracy(preds, labels)} else: raise KeyError(task_name) fedlab-nlp/tools/glue_scripts/glue_utils.py000066400000000000000000000403001427741542300214610ustar00rootroot00000000000000import os from enum import Enum from loguru import logger from transformers import DataProcessor, InputExample class OutputMode(Enum): classification = "classification" regression = "regression" class MrpcProcessor(DataProcessor): """Processor for the MRPC data set (GLUE version).""" def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) def get_example_from_tensor_dict(self, tensor_dict): """See base class.""" return InputExample( tensor_dict["idx"].numpy(), tensor_dict["sentence1"].numpy().decode("utf-8"), tensor_dict["sentence2"].numpy().decode("utf-8"), str(tensor_dict["label"].numpy()), ) def get_train_examples(self, data_dir): """See base class.""" logger.info(f"LOOKING AT {os.path.join(data_dir, 'train.tsv')}") return self._create_examples(self._read_tsv(os.path.join(data_dir, "train.tsv")), "train") def get_dev_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev") def get_test_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "test.tsv")), "test") def get_labels(self): """See base class.""" return ["0", "1"] def _create_examples(self, lines, set_type): """Creates examples for the training, dev and test sets.""" examples = [] for i, line in enumerate(lines): if i == 0: continue guid = f"{set_type}-{i}" text_a = line[3] text_b = line[4] label = None if set_type == "test" else line[0] examples.append(InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label)) return examples class MnliProcessor(DataProcessor): """Processor for the MultiNLI data set (GLUE version).""" def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) def get_example_from_tensor_dict(self, tensor_dict): """See base class.""" return InputExample( tensor_dict["idx"].numpy(), tensor_dict["premise"].numpy().decode("utf-8"), tensor_dict["hypothesis"].numpy().decode("utf-8"), str(tensor_dict["label"].numpy()), ) def get_train_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "train.tsv")), "train") def get_dev_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "dev_matched.tsv")), "dev_matched") def get_test_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "test_matched.tsv")), "test_matched") def get_labels(self): """See base class.""" return ["contradiction", "entailment", "neutral"] def _create_examples(self, lines, set_type): """Creates examples for the training, dev and test sets.""" examples = [] for i, line in enumerate(lines): if i == 0: continue guid = f"{set_type}-{line[0]}" text_a = line[8] text_b = line[9] label = None if set_type.startswith("test") else line[-1] examples.append(InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label)) return examples class MnliMismatchedProcessor(MnliProcessor): """Processor for the MultiNLI Mismatched data set (GLUE version).""" def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) def get_dev_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "dev_mismatched.tsv")), "dev_mismatched") def get_test_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "test_mismatched.tsv")), "test_mismatched") class ColaProcessor(DataProcessor): """Processor for the CoLA data set (GLUE version).""" def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) def get_example_from_tensor_dict(self, tensor_dict): """See base class.""" return InputExample( tensor_dict["idx"].numpy(), tensor_dict["sentence"].numpy().decode("utf-8"), None, str(tensor_dict["label"].numpy()), ) def get_train_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "train.tsv")), "train") def get_dev_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev") def get_test_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "test.tsv")), "test") def get_labels(self): """See base class.""" return ["0", "1"] def _create_examples(self, lines, set_type): """Creates examples for the training, dev and test sets.""" test_mode = set_type == "test" if test_mode: lines = lines[1:] text_index = 1 if test_mode else 3 examples = [] for i, line in enumerate(lines): guid = f"{set_type}-{i}" text_a = line[text_index] label = None if test_mode else line[1] examples.append(InputExample(guid=guid, text_a=text_a, text_b="", label=label)) return examples class Sst2Processor(DataProcessor): """Processor for the SST-2 data set (GLUE version).""" def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) def get_example_from_tensor_dict(self, tensor_dict): """See base class.""" return InputExample( tensor_dict["idx"].numpy(), tensor_dict["sentence"].numpy().decode("utf-8"), None, str(tensor_dict["label"].numpy()), ) def get_train_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "train.tsv")), "train") def get_dev_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev") def get_test_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "test.tsv")), "test") def get_labels(self): """See base class.""" return ["0", "1"] def _create_examples(self, lines, set_type): """Creates examples for the training, dev and test sets.""" examples = [] text_index = 1 if set_type == "test" else 0 for i, line in enumerate(lines): if i == 0: continue guid = f"{set_type}-{i}" text_a = line[text_index] label = None if set_type == "test" else line[1] examples.append(InputExample(guid=guid, text_a=text_a, text_b="", label=label)) return examples class StsbProcessor(DataProcessor): """Processor for the STS-B data set (GLUE version).""" def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) def get_example_from_tensor_dict(self, tensor_dict): """See base class.""" return InputExample( tensor_dict["idx"].numpy(), tensor_dict["sentence1"].numpy().decode("utf-8"), tensor_dict["sentence2"].numpy().decode("utf-8"), str(tensor_dict["label"].numpy()), ) def get_train_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "train.tsv")), "train") def get_dev_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev") def get_test_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "test.tsv")), "test") def get_labels(self): """See base class.""" return [None] def _create_examples(self, lines, set_type): """Creates examples for the training, dev and test sets.""" examples = [] for i, line in enumerate(lines): if i == 0: continue guid = f"{set_type}-{line[0]}" text_a = line[7] text_b = line[8] label = None if set_type == "test" else line[-1] examples.append(InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label)) return examples class QqpProcessor(DataProcessor): """Processor for the QQP data set (GLUE version).""" def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) def get_example_from_tensor_dict(self, tensor_dict): """See base class.""" return InputExample( tensor_dict["idx"].numpy(), tensor_dict["question1"].numpy().decode("utf-8"), tensor_dict["question2"].numpy().decode("utf-8"), str(tensor_dict["label"].numpy()), ) def get_train_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "train.tsv")), "train") def get_dev_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev") def get_test_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "test.tsv")), "test") def get_labels(self): """See base class.""" return ["0", "1"] def _create_examples(self, lines, set_type): """Creates examples for the training, dev and test sets.""" test_mode = set_type == "test" q1_index = 1 if test_mode else 3 q2_index = 2 if test_mode else 4 examples = [] for i, line in enumerate(lines): if i == 0: continue guid = f"{set_type}-{line[0]}" try: text_a = line[q1_index] text_b = line[q2_index] label = None if test_mode else line[5] except IndexError: continue examples.append(InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label)) return examples class QnliProcessor(DataProcessor): """Processor for the QNLI data set (GLUE version).""" def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) def get_example_from_tensor_dict(self, tensor_dict): """See base class.""" return InputExample( tensor_dict["idx"].numpy(), tensor_dict["question"].numpy().decode("utf-8"), tensor_dict["sentence"].numpy().decode("utf-8"), str(tensor_dict["label"].numpy()), ) def get_train_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "train.tsv")), "train") def get_dev_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev") def get_test_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "test.tsv")), "test") def get_labels(self): """See base class.""" return ["entailment", "not_entailment"] def _create_examples(self, lines, set_type): """Creates examples for the training, dev and test sets.""" examples = [] for i, line in enumerate(lines): if i == 0: continue guid = f"{set_type}-{line[0]}" text_a = line[1] text_b = line[2] label = None if set_type == "test" else line[-1] examples.append(InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label)) return examples class RteProcessor(DataProcessor): """Processor for the RTE data set (GLUE version).""" def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) def get_example_from_tensor_dict(self, tensor_dict): """See base class.""" return InputExample( tensor_dict["idx"].numpy(), tensor_dict["sentence1"].numpy().decode("utf-8"), tensor_dict["sentence2"].numpy().decode("utf-8"), str(tensor_dict["label"].numpy()), ) def get_train_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "train.tsv")), "train") def get_dev_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev") def get_test_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "test.tsv")), "test") def get_labels(self): """See base class.""" return ["entailment", "not_entailment"] def _create_examples(self, lines, set_type): """Creates examples for the training, dev and test sets.""" examples = [] for i, line in enumerate(lines): if i == 0: continue guid = f"{set_type}-{line[0]}" text_a = line[1] text_b = line[2] label = None if set_type == "test" else line[-1] examples.append(InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label)) return examples class WnliProcessor(DataProcessor): """Processor for the WNLI data set (GLUE version).""" def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) def get_example_from_tensor_dict(self, tensor_dict): """See base class.""" return InputExample( tensor_dict["idx"].numpy(), tensor_dict["sentence1"].numpy().decode("utf-8"), tensor_dict["sentence2"].numpy().decode("utf-8"), str(tensor_dict["label"].numpy()), ) def get_train_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "train.tsv")), "train") def get_dev_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev") def get_test_examples(self, data_dir): """See base class.""" return self._create_examples(self._read_tsv(os.path.join(data_dir, "test.tsv")), "test") def get_labels(self): """See base class.""" return ["0", "1"] def _create_examples(self, lines, set_type): """Creates examples for the training, dev and test sets.""" examples = [] for i, line in enumerate(lines): if i == 0: continue guid = f"{set_type}-{line[0]}" text_a = line[1] text_b = line[2] label = None if set_type == "test" else line[-1] examples.append(InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label)) return examples glue_tasks_num_labels = { "cola": 2, "mnli": 3, "mrpc": 2, "sst-2": 2, "sts-b": 1, "qqp": 2, "qnli": 2, "rte": 2, "wnli": 2, } glue_processors = { "cola": ColaProcessor, "mnli": MnliProcessor, "mnli-mm": MnliMismatchedProcessor, "mrpc": MrpcProcessor, "sst-2": Sst2Processor, "sts-b": StsbProcessor, "qqp": QqpProcessor, "qnli": QnliProcessor, "rte": RteProcessor, "wnli": WnliProcessor, } glue_output_modes = { "cola": "classification", "mnli": "classification", "mnli-mm": "classification", "mrpc": "classification", "sst-2": "classification", "sts-b": "regression", "qqp": "classification", "qnli": "classification", "rte": "classification", "wnli": "classification", }fedlab-nlp/tools/glue_scripts/partition.py000066400000000000000000000050351427741542300213240ustar00rootroot00000000000000import os import numpy as np from loguru import logger from fedlab.utils.dataset import DataPartitioner import fedlab.utils.dataset.functional as F from tools.partitions import label_skew_process class GlueDataPartition(DataPartitioner): def __init__(self, targets, num_clients, num_classes, label_vocab, balance=True, partition="iid", unbalance_sgm=0, num_shards=None, dir_alpha=None, verbose=True, seed=42): self.targets = np.array(targets) # with shape (num_samples,) self.num_samples = self.targets.shape[0] self.num_clients = num_clients self.label_vocab = label_vocab self.client_dict = dict() self.partition = partition self.balance = balance self.dir_alpha = dir_alpha self.num_shards = num_shards self.unbalance_sgm = unbalance_sgm self.verbose = verbose self.num_classes = num_classes # self.rng = np.random.default_rng(seed) # rng currently not supports randint np.random.seed(seed) # partition scheme check if balance is None: assert partition in ["dirichlet", "shards"], f"When balance=None, 'partition' only " \ f"accepts 'dirichlet' and 'shards'." elif isinstance(balance, bool): assert partition in ["iid", "dirichlet"], f"When balance is bool, 'partition' only " \ f"accepts 'dirichlet' and 'iid'." else: raise ValueError(f"'balance' can only be NoneType or bool, not {type(balance)}.") # perform partition according to setting self.client_dict = self._perform_partition() # get sample number count for each client # self.client_sample_count = F.samples_num_count(self.client_dict, self.num_clients) def _perform_partition(self): client_dict = label_skew_process( label_vocab=self.label_vocab, label_assignment=self.targets, client_num=self.num_clients, alpha=self.dir_alpha, data_length=len(self.targets) ) return client_dict def __getitem__(self, index): """Obtain sample indices for client ``index``. Args: index (int): BaseClient ID. Returns: list: List of sample indices for client ID ``index``. """ return self.client_dict[index] def __len__(self): """Usually equals to number of clients.""" return len(self.client_dict) fedlab-nlp/tools/glue_scripts/partition.sh000066400000000000000000000004141427741542300213020ustar00rootroot00000000000000#!/usr/bin/env bash run_dir=$1 task=$2 clients_num=$3 alpha=$4 python ${run_dir}/code/fednlp-dev/tools/glue_scripts/glue.py \ --data_dir ${run_dir}/data/glue/ \ --task ${task} \ --output_dir ${run_dir}/data/fedglue/ \ --clients_num ${clients_num} \ --alpha ${alpha} fedlab-nlp/tools/partitions.py000066400000000000000000000217071427741542300170100ustar00rootroot00000000000000import math import numpy as np """ All Non-IID code snippets come from https://github.com/FedML-AI/FedNLP/blob/master/data/advanced_partition/niid_label.py Tanks to FedNLP & FedML https://fedml.ai/ """ def dynamic_batch_fill(label_index_tracker, label_index_matrix, remaining_length, current_label_id): """ params ------------------------------------------------------------------------ label_index_tracker : 1d numpy array track how many data each label has used label_index_matrix : 2d array list of indexs of each label remaining_length : int remaining empty space in current partition client list current_label_id : int current round label id ------------------------------------------------------------------------ return --------------------------------------------------------- label_index_offset: dict dictionary key is label id and value is the offset associated with this key ---------------------------------------------------------- """ remaining_unfiled = remaining_length label_index_offset = {} label_remain_length_dict = {} total_label_remain_length = 0 # calculate total number of all the remaing labels and each label's remaining length for label_id, label_list in enumerate(label_index_matrix): if label_id == current_label_id: label_remain_length_dict[label_id] = 0 continue label_remaining_count = len(label_list) - label_index_tracker[label_id] if label_remaining_count > 0: total_label_remain_length = (total_label_remain_length + label_remaining_count) else: label_remaining_count = 0 label_remain_length_dict[label_id] = label_remaining_count length_pointer = remaining_unfiled if total_label_remain_length > 0: label_sorted_by_length = { k: v for k, v in sorted(label_remain_length_dict.items(), key=lambda item: item[1]) } else: label_index_offset = label_remain_length_dict return label_index_offset # for each label calculate the offset move forward by distribution of remaining labels for label_id in label_sorted_by_length.keys(): fill_count = math.ceil(label_remain_length_dict[label_id] / total_label_remain_length * remaining_length) fill_count = min(fill_count, label_remain_length_dict[label_id]) offset_forward = fill_count # if left room not enough for all offset set it to 0 if length_pointer - offset_forward <= 0 and length_pointer > 0: label_index_offset[label_id] = length_pointer length_pointer = 0 break else: length_pointer -= offset_forward label_remain_length_dict[label_id] -= offset_forward label_index_offset[label_id] = offset_forward # still has some room unfilled if length_pointer > 0: for label_id in label_sorted_by_length.keys(): # make sure no infinite loop happens fill_count = math.ceil(label_sorted_by_length[label_id] / total_label_remain_length * length_pointer) fill_count = min(fill_count, label_remain_length_dict[label_id]) offset_forward = fill_count if length_pointer - offset_forward <= 0 and length_pointer > 0: label_index_offset[label_id] += length_pointer length_pointer = 0 break else: length_pointer -= offset_forward label_remain_length_dict[label_id] -= offset_forward label_index_offset[label_id] += offset_forward return label_index_offset def label_skew_process(label_vocab, label_assignment, client_num, alpha, data_length): """ params ------------------------------------------------------------------- label_vocab : dict label vocabulary of the dataset label_assignment : 1d list a list of label, the index of list is the index associated to label client_num : int number of clients alpha : float similarity of each client, the larger the alpha the similar data for each client ------------------------------------------------------------------- return ------------------------------------------------------------------ partition_result : 2d array list of partition index of each client ------------------------------------------------------------------ """ label_index_matrix = [[] for _ in label_vocab] label_proportion = [] partition_result = [[] for _ in range(client_num)] client_length = 0 print("client_num", client_num) # shuffle indexs and calculate each label proportion of the dataset for index, value in enumerate(label_vocab): label_location = np.where(label_assignment == value)[0] label_proportion.append(len(label_location) / data_length) np.random.shuffle(label_location) label_index_matrix[index].extend(label_location[:]) print(label_proportion) # calculate size for each partition client label_index_tracker = np.zeros(len(label_vocab), dtype=int) total_index = data_length each_client_index_length = int(total_index / client_num) print("each index length", each_client_index_length) client_dir_dis = np.array([alpha * l for l in label_proportion]) print("alpha", alpha) print("client dir dis", client_dir_dis) proportions = np.random.dirichlet(client_dir_dis) print("dir distribution", proportions) # add all the unused data to the client for client_id in range(len(partition_result)): each_client_partition_result = partition_result[client_id] proportions = np.random.dirichlet(client_dir_dis) client_length = min(each_client_index_length, total_index) if total_index < client_length * 2: client_length = total_index total_index -= client_length client_length_pointer = client_length # for each label calculate the offset length assigned to by Dir distribution and then extend assignment for label_id, _ in enumerate(label_vocab): offset = round(proportions[label_id] * client_length) if offset >= client_length_pointer: offset = client_length_pointer client_length_pointer = 0 else: if label_id == (len(label_vocab) - 1): offset = client_length_pointer client_length_pointer -= offset start = int(label_index_tracker[label_id]) end = int(label_index_tracker[label_id] + offset) label_data_length = len(label_index_matrix[label_id]) # if the the label is assigned to a offset length that is more than what its remaining length if end > label_data_length: each_client_partition_result.extend( label_index_matrix[label_id][start:]) label_index_tracker[label_id] = label_data_length label_index_offset = dynamic_batch_fill( label_index_tracker, label_index_matrix, end - label_data_length, label_id) for fill_label_id in label_index_offset.keys(): start = label_index_tracker[fill_label_id] end = (label_index_tracker[fill_label_id] + label_index_offset[fill_label_id]) each_client_partition_result.extend( label_index_matrix[fill_label_id][start:end]) label_index_tracker[fill_label_id] = ( label_index_tracker[fill_label_id] + label_index_offset[fill_label_id]) else: each_client_partition_result.extend( label_index_matrix[label_id][start:end]) label_index_tracker[ label_id] = label_index_tracker[label_id] + offset # if last client still has empty rooms, fill empty rooms with the rest of the unused data if client_id == len(partition_result) - 1: print("last id length", len(each_client_partition_result)) print("Last client fill the rest of the unfilled lables.") for not_fillall_label_id in range(len(label_vocab)): if label_index_tracker[not_fillall_label_id] < len( label_index_matrix[not_fillall_label_id]): print("fill more id", not_fillall_label_id) start = label_index_tracker[not_fillall_label_id] each_client_partition_result.extend( label_index_matrix[not_fillall_label_id][start:]) label_index_tracker[not_fillall_label_id] = len( label_index_matrix[not_fillall_label_id]) partition_result[client_id] = each_client_partition_result return partition_result fedlab-nlp/tools/plot/000077500000000000000000000000001427741542300152115ustar00rootroot00000000000000fedlab-nlp/tools/plot/log_line.py000066400000000000000000000023451427741542300173570ustar00rootroot00000000000000import os import sys import json import pickle import numpy as np import pandas as pd import itertools from itertools import product from scipy.interpolate import make_interp_spline import seaborn as sns import matplotlib.pyplot as plt import matplotlib.ticker as ticker file_path_dict = { "rte": "/workspace/output/fedglue/rte/fedchild_n/20220728032430_bert.eval.log", "qnli": "/workspace/output/fedglue/qnli/fedchild_n/20220727161306_bert.eval.log", "cola": "/workspace/output/fedglue/cola/fedchild_n/20220727161054_bert.eval.log", "mrpc": "/workspace/output/fedglue/mrpc/fedchild_n/20220728032211_bert.eval.log", # "qqp": } task = sys.argv[1] file_path = file_path_dict[task] with open(file_path, "rb") as file: data = pickle.load(file) loss_log, acc_log = [], [] for round_metric in data["logs"]: for rd, metric in round_metric.items(): loss_log.append(round(float(metric["loss"]), 3)) acc_log.append(round(float(metric["acc"][0:4]), 3)) fig = plt.figure(figsize=(10, 6)) my_x_ticks = np.arange(1, len(acc_log) + 1, 1) plt.plot(my_x_ticks, loss_log, label="loss", markersize=8) plt.plot(my_x_ticks, acc_log, label="acc", markersize=8) plt.title(f"Task {task.upper()}") plt.legend() plt.show() fedlab-nlp/trainers/000077500000000000000000000000001427741542300147225ustar00rootroot00000000000000fedlab-nlp/trainers/BaseClient/000077500000000000000000000000001427741542300167335ustar00rootroot00000000000000fedlab-nlp/trainers/BaseClient/__init__.py000066400000000000000000000003621427741542300210450ustar00rootroot00000000000000from trainers.BaseClient.base_client import BaseClientTrainer, BaseClientManager # from trainers.BaseClient.base_local_trainer import BaseLocalTrainer __all__ = [ "BaseClientTrainer", "BaseClientManager", # "BaseLocalTrainer", ] fedlab-nlp/trainers/BaseClient/base_client.py000066400000000000000000000307351427741542300215650ustar00rootroot00000000000000"""BaseClientTrainer for FedLab-NLP""" from abc import ABC import torch from transformers import get_linear_schedule_with_warmup, AdamW from utils import registry from fedlab.utils import MessageCode, SerializationTool from fedlab.core.client.trainer import ClientTrainer from fedlab.core.client.manager import PassiveClientManager from fedlab.core.client.manager import ORDINARY_TRAINER, SERIAL_TRAINER class BaseClientTrainer(ClientTrainer, ABC): def __init__(self, model, train_dataset, valid_dataset): self._model = model self.train_dataset = train_dataset self.valid_dataset = valid_dataset self._before_training() def _before_training(self): """before training function""" # self.client_num = client_num # default is 100. self.type = SERIAL_TRAINER # represent serial trainer config = registry.get("config") self.model_config = config.M self.data_config = config.D self.training_config = config.T self.federated_config = config.F self.client_num = len(config.F.clients_id_list) self.device = config.training_config.device self.rank = config.federated_config.rank self.param_list = [] self.logger = registry.get("logger") self._build_metric() self._build_eval() # key: client idx, value: valid metric self.loc_best_metric = {} # key: client idx, value: test metric self.loc_test_metric = {} # key: client idx, value: serialized params self.loc_best_params = {} self.metric_name = self.metric.metric_name self._model.to(self.device) @property def uplink_package(self): return self.param_list def _train_alone(self, idx, model_parameters, *args, **kwargs): """local training for Client""" train_loader = self._get_dataloader(dataset=self.train_dataset, client_id=idx) if model_parameters is not None: SerializationTool.deserialize_model(self._model, model_parameters) # build optimizer,scheduler,loss optimizer, scheduler = self._build_optimizer(self._model, len(train_loader)) self._model, optimizer = self._mixed_train_model(self._model, optimizer) self._build_loss() for epoch in range(0, int(self.training_config.num_train_epochs)): self._on_epoch_begin() self._on_epoch(train_loader, optimizer, scheduler) self._on_epoch_end(idx) def _get_dataloader(self, dataset, client_id): """Get :class:`DataLoader` for ``client_id``.""" if isinstance(dataset, dict): data_loader = dataset[client_id] else: data_loader = dataset return data_loader def local_process(self, id_list, payload): """local process for Federated Learning""" model_parameters = payload[0] self.param_list = self.fed_train(model_parameters, id_list) return self.param_list def fed_train(self, model_parameters, id_list): param_list = [] for idx in id_list: self._train_alone( idx=idx, model_parameters=model_parameters ) param_list.append(self.model_parameters) return param_list def cen_train(self): self._train_alone( idx=-1, model_parameters=None, ) # Local Training Functions def _build_loss(self): self.criterion = registry.get_loss_class(self.training_config.loss_name)( config=self.training_config ) def _build_optimizer(self, model, train_dl_len): if self.training_config.max_steps > 0: t_total = self.training_config.max_steps self.training_config.num_train_epochs = \ self.training_config.max_steps // (train_dl_len // self.training_config.gradient_accumulation_steps) + 1 else: t_total = \ train_dl_len // self.training_config.gradient_accumulation_steps * self.training_config.num_train_epochs # Prepare optimizer and schedule (linear warmup and decay) no_decay = ['bias', 'LayerNorm.weight'] optimizer_grouped_parameters = [ {'params': [p for n, p in model.backbone.bert.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': self.training_config.weight_decay}, {'params': [p for n, p in model.backbone.bert.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}, ] optimizer = AdamW( optimizer_grouped_parameters, lr=self.training_config.learning_rate, eps=self.training_config.adam_epsilon ) scheduler = get_linear_schedule_with_warmup( optimizer, num_warmup_steps=self.training_config.warmup_steps, num_training_steps=t_total ) return optimizer, scheduler def _mixed_train_model(self, model, optimizer): if self.training_config.fp16: try: from apex import amp except ImportError: raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.") model, optimizer = amp.initialize(model, optimizer, opt_level=self.training_config.fp16_opt_level) # multi-gpu training (should be after apex fp16 initialization) if self.training_config.n_gpu > 1: self.logger.warning("We haven't tested our model under multi-gpu. Please be aware!") model = torch.nn.DataParallel(model) return model, optimizer # Local Test Function def _build_metric(self): self.metric = registry.get_metric_class(self.training_config.metric_name)( self.data_config.task_name, self.training_config.is_decreased_valid_metric ) def _build_eval(self): self.eval = registry.get_eval_class(self.training_config.metric_name)( self.device, self.metric ) def test_on_client(self, test_dataloader): for idx in self.loc_best_params: loc_best_params = self.loc_best_params[idx] SerializationTool.deserialize_model(self._model, loc_best_params) result = self.eval.test_and_eval( model=self._model, valid_dl=test_dataloader, model_type=self.model_config.model_type, model_output_mode=self.model_config.model_output_mode ) test_metric, test_loss = result[self.metric_name], result["eval_loss"] self.logger.critical( f"{self.data_config.task_name.upper()} Test, " f"Client:{idx}, Test loss:{test_loss:.3f}, " f"Test {self.metric_name}:{test_metric:.3f}" ) self.loc_test_metric[idx] = test_metric # Local Epoch Function def _on_epoch_begin(self): self.global_step = 0 self.tr_loss, self.logging_loss = 0.0, 0.0 self.total, self.correct = 0, 0 def _on_epoch(self, train_loader, optimizer, scheduler): for step, batch in enumerate(train_loader): self._model.train() batch = tuple(t.to(self.device) for t in batch) inputs = {'input_ids': batch[0], 'attention_mask': batch[1], 'labels': batch[3] } label = inputs['labels'] if self.model_config.model_type != 'distilbert': # XLM, DistilBERT and RoBERTa don't use segment_ids inputs['token_type_ids'] = batch[2] \ if self.model_config.model_type in ['bert', 'xlnet'] else None outputs = self._model(inputs) loss, logits = outputs[:2] _, predicted = torch.max(logits, 1) optimizer.zero_grad() if self.training_config.n_gpu > 1: loss = loss.mean() # mean() to average on multi-gpu parallel training if self.training_config.gradient_accumulation_steps > 1: loss = loss / self.training_config.gradient_accumulation_steps if self.training_config.fp16: try: from apex import amp except ImportError: raise ImportError( "Please install apex from https://www.github.com/nvidia/apex to use fp16 training.") with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() else: loss.backward() self.tr_loss += loss.item() if (step + 1) % self.training_config.gradient_accumulation_steps == 0: if self.training_config.fp16: torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), self.training_config.max_grad_norm) else: torch.nn.utils.clip_grad_norm_(self._model.parameters(), self.training_config.max_grad_norm) optimizer.step() scheduler.step() # Update learning rate schedule self.global_step += 1 self.total += label.size(0) self.correct += (predicted == label).sum().item() def _on_epoch_end(self, idx): """on epoch end""" self.logger.info(f"{self.data_config.task_name.upper()} Train, " f"client:{idx}, loss:{self.tr_loss/self.global_step:.3f}, " f"accuracy:{self.correct/self.total:.3f}") if not self.federated_config.pson: # not need for local test return valid_data = self._get_dataloader(dataset=self.valid_dataset, client_id=idx) result = self.eval.test_and_eval( model=self._model, valid_dl=valid_data, model_type=self.model_config.model_type, model_output_mode=self.model_config.model_output_mode ) test_metric, test_loss = result[self.metric_name], result["eval_loss"] # TODO hard code if not self.loc_best_metric.get(idx, None): self.loc_best_metric[idx] = 0.0 if self.loc_best_metric[idx] < test_metric: self.loc_best_metric[idx] = test_metric self.loc_best_params[idx] = SerializationTool.serialize_model(self._model) self.logger.debug(f"{self.data_config.task_name.upper()} Eval, " f"Client:{idx}, Current loss:{test_loss:.3f}, " f"Current {self.metric_name}:{test_metric:.3f}, " f"Best {self.metric_name}:{self.loc_best_metric[idx]:.3f}") class BaseClientManager(PassiveClientManager, ABC): def __init__(self, network, trainer): self.logger = registry.get("logger") super().__init__(network, trainer, self.logger) def main_loop(self): """Actions to perform when receiving a new message, including local trainers. Main procedure of each client: 1. client waits for data from server (PASSIVELY). 2. after receiving data, client start local model trainers procedure. 3. client synchronizes with server actively. """ while True: sender_rank, message_code, payload = self._network.recv(src=0) if message_code == MessageCode.Exit: # client exit feedback if self._network.rank == self._network.world_size - 1: self._network.send(message_code=MessageCode.Exit, dst=0) break elif message_code == MessageCode.ParameterUpdate: id_list, payload = payload[0].to( torch.int32).tolist(), payload[1:] # check the trainer type if self._trainer.type == SERIAL_TRAINER: # serial self._trainer.local_process( id_list=id_list, payload=payload ) elif self._trainer.type == ORDINARY_TRAINER: # ordinary assert len(id_list) == 1 self._trainer.local_process(payload=payload) self.synchronize() else: raise ValueError(f"Invalid MessageCode {message_code}. Please check MessageCode list.") def synchronize(self): """Synchronize with server""" self.logger.info("Uploading information to server.") self._network.send( content=self._trainer.uplink_package, message_code=MessageCode.ParameterUpdate, dst=0 ) fedlab-nlp/trainers/BaseClient/base_local_trainer.py000066400000000000000000000043521427741542300231210ustar00rootroot00000000000000# import torch # from abc import ABC # from utils import registry # # Discard LocalTrainer # # class BaseLocalTrainer(ABC): # def __init__(self): # # config = registry.get("config") # self.model_config = config.model_config # self.data_config = config.data_config # self.training_config = config.training_config # self.federated_config = config.federated_config # # self.device = config.training_config.device # self.metric = registry.get_metric_class(self.training_config.metric_name) # # self.logger = registry.get("logger") # # self._before_training() # # def train_model(self, *args): # raise NotImplementedError # # def eval_model(self, model, valid_dl): # raise NotImplementedError # # def infer_model(self, model, test_dl): # raise NotImplementedError # # def _build_loss(self): # return registry.get_loss_class(self.training_config.loss_name)( # config=self.training_config # ) # # def _build_metric(self): # self.metric = registry.get_metric_class(self.training_config.metric_name)( # self.data_config.task_name, self.training_config.is_decreased_valid_metric # ) # # def _build_optimizer(self, model, iteration_in_total): # raise NotImplementedError # # def _mixed_train_model(self, model, optimizer): # if self.training_config.fp16: # try: # from apex import amp # except ImportError: # raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.") # model, optimizer = amp.initialize(model, optimizer, opt_level=self.training_config.fp16_opt_level) # # # multi-gpu training (should be after apex fp16 initialization) # if self.training_config.n_gpu > 1: # self.logger.warning("We haven't tested our model under multi-gpu. Please be aware!") # model = torch.nn.DataParallel(model) # # return model, optimizer # # def _freeze_model_parameters(self, model): # raise NotImplementedError # # def _before_training(self): # self.logger.info("build metric ...") # self._build_metric() fedlab-nlp/trainers/BaseServer/000077500000000000000000000000001427741542300167635ustar00rootroot00000000000000fedlab-nlp/trainers/BaseServer/__init__.py000066400000000000000000000002311427741542300210700ustar00rootroot00000000000000from trainers.BaseServer.base_server import BaseSyncServerHandler, BaseServerManager __all__ = [ "BaseSyncServerHandler", "BaseServerManager" ] fedlab-nlp/trainers/BaseServer/base_server.py000066400000000000000000000243141427741542300216410ustar00rootroot00000000000000""" BaseServer for FedLab-NLP """ import os import random import threading from abc import ABC import torch from utils.register import registry from fedlab.core.server.handler import Aggregators from fedlab.core.server.handler import ParameterServerBackendHandler from fedlab.core.server.manager import ServerManager from fedlab.utils.serialization import SerializationTool from fedlab.utils import MessageCode from fedlab.core.coordinator import Coordinator class BaseSyncServerHandler(ParameterServerBackendHandler, ABC): def __init__(self, model, valid_data, test_data): self.valid_data = valid_data self.test_data = test_data config = registry.get("config") self.model_config = config.model_config self.data_config = config.data_config self.training_config = config.training_config self.federated_config = config.federated_config self.logger = registry.get("logger") self.device = config.training_config.device self._model = model.to(self.device) # basic setting self.client_num_in_total = config.federated_config.clients_num self.sample_ratio = config.federated_config.sample # client buffer self.client_buffer_cache = [] self.cache_cnt = 0 # stop condition self.global_round = config.federated_config.rounds self.round = 0 # metrics & eval self._build_metric() self._build_eval() self.global_valid_best_metric = \ float("inf") if self.training_config.is_decreased_valid_metric else -float("inf") self.global_test_best_metric = 0.0 self.metric_log = { "model_type": self.model_config.model_type, "clients_num": self.federated_config.clients_num, "alpha": self.federated_config.alpha, "task": self.data_config.task_name, "fl_algorithm": self.federated_config.fl_algorithm, "info": f"{self.model_config.model_type}_{self.federated_config.fl_algorithm}_" f"{self.federated_config.clients_num}_{self.federated_config.alpha}", "logs": [] } # metric line self.metric_name = self.metric.metric_name times = registry.get("run_time") self.metric_line = f"{times}_{self.model_config.model_type}_client={self.federated_config.clients_num}_" \ f"alpha={self.federated_config.alpha}_ci={self.federated_config.sample}_" # global model self.glo_save_file = os.path.join( self.training_config.checkpoint_dir, f"{times}_{self.model_config.model_type}.pth") self.best_glo_params = None def _build_eval(self): self.eval = registry.get_eval_class(self.training_config.metric_name)( self.device, self.metric ) def _build_metric(self): self.metric = registry.get_metric_class(self.training_config.metric_name)( self.data_config.task_name, self.training_config.is_decreased_valid_metric ) def stop_condition(self) -> bool: return self.round >= self.global_round def sample_clients(self): selection = random.sample( range(self.client_num_in_total), self.client_num_per_round ) return selection def _update_global_model(self, payload): assert len(payload) > 0 if len(payload) == 1: self.client_buffer_cache.append(payload[0].clone()) else: self.client_buffer_cache += payload # serial trainer assert len(self.client_buffer_cache) <= self.client_num_per_round if len(self.client_buffer_cache) == self.client_num_per_round: model_parameters_list = self.client_buffer_cache self.logger.debug( f"Model parameters aggregation, number of aggregation elements {len(model_parameters_list)}" ) # use aggregator serialized_parameters = Aggregators.fedavg_aggregate(model_parameters_list) SerializationTool.deserialize_model(self._model, serialized_parameters) self.round += 1 self.valid_on_server() # reset cache cnt self.client_buffer_cache = [] return True # return True to end this round. else: return False @property def client_num_per_round(self): return max(1, int(self.sample_ratio * self.client_num_in_total)) @property def downlink_package(self): """Property for manager layer. BaseServer manager will call this property when activates clients.""" return [self.model_parameters] @property def if_stop(self): """ class:`NetworkManager` keeps monitoring this attribute, and it will stop all related processes and threads when ``True`` returned. """ return self.round >= self.global_round def valid_on_server(self): result = self.eval.test_and_eval( model=self._model, valid_dl=self.valid_data, model_type=self.model_config.model_type, model_output_mode=self.model_config.model_output_mode ) self.on_round_end(result) def test_on_server(self): SerializationTool.deserialize_model(self._model, self.best_glo_params) result = self.eval.test_and_eval( model=self._model, valid_dl=self.test_data, model_type=self.model_config.model_type, model_output_mode=self.model_config.model_output_mode ) self.logger.critical(f"task: {self.data_config.task_name}, Setting: {self.metric_log['info']}, " f"Test {self.metric_name.upper()}: {result[self.metric_name]:.3f}") self.global_test_best_metric = result[self.metric_name] return result def save_global_model(self): torch.save(self.best_glo_params, self.glo_save_file) def on_round_end(self, result): test_metric, test_loss = result[self.metric_name], result["eval_loss"] # TODO hard code if self.global_valid_best_metric < test_metric: self.global_valid_best_metric = test_metric self.best_glo_params = SerializationTool.serialize_model(self._model) self.logger.info(f"{self.data_config.task_name}-{self.model_config.model_type} " f"train with client={self.federated_config.clients_num}_" f"alpha={self.federated_config.alpha}_" f"epoch={self.training_config.num_train_epochs}_" f"seed={self.training_config.seed}_" f"comm_round={self.federated_config.rounds}") self.logger.debug(f"{self.federated_config.fl_algorithm} Validating " f"Round: {self.round}, Current Loss: {test_loss:.3f}, " f"Current {self.metric_name}: {test_metric:.3f}, " f"Best {self.metric_name}: {self.global_valid_best_metric:.3f}") self.metric_log["logs"].append( {f"round_{self.round}": { "loss": f"{test_loss:.3f}", f"{self.metric.metric_name}": f"{test_metric:.3f}" } } ) class BaseServerManager(ServerManager): """Synchronous communication BaseServerManager.run() setup() main_loop() shut_down() """ def __init__(self, network, handler): super(BaseServerManager, self).__init__(network, handler) self.logger = registry.get("logger") def setup(self): self._network.init_network_connection() rank_client_id_map = {} for rank in range(1, self._network.world_size): _, _, content = self._network.recv(src=rank) rank_client_id_map[rank] = content[0].item() self.coordinator = Coordinator(rank_client_id_map, mode='GLOBAL') # mode='GLOBAL' if self._handler is not None: self._handler.client_num_in_total = self.coordinator.total def main_loop(self): while self._handler.if_stop is not True: activate = threading.Thread(target=self.activate_clients) activate.start() while True: sender_rank, message_code, payload = self._network.recv() if message_code == MessageCode.ParameterUpdate: if self._handler._update_global_model(payload): break else: raise Exception( "Unexpected message code {}".format(message_code)) def shutdown(self): """Shutdown stage.""" self.shutdown_clients() super().shutdown() def activate_clients(self): self.logger.info("BaseClient activation procedure") clients_this_round = self._handler.sample_clients() rank_dict = self.coordinator.map_id_list(clients_this_round) self.logger.info("BaseClient id list: {}".format(clients_this_round)) for rank, values in rank_dict.items(): downlink_package = self._handler.downlink_package id_list = torch.Tensor(values).to(downlink_package[0].dtype) self._network.send( content=[id_list] + downlink_package, message_code=MessageCode.ParameterUpdate, dst=rank ) def shutdown_clients(self): """Shutdown all clients. Send package to each client with :attr:`MessageCode.Exit`. Note: Communication agreements related: User can overwrite this function to define package for exiting information. """ client_list = range(self._handler.client_num_in_total) rank_dict = self.coordinator.map_id_list(client_list) for rank, values in rank_dict.items(): downlink_package = self._handler.downlink_package id_list = torch.Tensor(values).to(downlink_package[0].dtype) self._network.send(content=[id_list] + downlink_package, message_code=MessageCode.Exit, dst=rank) # wait for client exit feedback _, message_code, _ = self._network.recv( src=self._network.world_size - 1 ) assert message_code == MessageCode.Exit fedlab-nlp/trainers/FedBaseTrainer.py000066400000000000000000000066471427741542300201270ustar00rootroot00000000000000"""BaseTrainer for FedLab-NLP""" from abc import ABC from utils import registry from utils import setup_seed from utils import global_metric_save from fedlab.core.network import DistNetwork class BaseTrainer(ABC): def __init__(self, *args): config = registry.get("config") self.model_config = config.M self.data_config = config.D self.training_config = config.T self.federated_config = config.F self.logger = registry.get("logger") # self._before_training() @property def role(self): if self.federated_config.rank == 0: return "server" elif self.federated_config.rank > 0: return f"sub-server_{self.federated_config.rank}" else: return "centralized" def _build_server(self): raise NotImplementedError def _build_client(self): raise NotImplementedError def _build_local_trainer(self, *args): raise NotImplementedError def _build_network(self): self.network = DistNetwork( address=(self.federated_config.ip, self.federated_config.port), world_size=self.federated_config.world_size, rank=self.federated_config.rank, ethernet=self.federated_config.ethernet) def _build_data(self): self.data = registry.get_data_class(self.data_config.dataset_name)() def _build_model(self): self.model = registry.get_model_class(self.model_config.model_output_mode)( task_name=self.data_config.task_name ) def _before_training(self): setup_seed(self.training_config.seed) self.logger.info(f"{self.role} building dataset ...") # set before build model self._build_data() self.logger.info(f"{self.role} building model ...") self._build_model() # self.logger.info(f"{self.role} building local trainer ...") # self._build_local_trainer() if self.federated_config.rank != -1: self.logger.info(f"{self.role} building network ...") self._build_network() if self.federated_config.rank == 0: self.logger.info("building server ...") self._build_server() else: self._build_client() if self.federated_config.rank > 0: self.logger.info(f"building client {self.federated_config.rank} ...") self.logger.info(f"local rank {self.federated_config.rank}'s client ids " f"is {list(self.data.train_dataloader_dict.keys())}") else: self.logger.info("building centralized training") def train(self): # TODO phase decides train or test if self.federated_config.rank == 0: self.logger.debug(f"Server Start ...") self.server_manger.run() self.on_server_end() elif self.federated_config.rank > 0: self.logger.debug(f"Sub-Server {self.federated_config.rank} Training Start ...") self.client_manager.run() self.on_client_end() else: self.logger.debug(f"Centralized Training Start ...") self.client_trainer.cen_train() self.on_client_end() def on_server_end(self): """on_server_end""" self.handler.test_on_server() global_metric_save(self.handler, self.training_config, self.logger) def on_client_end(self): ... fedlab-nlp/trainers/__init__.py000066400000000000000000000004601427741542300170330ustar00rootroot00000000000000"""FedLab-NLP's trainers registry in trainer.__init__.py -- IMPORTANT!""" from trainers.FedBaseTrainer import BaseTrainer from run.fedavg.trainer import FedAvgTrainer from run.centralized.trainer import CentralizedTrainer __all__ = [ "BaseTrainer", "FedAvgTrainer", "CentralizedTrainer" ] fedlab-nlp/utils/000077500000000000000000000000001427741542300142335ustar00rootroot00000000000000fedlab-nlp/utils/__init__.py000066400000000000000000000016631427741542300163520ustar00rootroot00000000000000""" Import core names of utils """ from utils.general import ( get_cpus, init_training_device, make_sure_dirs, get_parameter_number, get_memory_usage, rm_dirs, rm_file, setup_imports, setup_seed, pickle_read, pickle_write, file_write, check_cached_data, global_metric_save, cen_metric_save ) from utils.logger import setup_logger from utils.loss import Loss from utils.config import build_config from utils.register import registry # from utils.transform import ss_tokenize, ms_tokenize __all__ = [ "get_cpus", "init_training_device", "get_parameter_number", "get_memory_usage", "setup_logger", "setup_imports", "setup_seed", "rm_dirs", "make_sure_dirs", "registry", "Loss", "build_config", "pickle_read", "pickle_write", "rm_file", "file_write", "check_cached_data", "global_metric_save", "cen_metric_save" ] fedlab-nlp/utils/callback.py000066400000000000000000000000001427741542300163270ustar00rootroot00000000000000fedlab-nlp/utils/config.py000066400000000000000000000110341427741542300160510ustar00rootroot00000000000000"""config for FedLab-NLP""" import os import time from abc import ABC from omegaconf import OmegaConf from transformers import HfArgumentParser from utils import make_sure_dirs, rm_file from utils.register import registry from configs import ModelArguments, DataTrainingArguments, TrainArguments, FederatedTrainingArguments class Config(ABC): def __init__(self, model_args, data_args, training_args, federated_args): self.model_config = model_args self.data_config = data_args self.training_config = training_args self.federated_config = federated_args def save_configs(self): ... def check_config(self): self.config_check_federated() def config_check_federated(self): if self.F.world_size is None: if self.F.rank == -1: self.F.world_size = 1 else: raise ValueError(f"Must set world_size, but find {self.F.world_size}") else: if self.F.clients_num % (self.F.world_size - 1): raise ValueError(f"{self.F.clients_num} % {(self.F.world_size - 1)} != 0") @property def M(self): return self.model_config @property def D(self): return self.data_config @property def T(self): return self.training_config @property def F(self): return self.federated_config def amend_config(model_args, data_args, training_args, federated_args): config = Config(model_args, data_args, training_args, federated_args) # load customer config (hard code) # TODO args in config.yaml can overwrite --arg root_folder = registry.get("root_folder") cust_config_path = os.path.join(root_folder, f"run/{config.F.fl_algorithm}/config.yaml") if os.path.isfile(cust_config_path): cust_config = OmegaConf.load(cust_config_path) for key, values in cust_config.items(): if values: args = getattr(config, key) for k, v in values.items(): setattr(args, k, v) # set training path config.T.output_dir = os.path.join(config.T.output_dir, config.D.task_name) make_sure_dirs(config.T.output_dir) if not config.D.cache_dir: cache_dir = os.path.join(config.T.output_dir, "cached_data") if config.F.rank != -1: config.D.cache_dir = os.path.join( cache_dir, f"cached_{config.M.model_type}_{config.F.clients_num}_{config.F.alpha}" ) else: config.D.cache_dir = os.path.join( cache_dir, f"cached_{config.M.model_type}_centralized" ) make_sure_dirs(config.D.cache_dir) # set training_args config.T.save_dir = os.path.join(config.T.output_dir, config.F.fl_algorithm.lower()) make_sure_dirs(config.T.save_dir) config.T.checkpoint_dir = os.path.join(config.T.save_dir, "saved_model") make_sure_dirs(config.T.checkpoint_dir) # set phase phase = "train" if config.T.do_train else "evaluate" registry.register("phase", phase) # set metric log path times = time.strftime("%Y%m%d%H%M%S", time.localtime()) registry.register("run_time", times) config.T.times = times config.T.metric_file = os.path.join(config.T.save_dir, f"{config.M.model_type}.eval") config.T.metric_log_file = os.path.join(config.T.save_dir, f"{times}_{config.M.model_type}.eval.log") # set federated_args if config.F.do_mimic and config.F.rank == 0: server_write_flag_path = os.path.join(config.D.cache_dir, "server_write.flag") rm_file(server_write_flag_path) if config.F.partition_method is None: config.F.partition_method = f"clients={config.F.clients_num}_alpha={config.F.alpha}" config.check_config() registry.register("config", config) return config def build_config(): # read parameters parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainArguments, FederatedTrainingArguments)) model_args, data_args, training_args, federated_args = parser.parse_args_into_dataclasses() # amend and register configs config = amend_config(model_args, data_args, training_args, federated_args) # logging fl & some path logger = registry.get("logger") logger.info(f"FL-Algorithm: {config.federated_config.fl_algorithm}") logger.info(f"output_dir: {config.training_config.output_dir}") logger.info(f"cache_dir: {config.data_config.cache_dir}") logger.info(f"save_dir: {config.training_config.save_dir}") logger.info(f"checkpoint_dir: {config.training_config.checkpoint_dir}") return config fedlab-nlp/utils/evaluations.py000066400000000000000000000044301427741542300171400ustar00rootroot00000000000000"""define evaluation scripts for FedLab-NLP """ import torch import numpy as np from abc import ABC from utils.register import registry class BaseEval(ABC): def __init__(self, device, metric): self.device = device self.metric = metric self.task_name = metric.task_name self.logger = registry.get("logger") def test_and_eval(self, valid_dl, model, *args): raise NotImplementedError @registry.register_eval("glue") class GlueEval(BaseEval, ABC): def __init__(self, device, metric): super(GlueEval, self).__init__(device, metric) def test_and_eval(self, valid_dl, model, model_type, model_output_mode): model.to(self.device) eval_loss, nb_eval_steps = 0.0, 0 preds, out_label_ids = None, None results = {} for batch in valid_dl: model.eval() batch = tuple(t.to(self.device) for t in batch) with torch.no_grad(): inputs = {'input_ids': batch[0], 'attention_mask': batch[1], 'labels': batch[3] } if model_type != 'distilbert': # XLM, DistilBERT and RoBERTa don't use segment_ids inputs['token_type_ids'] = \ batch[2] if model_type in ['bert', 'xlnet'] else None outputs = model(inputs) tmp_eval_loss, logits = outputs[:2] eval_loss += tmp_eval_loss.mean().item() nb_eval_steps += 1 if preds is None: preds = logits.detach().cpu().numpy() out_label_ids = inputs['labels'].detach().cpu().numpy() else: preds = np.append(preds, logits.detach().cpu().numpy(), axis=0) out_label_ids = np.append(out_label_ids, inputs['labels'].detach().cpu().numpy(), axis=0) eval_loss = eval_loss / nb_eval_steps results["eval_loss"] = eval_loss if model_output_mode == "seq_classification": preds = np.argmax(preds, axis=1) elif model_output_mode == "regression": preds = np.squeeze(preds) result = self.metric.calculate_metric(preds, out_label_ids, False) results.update(result) return results fedlab-nlp/utils/general.py000066400000000000000000000121651427741542300162270ustar00rootroot00000000000000""" general functions for fedlab-nlp """ import os import pickle import importlib import math import multiprocessing import random import shutil from glob import glob import numpy as np import psutil import torch def pickle_read(path, read_format="rb"): with open(path, read_format) as file: obj = pickle.load(file) return obj def pickle_write(obj, path, write_format="wb"): with open(path, write_format) as file: pickle.dump(obj, file) def file_write(line, path, mode): with open(path, mode) as file: file.write(line+"\n") def make_sure_dirs(path: str): """Create dir if not exists Args: path (str): path """ if not os.path.exists(path): os.makedirs(path) return path def rm_dirs(path: str): """ remove file existing check. Args: path (str): path """ if os.path.exists(path): shutil.rmtree(path) def rm_file(file_path: str): if os.path.isfile(file_path): os.unlink(file_path) def get_cpus(): """return total num of cpus in current machine.""" try: with open("/sys/fs/cgroup/cpu/cpu.cfs_quota_us") as f: cfs_quota_us = int(f.readline()) with open("/sys/fs/cgroup/cpu/cpu.cfs_period_us") as f: cfs_period_us = int(f.readline()) if cfs_quota_us > 0 and cfs_period_us > 0: return int(math.ceil(cfs_quota_us / cfs_period_us)) except Exception: pass return multiprocessing.cpu_count() def get_memory_usage(): """ return total memory been used. memory use in GB """ pid = os.getpid() py = psutil.Process(pid) memory_use = py.memory_info()[0] / 2.0 ** 30 return memory_use def get_parameter_number(net): total_num = sum(p.numel() for p in net.parameters()) trainable_num = sum(p.numel() for p in net.parameters() if p.requires_grad) return {'Total': total_num / 1e6, 'Trainable': trainable_num / 1e6} def check_cached_data(path): data_processed_flag_file = os.path.join(path, "server_write.flag") if not os.path.isfile(data_processed_flag_file): return False with open(data_processed_flag_file) as file: lines = [] for line in file: lines.append(line) if not lines: return False return True def setup_seed(seed: int): os.environ["PYTHONHASHSEED"] = str(seed) np.random.seed(seed) random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.enabled = True # torch.backends.cudnn.benchmark = False # torch.backends.cudnn.deterministic = True def init_training_device(gpu): device = torch.device(f"cuda:{gpu}" if torch.cuda.is_available() else "cpu") return device def global_metric_save(handler, training_config, logger): pickle_write(handler.metric_log, training_config.metric_log_file) handler.metric_line += f"valid_{handler.metric_name}={handler.global_valid_best_metric:.3f}_" handler.metric_line += f"test_{handler.metric_name}={handler.global_test_best_metric:.3f}" file_write(handler.metric_line, training_config.metric_file, "a+") logger.info(f"watch training logs --> {training_config.metric_log_file}") logger.info(f"training metric --> {training_config.metric_file}") def cen_metric_save(loc_trainer, training_config, logger): test_metric = loc_trainer.loc_test_metric[-1] valid_metric = loc_trainer.loc_best_metric[-1] times = training_config.times line = f"{times}_valid_{loc_trainer.metric_name}={valid_metric:.3f}_" line += f"test_{loc_trainer.metric_name}={test_metric:.3f}" file_write(line, training_config.metric_file, "a+") logger.info(f"training metric --> {training_config.metric_file}") def setup_imports(): from utils.register import registry # First, check if imports are already setup has_already_setup = registry.get("imports_setup", no_warning=True) if has_already_setup: return # Automatically load all of the modules, so that they register with registry root_folder = os.path.dirname(os.path.abspath(__file__)) project_name = root_folder.split(os.sep)[-2] root_folder = os.path.join(root_folder, "..") # check here files = [] for package_name in ["trainers", "data", "models", "utils", "configs"]: folder = os.path.join(root_folder, package_name) pattern = os.path.join(folder, "**", "*.py") files.extend(glob(pattern, recursive=True)) for f in files: f = os.path.realpath(f) if f.endswith(".py") and not f.endswith("__init__.py"): splits = f.split(os.sep) import_prefix_index = 0 for idx, split in enumerate(splits): if split == project_name: import_prefix_index = idx + 1 file_name = splits[-1] module_name = file_name[: file_name.find(".py")] module = ".".join( splits[import_prefix_index:-1] + [module_name] ) importlib.import_module(module) registry.register("root_folder", root_folder) registry.register("imports_setup", True) fedlab-nlp/utils/globalhost.py000066400000000000000000000000001427741542300167310ustar00rootroot00000000000000fedlab-nlp/utils/logger.py000066400000000000000000000013151427741542300160640ustar00rootroot00000000000000""" pretty logging for fedlab-nlp """ import sys from loguru import logger from utils.register import registry def formatter(record): # default format time_format = "{time:MM-DD/HH:mm:ss}" lvl_format = "{level:^5}" rcd_format = "{file}:{line:}" msg_format = "{message}" if record["level"].name in ["WARNING", "CRITICAL"]: lvl_format = "" + lvl_format + "" return "|".join([time_format, lvl_format, rcd_format, msg_format]) + "\n" def setup_logger(): logger.remove() logger.add( sys.stderr, format=formatter, colorize=True, enqueue=True ) registry.register("logger", logger) return logger fedlab-nlp/utils/loss.py000066400000000000000000000206301427741542300155660ustar00rootroot00000000000000""" loss for fedlab-nlp """ from typing import Dict, List import torch import torch.nn as nn import torch.nn.functional as F from loguru import logger from torch.autograd import Variable from utils.register import registry class BaseLoss(nn.Module): def __init__(self, config): super().__init__() self.config = config self.logger = registry.get("logger") def forward(self, x, target): raise NotImplementedError @registry.register_loss("focal") class FocalLoss(BaseLoss): """FocalLoss.""" def __init__(self, config): super().__init__(config) self.gamma = config.gamma self.alpha = config.alpha if isinstance(self.alpha, (float, int)): self.alpha = torch.Tensor([self.alpha, 1 - self.alpha]) if isinstance(self.alpha, list): self.alpha = torch.Tensor(self.alpha) self.size_average = config.size_average def forward(self, x, target): target = target.squeeze().long() if x.dim() > 2: x = x.view(x.size(0), x.size(1), -1) # N,C,H,W => N,C,H*W x = x.transpose(1, 2) # N,C,H*W => N,H*W,C x = x.contiguous().view(-1, x.size(2)) # N,H*W,C => N*H*W,C target = target.view(-1, 1) logpt = F.log_softmax(x) logpt = logpt.gather(1, target) logpt = logpt.view(-1) pt = Variable(logpt.data.exp()) if self.alpha is not None: if self.alpha.type() != x.data.type(): self.alpha = self.alpha.type_as(x.data) at = self.alpha.gather(0, target.data.view(-1)) logpt = logpt * Variable(at) loss = -1 * (1 - pt) ** self.gamma * logpt if self.size_average: return loss.mean() else: return loss.sum() @registry.register_loss("label_smooth") class LabelSmoothLoss(BaseLoss): """Label Smoothing FocalLoss.""" def __init__(self, config): super().__init__(config) self.gamma = config.get("gamma", 2) self.smooth_eps = config.get("smooth_eps", 0.25) self.size_average = config.get("size_average", True) def forward(self, x, target): cls_num = x.size(1) target = target.squeeze().long() if x.dim() > 2: x = x.view(x.size(0), x.size(1), -1) # N,C,H,W => N,C,H*W x = x.transpose(1, 2) # N,C,H*W => N,H*W,C x = x.contiguous().view(-1, x.size(2)) # N,H*W,C => N*H*W,C target = target.view(-1, 1) pred = F.log_softmax(x) logpt = pred.gather(1, target) logpt = logpt.view(-1) pt = Variable(logpt.data.exp()) focal_loss = -1 * (1 - pt) ** self.gamma * logpt lce = -1 * torch.sum(pred, dim=1) / cls_num loss = (1 - self.smooth_eps) * focal_loss + self.smooth_eps * lce if self.size_average: return loss.mean() else: return loss.sum() @registry.register_loss("kl") class DistillationLoss(BaseLoss): """DistillationLoss.""" def __init__(self, config): super().__init__(config) self.temp = self.config.temp # self.reduction_kd = self.configs.get("reduction_kd", "batchmean") self.criterion = nn.KLDivLoss(reduction="batchmean") def forward(self, x, y): student = F.log_softmax(x / self.temp, dim=1) teacher = F.softmax(y / self.temp, dim=1) klloss = self.criterion(student, teacher) * self.temp * self.temp return klloss def normalize(x: torch.Tensor, axis: int = -1) -> torch.Tensor: """Normalizing to unit length along the specified dimension. Args: x (torch.Tensor): x axis (int): axis Returns: torch.Tensor: """ x = 1.0 * x / (torch.norm(x, 2, axis, keepdim=True).expand_as(x) + 1e-12) return x def euclidean_dist(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: """euclidean_dist. Args: x (torch.Tensor): in shape of [m, d] y (torch.Tensor): in shape of [n, d] Returns: torch.Tensor: in shape of [m, n] """ m, n = x.size(0), y.size(0) xx = torch.pow(x, 2).sum(1, keepdim=True).expand(m, n) yy = torch.pow(y, 2).sum(1, keepdim=True).expand(n, m).t() dist = xx + yy dist.addmm_(1, -2, x, y.t()) dist = dist.clamp(min=1e-12).sqrt() # for numerical stability return dist def hard_example_mining(dist_mat, labels, return_inds=False): """For each anchor, find the hardest positive and negative sample. Args: dist_mat: pytorch Variable, pair wise distance between samples, shape [N, N] labels: pytorch LongTensor, with shape [N] return_inds: whether to return the indices. Save time if `False`(?) Returns: dist_ap: pytorch Variable, distance(anchor, positive); shape [N] dist_an: pytorch Variable, distance(anchor, negative); shape [N] p_inds: pytorch LongTensor, with shape [N]; indices of selected hard positive samples; 0 <= p_inds[i] <= N - 1 n_inds: pytorch LongTensor, with shape [N]; indices of selected hard negative samples; 0 <= n_inds[i] <= N - 1 NOTE: Only consider the case in which all labels have same num of samples, thus we can cope with all anchors in parallel. """ assert len(dist_mat.size()) == 2 assert dist_mat.size(0) == dist_mat.size(1) N = dist_mat.size(0) # shape [N, N] is_pos = labels.expand(N, N).eq(labels.expand(N, N).t()) is_neg = labels.expand(N, N).ne(labels.expand(N, N).t()) # `dist_ap` means distance(anchor, positive) # both `dist_ap` and `relative_p_inds` with shape [N, 1] dist_ap, relative_p_inds = torch.max( dist_mat[is_pos].contiguous().view(N, -1), 1, keepdim=True ) # `dist_an` means distance(anchor, negative) # both `dist_an` and `relative_n_inds` with shape [N, 1] dist_an, relative_n_inds = torch.min( dist_mat[is_neg].contiguous().view(N, -1), 1, keepdim=True ) # shape [N] dist_ap = dist_ap.squeeze(1) dist_an = dist_an.squeeze(1) if return_inds: # shape [N, N] ind = ( labels.new() .resize_as_(labels) .copy_(torch.arange(0, N).long()) .unsqueeze(0) .expand(N, N) ) # shape [N, 1] p_inds = torch.gather( ind[is_pos].contiguous().view(N, -1), 1, relative_p_inds.data ) n_inds = torch.gather( ind[is_neg].contiguous().view(N, -1), 1, relative_n_inds.data ) # shape [N] p_inds = p_inds.squeeze(1) n_inds = n_inds.squeeze(1) return dist_ap, dist_an # , p_inds, n_inds return dist_ap, dist_an @registry.register_loss("xent") class CrossEntropyLoss(BaseLoss): def __init__(self, config): super().__init__(config) self.loss_fn = nn.CrossEntropyLoss() def forward(self, x, target): target = target.squeeze().long() return self.loss_fn(x, target) @registry.register_loss("bce") class BCELoss(BaseLoss): def __init__(self, config): super().__init__(config) self.loss_fn = nn.BCELoss() def forward(self, x, target): return self.loss_fn(x, target) @registry.register_loss("bcewlogits") class BCEWithLogitsLoss(BaseLoss): def __init__(self, config): super().__init__(config) self.loss_fn = nn.BCEWithLogitsLoss() def forward(self, x, target): return self.loss_fn(x, target) @registry.register_loss("l1") class L1Loss(BaseLoss): def __init__(self, config): super().__init__(config) self.loss_fn = nn.L1Loss() def forward(self, x, target): return self.loss_fn(x, target) @registry.register_loss("mse") class MSELoss(BaseLoss): def __init__(self, config): super().__init__(config) self.loss_fn = nn.MSELoss() def forward(self, x, target): return self.loss_fn(x, target) @registry.register_loss("smoothl1") class SmoothL1Loss(BaseLoss): def __init__(self, config): super().__init__(config) self.loss_fn = nn.SmoothL1Loss() def forward(self, x, target): return self.loss_fn(x, target) class Loss(nn.Module): def __init__(self): super().__init__() self.config = registry.get("loss_config") self._build() def _build(self): logger.debug("create criterion") self.losses = registry.get_loss_class(self.config.type)(self.config) def forward(self, x, target: torch.Tensor): return self.losses(x, target) fedlab-nlp/utils/metrics.py000066400000000000000000000036631427741542300162630ustar00rootroot00000000000000from abc import ABC from utils.register import registry from tools.glue_scripts.glue_metric import glue_compute_metrics class BaseMetric(ABC): def __init__(self, task_name, is_decreased_valid_metric=False): super().__init__() self.task_name = task_name self.is_decreased_valid_metric = is_decreased_valid_metric self.best_valid_metric = float("inf") if self.is_decreased_valid_metric else -float("inf") self.results = {} def calculate_metric(self, *args): raise NotImplementedError def update_metrics(self, *args): raise NotImplementedError @property def best_metric(self): return self.results @property def metric_name(self): raise NotImplementedError @registry.register_metric("glue") class GlueMetric(BaseMetric): def __init__(self, task_name, is_decreased_valid_metric=False): super().__init__(task_name, is_decreased_valid_metric) def calculate_metric(self, preds, labels, updated=True): results = glue_compute_metrics(self.task_name, preds, labels) if updated: self.update_metrics(results) return results def update_metrics(self, results): cur_valid_metric = results[self.metric_name] if self.is_decreased_valid_metric: is_best = cur_valid_metric < self.best_valid_metric else: is_best = cur_valid_metric > self.best_valid_metric if is_best: self.results.update(results) self.best_valid_metric = cur_valid_metric @property def metric_name(self): glue_metric_name = { "cola": "mcc", "sst-2": "acc", "mrpc": "f1", "sts-b": "acc", "qqp": "acc", "mnli": "acc", "mnli-mm": "acc", "qnli": "acc", "rte": "acc", "wnli": "acc" } return glue_metric_name[self.task_name] fedlab-nlp/utils/register.py000066400000000000000000000102111427741542300164240ustar00rootroot00000000000000 class Registry: mapping = { "state": {}, "loss_name_mapping": {}, "data_name_mapping": {}, "model_name_mapping": {}, "fl_name_mapping": {}, "metric_name_mapping": {}, "eval_name_mapping": {} } @classmethod def register(cls, name, obj): cls.mapping["state"][name] = obj @classmethod def get(cls, name, default=None, no_warning=False): original_name = name name = name.split(".") value = cls.mapping["state"] for subname in name: value = value.get(subname, default) if value is default: break if ( "writer" in cls.mapping["state"] and value == default and no_warning is False ): cls.mapping["state"]["writer"].warning( "Key {} is not present in registry, returning default value " "of {}".format(original_name, default) ) return value @classmethod def register_loss(cls, name): def wrap(func): from utils.loss import BaseLoss assert issubclass( func, BaseLoss ), "All loss must inherit utils.loss.BaseLoss class" cls.mapping["loss_name_mapping"][name] = func return func return wrap @classmethod def register_data(cls, name): def wrap(func): from data.base_dataloader import BaseDataLoader assert issubclass( func, BaseDataLoader ), "All dataset must inherit data.base_data_loader.BaseDataLoader class" cls.mapping["data_name_mapping"][name] = func return func return wrap @classmethod def register_model(cls, name): def wrap(func): from models.base_models import BaseModels assert issubclass( func, BaseModels ), "All model must inherit models.base_models.BaseModels class" cls.mapping["model_name_mapping"][name] = func return func return wrap @classmethod def register_fl_algorithm(cls, name): def wrap(func): from trainers.FedBaseTrainer import BaseTrainer assert issubclass( func, BaseTrainer ), "All federated algorithm must inherit trainers.base_fed_trainer.BaseTrainer class" cls.mapping["fl_name_mapping"][name] = func return func return wrap @classmethod def register_metric(cls, name): def wrap(func): from utils.metrics import BaseMetric assert issubclass( func, BaseMetric ), "All metric must inherit utils.metrics.BaseMetric class" cls.mapping["metric_name_mapping"][name] = func return func return wrap @classmethod def register_eval(cls, name): def wrap(func): from utils.evaluations import BaseEval assert issubclass( func, BaseEval ), "All evaluation must inherit utils.evaluations.BaseEval class" cls.mapping["eval_name_mapping"][name] = func return func return wrap @classmethod def get_loss_class(cls, name): return cls.mapping["loss_name_mapping"].get(name, None) @classmethod def get_data_class(cls, name): return cls.mapping["data_name_mapping"].get(name, None) @classmethod def get_model_class(cls, name): return cls.mapping["model_name_mapping"].get(name, None) @classmethod def get_fl_class(cls, name): return cls.mapping["fl_name_mapping"].get(name, None) @classmethod def get_metric_class(cls, name): return cls.mapping["metric_name_mapping"].get(name, None) @classmethod def get_eval_class(cls, name): return cls.mapping["eval_name_mapping"].get(name, None) @classmethod def unregister(cls, name): return cls.mapping["state"].pop(name, None) @classmethod def get_keys(cls): keys = list(cls.mapping["state"].keys()) return keys registry = Registry()