add federated learning inplementation by plato

Signed-off-by: XinYao1994 <xyao@cs.hku.hk>

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"

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

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"]

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"]

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
 
 

examples/federated_learning/yolov5_coco128_mistnet/README.md

@@ -32,37 +32,32 @@ 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) 
 
-```
-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
-```
+Create data interface for ```EDGE1_NODE```.
 
-move ```./data/1``` to `/data` of ```EDGE1_NODE```.
-
-```
-mkdir -p /data
-cd /data
-mv ./data/1 ./
+```shell
+mkdir -p /data/1
+cd /data/1
+wget https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip
+unzip coco128.zip -d COCO
 ```
 
-move ```./data/2``` to `/data` of ```EDGE2_NODE```.
+Create data interface for ```EDGE2_NODE```.
 
-```
-mkdir -p /data
-cd /data
-mv ./data/2 ./
+```shell
+mkdir -p /data/2
+cd /data/2
+wget https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip
+unzip coco128.zip -d 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 +65,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 - <<EOF
 apiVersion: sedna.io/v1alpha1
 kind: Dataset
 metadata:
-  name: "coco-dataset"
+  name: "coco-dataset-1"
 spec:
-  url: "/data/test.txt"
-  format: "txt"
-  nodeName: edge-node
+  url: "/data/1/COCO"
+  format: "dir"
+  nodeName: $EDGE1_NODE
 EOF
 ```
 
-create dataset for `$EDGE2_NODE`
-
-```
+```bash
 kubectl create -f - <<EOF
 apiVersion: sedna.io/v1alpha1
 kind: Dataset
 metadata:
-  name: "coco-dataset"
+  name: "coco-dataset-2"
 spec:
-  url: "/data/test.txt"
-  format: "txt"
-  nodeName: edge-node
+  url: "/data/2/COCO"
+  format: "dir"
+  nodeName: $EDGE2_NODE
 EOF
 ```
 
 #### Create Model
-
-create the directory `/model` in the host of `$EDGE1_NODE`
-
-```
-mkdir /model
+create the directory `/model` and `/pretrained` in `$EDGE1_NODE` and `$EDGE2_NODE`.
+```bash
+mkdir -p /model
+mkdir -p /pretrained
 ```
 
-create the directory `/model` in the host of `$EDGE2_NODE`
+create the directory `/model` and `/pretrained` in the host of `$CLOUD_NODE` (download links [here](https://kubeedge.obs.cn-north-1.myhuaweicloud.com/examples/yolov5_coco128_mistnet/yolov5.pth))
 
-```
-mkdir /model
-```
 
-```
-TODO: put pretrained model on nodes.
+```bash
+# on the cloud side
+mkdir -p /model
+mkdir -p /pretrained
+cd /pretrained
+wget https://kubeedge.obs.cn-north-1.myhuaweicloud.com/examples/yolov5_coco128_mistnet/yolov5.pth
 ```
 
 create model
 
-```
+```bash
 kubectl create -f - <<EOF
 apiVersion: sedna.io/v1alpha1
 kind: Model
 metadata:
   name: "yolo-v5-model"
 spec:
-  url: "/model/yolo.pb"
-  format: "pb"
+  url: "/model/yolov5.pth"
+  format: "pth"
+EOF
+
+kubectl create -f - <<EOF
+apiVersion: sedna.io/v1alpha1
+kind: Model
+metadata:
+  name: "yolo-v5-pretrained-model"
+spec:
+  url: "/pretrained/yolov5.pth"
+  format: "pth"
 EOF
 ```
 
-#### Start Federated Learning Job
+### Create a secret with your S3 user credential. (Optional)
 
+```shell
+kubectl create -f - <<EOF
+apiVersion: v1
+kind: Secret
+metadata:
+  name: mysecret
+  annotations:
+    s3-endpoint: s3.amazonaws.com 
+    s3-usehttps: "1" 
+stringData: 
+  ACCESS_KEY_ID: XXXX
+  SECRET_ACCESS_KEY: XXXXXXXX
+EOF
 ```
+
+#### Start Federated Learning Job
+
+```bash
 kubectl create -f - <<EOF
 apiVersion: sedna.io/v1alpha1
 kind: FederatedLearningJob
 metadata:
-  name: mistnet-on-mnist-dataset
+  name: yolo-v5
 spec:
-  stopCondition:
-    operator: "or" # and
-      conditions:
-        - operator: ">"
-          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 +186,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`.

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.pth", "")
+
+    server = AggregationServerV2(
         data=data,
         estimator=estimator,
         aggregation=mistnet,

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' ]

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)

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 @@ mistnet = MistNet(cut_layer=Context.get_parameters("cut_layer"),
                   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 @@ class Dataset:
             "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 @@ class Dataset:
 class Estimator:
     def __init__(self) -> None:
         self.model = None
+        self.pretrained = None
         self.hyperparameters = {
             "type": "yolov5",
             "rounds": 1,

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()

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)
+  ]

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

lib/sedna/algorithms/aggregation/aggregation.py

@@ -110,8 +110,18 @@ class FedAvg(BaseAggregation, abc.ABC):
 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

lib/sedna/algorithms/client_choose/client_choose.py

@@ -33,4 +33,6 @@ class SimpleClientChoose(AbstractClientChoose):
 
     def __init__(self, per_round=1):
         super().__init__()
-        self.per_round = per_round
+        self.parameters = {
+            "per_round": per_round
+        }

lib/sedna/algorithms/transmitter/transmitter.py

@@ -35,6 +35,9 @@ class WSTransmitter(AbstractTransmitter, ABC):
     An implementation of Transmitter based on WebSocket.
     """
 
+    def __init__(self):
+        self.parameters = {}
+
     def recv(self):
         pass
 
@@ -52,10 +55,12 @@ class S3Transmitter(AbstractTransmitter, ABC):
                  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

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

lib/sedna/core/federated_learning/federated_learning.py

@@ -17,9 +17,6 @@ import asyncio
 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.common.file_ops import FileOps
 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 @@ class FederatedLearningV0(JobBase):
                     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 @@ class FederatedLearning:
         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)

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 @@ class AggregationServer(BaseServer):
         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()