diff --git a/build/crd-samples/sedna/federatedlearningjob_yolo_v1alpha1.yaml b/build/crd-samples/sedna/federatedlearningjob_yolo_v1alpha1.yaml index 9864951a9..8230be0d6 100644 --- a/build/crd-samples/sedna/federatedlearningjob_yolo_v1alpha1.yaml +++ b/build/crd-samples/sedna/federatedlearningjob_yolo_v1alpha1.yaml @@ -5,7 +5,7 @@ metadata: spec: pretrainedModel: # option name: "yolo-v5-pretrained-model" - transimitter: # option + transmitter: # option ws: { } # option, by default s3: # option, but at least one aggDataPath: "s3://sedna/fl/aggregation_data" @@ -17,7 +17,7 @@ spec: spec: nodeName: "sedna-control-plane" containers: - - image: kubeedge/sedna-fl-aggregation:mistnetyolo + - image: kubeedge/sedna-example-federated-learning-mistnet-yolo-aggregator:v0.4.0 name: agg-worker imagePullPolicy: IfNotPresent env: # user defined environments @@ -28,21 +28,54 @@ spec: - name: "aggregation_algorithm" value: "mistnet" - name: "batch_size" + value: "32" resources: # user defined resources limits: memory: 8Gi trainingWorkers: - dataset: - name: "coco-dataset" + name: "coco-dataset-1" template: spec: nodeName: "edge-node" containers: - - image: kubeedge/sedna-fl-train:mistnetyolo + - image: kubeedge/sedna-example-federated-learning-mistnet-yolo-client:v0.4.0 name: train-worker imagePullPolicy: IfNotPresent args: [ "-i", "1" ] env: # user defined environments + - name: "cut_layer" + value: "4" + - name: "epsilon" + value: "100" + - name: "aggregation_algorithm" + value: "mistnet" + - name: "batch_size" + value: "32" + - name: "learning_rate" + value: "0.001" + - name: "epochs" + value: "1" + resources: # user defined resources + limits: + memory: 2Gi + - dataset: + name: "coco-dataset-2" + template: + spec: + nodeName: "edge-node" + containers: + - image: kubeedge/sedna-example-federated-learning-mistnet-yolo-client:v0.4.0 + name: train-worker + imagePullPolicy: IfNotPresent + args: [ "-i", "2" ] + env: # user defined environments + - name: "cut_layer" + value: "4" + - name: "epsilon" + value: "100" + - name: "aggregation_algorithm" + value: "mistnet" - name: "batch_size" value: "32" - name: "learning_rate" diff --git a/examples/build_image.sh b/examples/build_image.sh index 6a154d845..fb05c2c3f 100644 --- a/examples/build_image.sh +++ b/examples/build_image.sh @@ -17,11 +17,13 @@ cd "$(dirname "${BASH_SOURCE[0]}")" IMAGE_REPO=${IMAGE_REPO:-kubeedge} -IMAGE_TAG=${IMAGE_TAG:-v0.3.0} +IMAGE_TAG=${IMAGE_TAG:-v0.4.0} EXAMPLE_REPO_PREFIX=${IMAGE_REPO}/sedna-example- dockerfiles=( +federated-learning-mistnet-yolo-aggregator.Dockerfile +federated-learning-mistnet-yolo-client.Dockerfile federated-learning-surface-defect-detection-aggregation.Dockerfile federated-learning-surface-defect-detection-train.Dockerfile incremental-learning-helmet-detection.Dockerfile diff --git a/examples/federated-learning-mistnet-yolo-aggregator.Dockerfile b/examples/federated-learning-mistnet-yolo-aggregator.Dockerfile new file mode 100644 index 000000000..e316f6eb3 --- /dev/null +++ b/examples/federated-learning-mistnet-yolo-aggregator.Dockerfile @@ -0,0 +1,23 @@ +FROM tensorflow/tensorflow:1.15.4 + +RUN apt update \ + && apt install -y libgl1-mesa-glx git + +COPY ./lib/requirements.txt /home + +RUN python -m pip install --upgrade pip + +RUN pip install -r /home/requirements.txt + +ENV PYTHONPATH "/home/lib:/home/plato:/home/plato/packages/yolov5" + +COPY ./lib /home/lib +RUN git clone https://github.com/TL-System/plato.git /home/plato + +RUN pip install -r /home/plato/requirements.txt +RUN pip install -r /home/plato/packages/yolov5/requirements.txt + +WORKDIR /home/work +COPY examples/federated_learning/yolov5_coco128_mistnet /home/work/ + +CMD ["/bin/sh", "-c", "ulimit -n 50000; python aggregate.py"] diff --git a/examples/federated-learning-mistnet-yolo-client.Dockerfile b/examples/federated-learning-mistnet-yolo-client.Dockerfile new file mode 100644 index 000000000..b1e7aa356 --- /dev/null +++ b/examples/federated-learning-mistnet-yolo-client.Dockerfile @@ -0,0 +1,23 @@ +FROM tensorflow/tensorflow:1.15.4 + +RUN apt update \ + && apt install -y libgl1-mesa-glx git + +COPY ./lib/requirements.txt /home + +RUN python -m pip install --upgrade pip + +RUN pip install -r /home/requirements.txt + +ENV PYTHONPATH "/home/lib:/home/plato:/home/plato/packages/yolov5" + +COPY ./lib /home/lib +RUN git clone https://github.com/TL-System/plato.git /home/plato + +RUN pip install -r /home/plato/requirements.txt +RUN pip install -r /home/plato/packages/yolov5/requirements.txt + +WORKDIR /home/work +COPY examples/federated_learning/yolov5_coco128_mistnet /home/work/ + +ENTRYPOINT ["python", "train.py"] diff --git a/examples/federated_learning/surface_defect_detection/training_worker/train.py b/examples/federated_learning/surface_defect_detection/training_worker/train.py index 4fd9a1122..37f21d0cc 100644 --- a/examples/federated_learning/surface_defect_detection/training_worker/train.py +++ b/examples/federated_learning/surface_defect_detection/training_worker/train.py @@ -11,7 +11,6 @@ # 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 numpy as np @@ -74,6 +73,7 @@ def main(): learning_rate=learning_rate, validation_split=validation_split ) + return train_jobs diff --git a/examples/federated_learning/yolov5_coco128_mistnet/README.md b/examples/federated_learning/yolov5_coco128_mistnet/README.md index 509734ae8..5214859b6 100644 --- a/examples/federated_learning/yolov5_coco128_mistnet/README.md +++ b/examples/federated_learning/yolov5_coco128_mistnet/README.md @@ -32,37 +32,43 @@ Follow the [Sedna installation document](/docs/setup/install.md) to install Sedn ### Prepare Dataset -Download [dataset](https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip) and do data partition +Download [dataset](https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip) -``` +Create data interface for ```EDGE1_NODE```. + +```shell +mkdir -p /data +mkdir -p /data/1 +cd /data/1 wget https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip -unzip coco128.zip -d data -rm coco128.zip -python partition.py ./data 2 +unzip coco128.zip -d COCO +touch COCO/robot.txt ``` -move ```./data/1``` to `/data` of ```EDGE1_NODE```. +Create data interface for ```EDGE2_NODE```. -``` +```shell mkdir -p /data -cd /data -mv ./data/1 ./ +mkdir -p /data/2 +cd /data/2 +wget https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip +unzip coco128.zip -d COCO +touch COCO/robot.txt ``` -move ```./data/2``` to `/data` of ```EDGE2_NODE```. -``` -mkdir -p /data -cd /data -mv ./data/2 ./ +create the directory `./data/COCO` in the host of `$CLOUD_NODE` for storing test data. +```bash +mkdir ./data +mkdir ./data/COCO ``` ### Prepare Images This example uses these images: -1. aggregation worker: ```kubeedge/sedna-example-federated-learning-mistnet:v0.3.0``` -2. train worker: ```kubeedge/sedna-example-federated-learning-mistnet-client:v0.3.0``` +1. aggregation worker: ```kubeedge/sedna-example-federated-learning-mistnet-yolo-aggregato:v0.4.0``` +2. train worker: ```kubeedge/sedna-example-federated-learning-mistnet-yolo-client:v0.4.0``` These images are generated by the script [build_images.sh](/examples/build_image.sh). @@ -70,103 +76,118 @@ These images are generated by the script [build_images.sh](/examples/build_image #### Create Dataset -create dataset for `$EDGE1_NODE` +create dataset for `$EDGE1_NODE` and `$EDGE2_NODE` -```n +```bash kubectl create -f - <" - threshold: 100 - metric: rounds - - operator: ">" - threshold: 0.95 - metric: targetAccuracy - - operator: "<" - threshold: 0.03 - metric: deltaLoss - aggregationTrigger: - condition: - operator: ">" - threshold: 5 - metric: num_of_ready_clients + pretrainedModel: # option + name: "yolo-v5-pretrained-model" + transmitter: # option + ws: { } # option, by default + s3: # optional, but at least one + aggDataPath: "s3://sedna/fl/aggregation_data" + credentialName: mysecret aggregationWorker: model: - name: "mistnet-on-mnist-model" + name: "yolo-v5-model" template: spec: nodeName: $CLOUD_NODE containers: - - image: kubeedge/sedna-example-federated-learning-mistnet-on-mnist-dataset-aggregation:v0.4.0 - name: agg-worker + - image: kubeedge/sedna-example-federated-learning-mistnet-yolo-aggregator:v0.4.0 + name: agg-worker imagePullPolicy: IfNotPresent env: # user defined environments - name: "cut_layer" @@ -176,63 +197,64 @@ spec: - name: "aggregation_algorithm" value: "mistnet" - name: "batch_size" - value: "10" - resources: # user defined resources + value: "32" + resources: # user defined resources limits: - memory: 2Gi + memory: 8Gi trainingWorkers: - dataset: - name: "edge1-surface-defect-detection-dataset" + name: "coco-dataset-1" template: spec: nodeName: $EDGE1_NODE containers: - - image: kubeedge/sedna-example-federated-learning-mistnet-on-mnist-dataset-train:v0.4.0 - name: train-worker + - image: kubeedge/sedna-example-federated-learning-mistnet-yolo-client:v0.4.0 + name: train-worker imagePullPolicy: IfNotPresent - env: # user defined environments + args: [ "-i", "1" ] + env: # user defined environments + - name: "cut_layer" + value: "4" + - name: "epsilon" + value: "100" + - name: "aggregation_algorithm" + value: "mistnet" - name: "batch_size" value: "32" - name: "learning_rate" value: "0.001" - name: "epochs" - value: "2" - resources: # user defined resources + value: "1" + resources: # user defined resources limits: memory: 2Gi - dataset: - name: "edge2-surface-defect-detection-dataset" + name: "coco-dataset-2" template: spec: nodeName: $EDGE2_NODE containers: - - image: kubeedge/sedna-example-federated-learning-mistnet-on-mnist-dataset-train:v0.4.0 - name: train-worker + - image: kubeedge/sedna-example-federated-learning-mistnet-yolo-client:v0.4.0 + name: train-worker imagePullPolicy: IfNotPresent - env: # user defined environments + args: [ "-i", "2" ] + env: # user defined environments + - name: "cut_layer" + value: "4" + - name: "epsilon" + value: "100" + - name: "aggregation_algorithm" + value: "mistnet" - name: "batch_size" value: "32" - name: "learning_rate" value: "0.001" - name: "epochs" - value: "2" - resources: # user defined resources + value: "1" + resources: # user defined resources limits: memory: 2Gi EOF ``` -``` -TODO: show the benifit of mistnet. for example, the compared results of fedavg & mistnet. - -``` - -### Check Federated Learning Status - -``` -kubectl get federatedlearningjob surface-defect-detection -``` - -### Check Federated Learning Train Result -After the job completed, you will find the model generated on the directory `/model` in `$EDGE1_NODE` and `$EDGE2_NODE`. diff --git a/examples/federated_learning/yolov5_coco128_mistnet/aggregate.py b/examples/federated_learning/yolov5_coco128_mistnet/aggregate.py index 0ba8d558f..015ef643d 100644 --- a/examples/federated_learning/yolov5_coco128_mistnet/aggregate.py +++ b/examples/federated_learning/yolov5_coco128_mistnet/aggregate.py @@ -14,14 +14,16 @@ from interface import mistnet, s3_transmitter, simple_chooser from interface import Dataset, Estimator -from sedna.service.server import AggregationServer - +from sedna.service.server import AggregationServerV2 +from sedna.common.config import BaseConfig def run_server(): data = Dataset() estimator = Estimator() - server = AggregationServer( + estimator.pretrained = BaseConfig.pretrained_model_url.replace("yolov5", "") + + server = AggregationServerV2( data=data, estimator=estimator, aggregation=mistnet, diff --git a/examples/federated_learning/yolov5_coco128_mistnet/coco128.yaml b/examples/federated_learning/yolov5_coco128_mistnet/coco128.yaml new file mode 100644 index 000000000..cabd6e133 --- /dev/null +++ b/examples/federated_learning/yolov5_coco128_mistnet/coco128.yaml @@ -0,0 +1,28 @@ +# COCO 2017 dataset http://cocodataset.org - first 128 training images +# Train command: python train.py --data coco128.yaml +# Default dataset location is next to YOLOv5: +# /parent_folder +# /coco128 +# /yolov5 + + +# download command/URL (optional) +download: https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip + +# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/] +train: ./data/COCO/coco128/images/train2017/ # 128 images +val: ./data/COCO/coco128/images/train2017/ # 128 images + +# number of classes +nc: 80 + +# class names +names: [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', + 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', + 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', + 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', + 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', + 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', + 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', + 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', + 'hair drier', 'toothbrush' ] diff --git a/examples/federated_learning/yolov5_coco128_mistnet/hyp.scratch.yaml b/examples/federated_learning/yolov5_coco128_mistnet/hyp.scratch.yaml new file mode 100644 index 000000000..44f26b665 --- /dev/null +++ b/examples/federated_learning/yolov5_coco128_mistnet/hyp.scratch.yaml @@ -0,0 +1,33 @@ +# Hyperparameters for COCO training from scratch +# python train.py --batch 40 --cfg yolov5m.yaml --weights '' --data coco.yaml --img 640 --epochs 300 +# See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials + + +lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3) +lrf: 0.2 # final OneCycleLR learning rate (lr0 * lrf) +momentum: 0.937 # SGD momentum/Adam beta1 +weight_decay: 0.0005 # optimizer weight decay 5e-4 +warmup_epochs: 3.0 # warmup epochs (fractions ok) +warmup_momentum: 0.8 # warmup initial momentum +warmup_bias_lr: 0.1 # warmup initial bias lr +box: 0.05 # box loss gain +cls: 0.5 # cls loss gain +cls_pw: 1.0 # cls BCELoss positive_weight +obj: 1.0 # obj loss gain (scale with pixels) +obj_pw: 1.0 # obj BCELoss positive_weight +iou_t: 0.20 # IoU training threshold +anchor_t: 4.0 # anchor-multiple threshold +# anchors: 3 # anchors per output layer (0 to ignore) +fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5) +hsv_h: 0.015 # image HSV-Hue augmentation (fraction) +hsv_s: 0.7 # image HSV-Saturation augmentation (fraction) +hsv_v: 0.4 # image HSV-Value augmentation (fraction) +degrees: 0.0 # image rotation (+/- deg) +translate: 0.1 # image translation (+/- fraction) +scale: 0.5 # image scale (+/- gain) +shear: 0.0 # image shear (+/- deg) +perspective: 0.0 # image perspective (+/- fraction), range 0-0.001 +flipud: 0.0 # image flip up-down (probability) +fliplr: 0.5 # image flip left-right (probability) +mosaic: 1.0 # image mosaic (probability) +mixup: 0.0 # image mixup (probability) diff --git a/examples/federated_learning/yolov5_coco128_mistnet/interface.py b/examples/federated_learning/yolov5_coco128_mistnet/interface.py index 6c654f58c..0f56a3a42 100644 --- a/examples/federated_learning/yolov5_coco128_mistnet/interface.py +++ b/examples/federated_learning/yolov5_coco128_mistnet/interface.py @@ -15,7 +15,7 @@ from sedna.algorithms.aggregation import MistNet from sedna.algorithms.client_choose import SimpleClientChoose from sedna.common.config import Context -from sedna.core.federated_learning import FederatedLearning +from sedna.core.federated_learning import FederatedLearningV2 simple_chooser = SimpleClientChoose(per_round=1) @@ -24,7 +24,7 @@ epsilon=Context.get_parameters("epsilon")) # The function `get_transmitter_from_config()` returns an object instance. -s3_transmitter = FederatedLearning.get_transmitter_from_config() +s3_transmitter = FederatedLearningV2.get_transmitter_from_config() class Dataset: @@ -44,6 +44,7 @@ def __init__(self) -> None: "num_classes": 80, # image size "image_size": 640, + "download_urls": ["https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip",], "classes": [ "person", @@ -134,6 +135,7 @@ def __init__(self) -> None: class Estimator: def __init__(self) -> None: self.model = None + self.pretrained = None self.hyperparameters = { "type": "yolov5", "rounds": 1, diff --git a/examples/federated_learning/yolov5_coco128_mistnet/train.py b/examples/federated_learning/yolov5_coco128_mistnet/train.py index 62886cb34..99406dd21 100644 --- a/examples/federated_learning/yolov5_coco128_mistnet/train.py +++ b/examples/federated_learning/yolov5_coco128_mistnet/train.py @@ -11,23 +11,25 @@ # 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 from interface import mistnet, s3_transmitter from interface import Dataset, Estimator -from sedna.core.federated_learning import FederatedLearning - +from sedna.common.config import BaseConfig +from sedna.core.federated_learning import FederatedLearningV2 def main(): data = Dataset() estimator = Estimator() - - fl_model = FederatedLearning( + data.parameters["data_path"] = BaseConfig.train_dataset_url.replace("robot.txt", "") + data.parameters["train_path"] = os.path.join(data.parameters["data_path"], "./coco128/images/train2017/") + data.parameters["test_path"] = data.parameters["train_path"] + fl_model = FederatedLearningV2( + data=data, estimator=estimator, aggregation=mistnet, transmitter=s3_transmitter) - fl_model.train(data) - + fl_model.train() if __name__ == '__main__': main() diff --git a/examples/federated_learning/yolov5_coco128_mistnet/yolov5s.yaml b/examples/federated_learning/yolov5_coco128_mistnet/yolov5s.yaml new file mode 100644 index 000000000..e4e9e4dde --- /dev/null +++ b/examples/federated_learning/yolov5_coco128_mistnet/yolov5s.yaml @@ -0,0 +1,48 @@ +# parameters +nc: 80 # number of classes +depth_multiple: 0.33 # model depth multiple +width_multiple: 0.50 # layer channel multiple + +# anchors +anchors: + - [ 10,13, 16,30, 33,23 ] # P3/8 + - [ 30,61, 62,45, 59,119 ] # P4/16 + - [ 116,90, 156,198, 373,326 ] # P5/32 + +# YOLOv5 backbone +backbone: + # [from, number, module, args] + [ [ -1, 1, Focus, [ 64, 3 ] ], # 0-P1/2 + [ -1, 1, Conv, [ 128, 3, 2 ] ], # 1-P2/4 + [ -1, 3, C3, [ 128 ] ], + [ -1, 1, Conv, [ 256, 3, 2 ] ], # 3-P3/8 + [ -1, 9, C3, [ 256 ] ], + [ -1, 1, Conv, [ 512, 3, 2 ] ], # 5-P4/16 + [ -1, 9, C3, [ 512 ] ], + [ -1, 1, Conv, [ 1024, 3, 2 ] ], # 7-P5/32 + [ -1, 1, SPP, [ 1024, [ 5, 9, 13 ] ] ], + [ -1, 3, C3, [ 1024, False ] ], # 9 + ] + +# YOLOv5 head +head: + [ [ -1, 1, Conv, [ 512, 1, 1 ] ], + [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], + [ [ -1, 6 ], 1, Concat, [ 1 ] ], # cat backbone P4 + [ -1, 3, C3, [ 512, False ] ], # 13 + + [ -1, 1, Conv, [ 256, 1, 1 ] ], + [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], + [ [ -1, 4 ], 1, Concat, [ 1 ] ], # cat backbone P3 + [ -1, 3, C3, [ 256, False ] ], # 17 (P3/8-small) + + [ -1, 1, Conv, [ 256, 3, 2 ] ], + [ [ -1, 14 ], 1, Concat, [ 1 ] ], # cat head P4 + [ -1, 3, C3, [ 512, False ] ], # 20 (P4/16-medium) + + [ -1, 1, Conv, [ 512, 3, 2 ] ], + [ [ -1, 10 ], 1, Concat, [ 1 ] ], # cat head P5 + [ -1, 3, C3, [ 1024, False ] ], # 23 (P5/32-large) + + [ [ 17, 20, 23 ], 1, Detect, [ nc, anchors ] ], # Detect(P3, P4, P5) + ] diff --git a/lib/sedna/algorithms/aggregation/__init__.py b/lib/sedna/algorithms/aggregation/__init__.py index 96d7da5e5..4725746ab 100644 --- a/lib/sedna/algorithms/aggregation/__init__.py +++ b/lib/sedna/algorithms/aggregation/__init__.py @@ -13,4 +13,4 @@ # limitations under the License. from . import aggregation -from .aggregation import FedAvg, MistNet +from .aggregation import FedAvg, MistNet, AggClient diff --git a/lib/sedna/algorithms/aggregation/aggregation.py b/lib/sedna/algorithms/aggregation/aggregation.py index 91121cf9a..b49c0198b 100644 --- a/lib/sedna/algorithms/aggregation/aggregation.py +++ b/lib/sedna/algorithms/aggregation/aggregation.py @@ -109,8 +109,18 @@ def aggregate(self, clients: List[AggClient]): class MistNet(BaseAggregation, abc.ABC): def __init__(self, cut_layer, epsilon=100): super().__init__() - self.cut_layer = cut_layer - self.epsilon = epsilon + self.parameters = { + "type": "mistnet", + "cut_layer": cut_layer, + "epsilon": epsilon + } + if isinstance(self.parameters["cut_layer"], str): + if self.parameters["cut_layer"].isdigit(): + self.parameters["cut_layer"] = int(cut_layer) + + if isinstance(self.parameters["epsilon"], str): + if self.parameters["epsilon"].isdigit(): + self.parameters["epsilon"] = int(cut_layer) def aggregate(self, clients: List[AggClient]): pass diff --git a/lib/sedna/algorithms/client_choose/client_choose.py b/lib/sedna/algorithms/client_choose/client_choose.py index 9e9fae603..17f94262e 100644 --- a/lib/sedna/algorithms/client_choose/client_choose.py +++ b/lib/sedna/algorithms/client_choose/client_choose.py @@ -26,4 +26,6 @@ class SimpleClientChoose(AbstractClientChoose): def __init__(self, per_round=1): super().__init__() - self.per_round = per_round + self.parameters = { + "per_round": per_round + } diff --git a/lib/sedna/algorithms/transmitter/transmitter.py b/lib/sedna/algorithms/transmitter/transmitter.py index 9ea01f61d..b0b4831f6 100644 --- a/lib/sedna/algorithms/transmitter/transmitter.py +++ b/lib/sedna/algorithms/transmitter/transmitter.py @@ -28,6 +28,9 @@ def send(self, data): class WSTransmitter(AbstractTransmitter, ABC): + def __init__(self): + self.parameters = {} + def recv(self): pass @@ -42,10 +45,12 @@ def __init__(self, access_key, secret_key, transmitter_url): - self.s3_endpoint_url = s3_endpoint_url - self.access_key = access_key - self.secret_key = secret_key - self.transmitter_url = transmitter_url + self.parameters = { + "s3_endpoint_url": s3_endpoint_url, + "s3_bucket": transmitter_url, + "access_key": access_key, + "secret_key": secret_key + } def recv(self): pass diff --git a/lib/sedna/core/federated_learning/__init__.py b/lib/sedna/core/federated_learning/__init__.py index c36fe3b80..e5eaad2d7 100644 --- a/lib/sedna/core/federated_learning/__init__.py +++ b/lib/sedna/core/federated_learning/__init__.py @@ -13,3 +13,4 @@ # limitations under the License. from .federated_learning import FederatedLearning +from .federated_learning import FederatedLearningV2 diff --git a/lib/sedna/core/federated_learning/federated_learning.py b/lib/sedna/core/federated_learning/federated_learning.py index f49825dda..999a153ed 100644 --- a/lib/sedna/core/federated_learning/federated_learning.py +++ b/lib/sedna/core/federated_learning/federated_learning.py @@ -17,9 +17,6 @@ import sys import time -from plato.clients import registry as client_registry -from plato.config import Config - from sedna.algorithms.transmitter import S3Transmitter, WSTransmitter from sedna.common.class_factory import ClassFactory, ClassType from sedna.common.config import BaseConfig, Context @@ -28,8 +25,10 @@ from sedna.core.base import JobBase from sedna.service.client import AggregationClient +__all__ = ('FederatedLearning', 'FederatedLearningV2') + -class FederatedLearningV0(JobBase): +class FederatedLearning(JobBase): """ Federated learning enables multiple actors to build a common, robust machine learning model without sharing data, thus allowing to address @@ -187,8 +186,12 @@ def train(self, train_data, task_info_res) -class FederatedLearning: - def __init__(self, data=None, estimator=None, aggregation=None, transmitter=None) -> None: +class FederatedLearningV2: + def __init__(self, data=None, estimator=None, + aggregation=None, transmitter=None) -> None: + + from plato.config import Config + from plato.clients import registry as client_registry # set parameters server = Config.server._asdict() clients = Config.clients._asdict() @@ -196,32 +199,27 @@ def __init__(self, data=None, estimator=None, aggregation=None, transmitter=None train = Config.trainer._asdict() if data is not None: - for xkey in data.parameters: - datastore[xkey] = data.parameters[xkey] + datastore.update(data.parameters) Config.data = Config.namedtuple_from_dict(datastore) self.model = None if estimator is not None: self.model = estimator.model - for xkey in estimator.hyperparameters: - train[xkey] = estimator.hyperparameters[xkey] + train.update(estimator.hyperparameters) Config.trainer = Config.namedtuple_from_dict(train) if aggregation is not None: - Config.algorithm = Config.namedtuple_from_dict(aggregation.parameters) + Config.algorithm = Config.namedtuple_from_dict( + aggregation.parameters) if aggregation.parameters["type"] == "mistnet": clients["type"] = "mistnet" server["type"] = "mistnet" - if isinstance(transmitter, S3Transmitter): - server["address"] = Context.get_parameters("AGG_IP") - server["port"] = Context.get_parameters("AGG_PORT") - server["s3_endpoint_url"] = transmitter.s3_endpoint_url - server["s3_bucket"] = transmitter.s3_bucket - server["access_key"] = transmitter.access_key - server["secret_key"] = transmitter.secret_key - elif isinstance(transmitter, WSTransmitter): - pass + server["address"] = Context.get_parameters("AGG_IP") + server["port"] = Context.get_parameters("AGG_PORT") + + if transmitter is not None: + server.update(transmitter.parameters) Config.server = Config.namedtuple_from_dict(server) Config.clients = Config.namedtuple_from_dict(clients) diff --git a/lib/sedna/service/server/aggregation.py b/lib/sedna/service/server/aggregation.py index 3529d6b48..717540fea 100644 --- a/lib/sedna/service/server/aggregation.py +++ b/lib/sedna/service/server/aggregation.py @@ -13,26 +13,26 @@ # limitations under the License. import time -from typing import List, Optional, Dict, Any - import uuid -from pydantic import BaseModel +from typing import Any, Dict, List, Optional + from fastapi import FastAPI, WebSocket from fastapi.routing import APIRoute +from pydantic import BaseModel +from starlette.endpoints import WebSocketEndpoint from starlette.requests import Request from starlette.responses import JSONResponse from starlette.routing import WebSocketRoute -from starlette.endpoints import WebSocketEndpoint from starlette.types import ASGIApp, Receive, Scope, Send +from sedna.algorithms.aggregation import AggClient +from sedna.common.config import BaseConfig, Context +from sedna.common.class_factory import ClassFactory, ClassType from sedna.common.log import LOGGER from sedna.common.utils import get_host_ip -from sedna.common.class_factory import ClassFactory, ClassType -from sedna.algorithms.aggregation import AggClient - from .base import BaseServer -__all__ = ('AggregationServer',) +__all__ = ('AggregationServer', 'AggregationServerV2') class WSClientInfo(BaseModel): # pylint: disable=too-few-public-methods @@ -266,3 +266,56 @@ async def client_info(self, request: Request): if client_id: return server.get_client(client_id) return WSClientInfoList(clients=server.client_list) + + +class AggregationServerV2(): + def __init__(self, data=None, estimator=None, + aggregation=None, transmitter=None, + chooser=None) -> None: + from plato.config import Config + from plato.servers import registry as server_registry + # set parameters + server = Config.server._asdict() + clients = Config.clients._asdict() + datastore = Config.data._asdict() + train = Config.trainer._asdict() + + if data is not None: + datastore.update(data.parameters) + Config.data = Config.namedtuple_from_dict(datastore) + + self.model = None + if estimator is not None: + self.model = estimator.model + if estimator.pretrained is not None: + LOGGER.info(estimator.pretrained) + Config.params['model_dir'] = estimator.pretrained + train.update(estimator.hyperparameters) + Config.trainer = Config.namedtuple_from_dict(train) + + server["address"] = Context.get_parameters("AGG_BIND_IP", "0.0.0.0") + server["port"] = Context.get_parameters("AGG_BIND_PORT", 7363) + if transmitter is not None: + server.update(transmitter.parameters) + + if aggregation is not None: + Config.algorithm = Config.namedtuple_from_dict( + aggregation.parameters) + if aggregation.parameters["type"] == "mistnet": + clients["type"] = "mistnet" + server["type"] = "mistnet" + + if chooser is not None: + clients["per_round"] = chooser.parameters["per_round"] + + LOGGER.info("address %s, port %s", server["address"], server["port"]) + + Config.server = Config.namedtuple_from_dict(server) + Config.clients = Config.namedtuple_from_dict(clients) + + # Config.store() + # create a server + self.server = server_registry.get(model=self.model) + + def start(self): + self.server.run()