# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
#
# 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 typing import List, Optional, Union
import torch.nn as nn
from nvflare.app_common.widgets.convert_to_fed_event import ConvertToFedEvent
from nvflare.app_common.workflows.fedavg import FedAvg
from nvflare.app_opt.pt.job_config.base_fed_job import BaseFedJob
[docs]
class FedAvgJob(BaseFedJob):
def __init__(
self,
initial_model: nn.Module,
n_clients: int,
num_rounds: int,
name: str = "fed_job",
min_clients: int = 1,
mandatory_clients: Optional[List[str]] = None,
key_metric: str = "accuracy",
convert_to_fed_event: Union[bool, ConvertToFedEvent, None] = None,
analytics_receiver: Union[bool, ConvertToFedEvent, None] = None,
):
"""PyTorch FedAvg Job.
Configures server side FedAvg controller, persistor with initial model, and widgets.
User must add executors.
Args:
initial_model (nn.Module): initial PyTorch Model
n_clients (int): number of clients for this job
num_rounds (int): number of rounds for FedAvg
name (name, optional): name of the job. Defaults to "fed_job"
min_clients (int, optional): the minimum number of clients for the job. Defaults to 1.
mandatory_clients (List[str], optional): mandatory clients to run the job. Default None.
key_metric (str, optional): Metric used to determine if the model is globally best.
if metrics are a `dict`, `key_metric` can select the metric used for global model selection.
Defaults to "accuracy".
convert_to_fed_event (Union[bool, ConvertToFedEvent, None]): A component to convert certain events to fed events.
If not provided, a ConvertToFedEvent object will be created and add to Client
If provided, a ConvertToFedEvent object will be add to Client
If set to True, a ConvertToFedEvent object will be created and add to Client
If set to False, no ConvertToFedEvent will be created.
analytics_receiver (Union[bool, AnalyticsReceiver, None]): A component to receive analytics.
If not provided, a TBAnalyticsReceiver object will be created and add to seerver
If provided, a AnalyticsReceiver object will be add to Server
If set to True, a TBAnalyticsReceiver object will be created and add to seerver
If set to False, no AnalyticsReceiver will be created.
"""
if not isinstance(initial_model, nn.Module):
raise ValueError(f"Expected initial model to be nn.Module, but got type f{type(initial_model)}.")
super().__init__(
initial_model,
name,
min_clients,
mandatory_clients,
key_metric,
convert_to_fed_event=convert_to_fed_event,
analytics_receiver=analytics_receiver,
)
controller = FedAvg(
num_clients=n_clients,
num_rounds=num_rounds,
persistor_id=self.comp_ids["persistor_id"],
)
self.to_server(controller)