[Q&A] job is not execute after submit the job

[Gozilasim]

Python version (python3 -V)

3.8.10

NVFlare version (python3 -m pip list | grep "nvflare")

2.5.0

NVFlare branch (if running examples, please use the branch that corresponds to the NVFlare version, git branch)

2.5

Operating system

Ubuntu 20.04.6LTS

Have you successfully run any of the following examples?

• hello-numpy-sag with simulator
• hello-pt with simulator
• hello-numpy-sag with POC
• hello-pt with POC

Please describe your question

Hello guys, i try to run the cifar10-real-world example, i had successfully connect server with client , when i try to submit the job via admin or the .sh file, the server prompt out this message

DefaultJobScheduler - INFO - [identity=secure_project, run=?]: Try to schedule job c8c22eab-9c9c-404a-9716-ae8389948583, get result: (not enough sites have enough resources (ok sites 0 < min sites 1)).

my device is jetson agx orin, when i set the resource,json for each site with the num_of_gpus = 1, it also error and came with this error

ConfigError: ConfigError: Error processing '/mnt/5958b632-9b58-4b1e-9b06-aaa3d99c9dcd/cifar10/cifar10-real-world/workspaces/secure_workspace/site-1/startup/../startup/fed_client.json' in element '{"id": "resource_manager", "path": "nvflare.app_common.resource_managers.gpu_resource_manager.GPUResourceManager", "args": {"num_of_gpus": 1, "mem_per_gpu_in_GiB": 1}}': path: 'components.#1', exception: 'ValueError: num_of_gpus specified (1) exceeds available GPUs: 0.'

the num_of_gpus should be the gpu id or the number of gpu available? also, i try to seek ChatGPT for help, i get the answer is that this GPUResourceManager its GPU detection method is not suitable for Jetson, and ChatGPT help me to solve it but other method, then i can set num_of_gpus = 1 in my rescoure,json in each site, after that i attempt to submit the job again, i still get same result , no site have enough resource , this is because of why ?

these is meta.json
{
"name": "cifar10_fedavg_stream_tb_alpha1.0",
"resource_spec": {
"site-1": {
"num_of_gpus": 1,
"mem_per_gpu_in_GiB": 1
},
"site-2": {
"num_of_gpus": 0,
"mem_per_gpu_in_GiB": 1
},
"site-3": {
"num_of_gpus": 0,
"mem_per_gpu_in_GiB": 1
},
"site-4": {
"num_of_gpus": 0,
"mem_per_gpu_in_GiB": 1
},
"site-5": {
"num_of_gpus": 0,
"mem_per_gpu_in_GiB": 1
},
"site-6": {
"num_of_gpus": 0,
"mem_per_gpu_in_GiB": 1
},
"site-7": {
"num_of_gpus": 0,
"mem_per_gpu_in_GiB": 1
},
"site-8": {
"num_of_gpus": 0,
"mem_per_gpu_in_GiB": 1
}
},
"deploy_map": {
"cifar10_fedavg_stream_tb": [
"@ALL"
]
},
"min_clients": 1
}

and these is resource.json

{
"format_version": 2,
"client": {
"retry_timeout": 30,
"compression": "Gzip"
},
"components": [
{
"id": "resource_manager",
"path": "nvflare.app_common.resource_managers.gpu_resource_manager.GPUResourceManager",
"args": {
"num_of_gpus": 1,
"mem_per_gpu_in_GiB": 8
}
},
{
"id": "resource_consumer",
"path": "nvflare.app_common.resource_consumers.gpu_resource_consumer.GPUResourceConsumer",
"args": {}
}

]
}

please help , tqvm

[YuanTingHsieh]

@Gozilasim thanks for the interest and the question.

So our GPUResourceManager implementation is using this method to determine number of GPU: https://github.com/NVIDIA/NVFlare/blob/main/nvflare/fuel/utils/gpu_utils.py#L60
which in the back is using "nvidia-smi" to find it.
When you type nvidia-smi in your Jetson environment what happened?
Did you install CUDA there?
Another thing you can try is just don't put all these components and leave "resource_spec" blank and then your job should run.

[Gozilasim]

@Gozilasim thanks for the interest and the question.

So our GPUResourceManager implementation is using this method to determine number of GPU: https://github.com/NVIDIA/NVFlare/blob/main/nvflare/fuel/utils/gpu_utils.py#L60 which in the back is using "nvidia-smi" to find it. When you type nvidia-smi in your Jetson environment what happened? Did you install CUDA there? Another thing you can try is just don't put all these components and leave "resource_spec" blank and then your job should run.

thanks for the reply @YuanTingHsieh, but i had resolved this problem last week, since jetson orin was designed using IGPU, therefore it doest not have pci bus, nvidia-smi is not work for it, i resolved this problem by overwrite my GPU detection code inside that file, for those who also faced same problem with nvidia-smi , i suggest to use pycuda.driver, but before this, please make sure cuda was installed in the device

for ur suggestion , i dont think it works, becuz my first idea is to skip the GPU detection and its memory by configure the code inside GPUResourceManager, then make the resource_spec all component = 0, then i able to connect between the server and client, but when i submit the job, it will become these
DefaultJobScheduler - INFO - [identity=secure_project, run=?]: Try to schedule job c8c22eab-9c9c-404a-9716-ae8389948583, get result: (not enough sites have enough resources (ok sites 0 < min sites 1))

So for my suggestion, modify this file is better
https://github.com/NVIDIA/NVFlare/blob/main/nvflare/fuel/utils/gpu_utils.py#L60

anyone need the code can leave the comment

[YuanTingHsieh]

Thanks for the reply!
We will be happy if you want to contribute your detection logic back as well!

If you make all resource_spec empty for the job's meta.json:

{
"name": "cifar10_fedavg_stream_tb_alpha1.0",
"resource_spec": {
},
"deploy_map": {
"cifar10_fedavg_stream_tb": [
"@ALL"
]
},
"min_clients": 1
}

And leave the "resource.json" at each site empty:

{
"format_version": 2,
"client": {
"retry_timeout": 30,
"compression": "Gzip"
},
"components": [
{
"id": "resource_manager",
"path": "nvflare.app_common.resource_managers.gpu_resource_manager.GPUResourceManager",
"args": {
}
},
{
"id": "resource_consumer",
"path": "nvflare.app_common.resource_consumers.gpu_resource_consumer.GPUResourceConsumer",
"args": {}
}

]
}

would this work for you?

[Gozilasim]

@YuanTingHsieh sorry , out of the topic, do you have any real-world project that using tensorflow framework with transfer learning, means the tensorflow network structure is crated by pretrained model + additional layer . i dont know how to implement it .

[YuanTingHsieh]

@Gozilasim I think you can have the pre-trained model load in Server side using something similar to https://github.com/NVIDIA/NVFlare/blob/main/nvflare/app_common/np/np_model_persistor.py#L59 or https://github.com/NVIDIA/NVFlare/blob/main/nvflare/app_opt/tf/model_persistor.py#L26

Depends on whether you want to fine-tune the whole pre-trained model or just the additional layer, you might need different model update approaches.

@holgerroth @ZiyueXu77 can you help answer this question, thanks!

[holgerroth]

Hi @Gozilasim, we have some more advanced tf examples here https://github.com/NVIDIA/NVFlare/tree/main/examples/advanced/job_api/tf but it doesn't cover fine-tuning. For that you need to follow the TF documentation https://www.tensorflow.org/tutorials/images/transfer_learning.

[Gozilasim]

@holgerroth tq for the references, but i am doing on real world deployment, not sure the reference is suitable for me, because i still a newbie here

i tried to implement FL with tensorflow by referring to the cifar10-real-world. It seems there are a lot of mistake, but the implementation still works. However, the model's performance after implementing FL was poor. may you help me to review my code to ensure i am doing on the right way ?

training.py

`

from custom.utils.data_splitter import load_dataset, data_split
from custom.networks.net import CustomNN

import os
import numpy as np
import tensorflow as tf
from tensorflow.keras import optimizers, losses
from tensorflow.keras.callbacks import TensorBoard
from typing import Union
import copy

from nvflare.apis.fl_constant import FLMetaKey, ReturnCode
from nvflare.app_common.abstract.fl_model import FLModel, ParamsType
from nvflare.app_common.app_constant import AppConstants, ModelName, ValidateType
from nvflare.app_common.abstract.model_learner import ModelLearner
from nvflare.app_common.utils.fl_model_utils import FLModelUtils

class CustomModelLearner(ModelLearner):
def init(
self,
dataset_path: str = "/mnt/5958b632-9b58-4b1e-9b06-aaa3d99c9dcd/cifar10/cnn/fire_dataset/client1",
aggregation_epochs: int = 1,
lr: float = 1e-2,
fedproxloss_mu: float = 0.0,
central: bool = False,
analytic_sender_id: str = "analytic_sender",
batch_size: int = 64,
num_workers: int = 0,
):
super().init()
self.dataset_path = dataset_path
self.aggregation_epochs = aggregation_epochs
self.lr = lr
self.fedproxloss_mu = fedproxloss_mu
self.best_acc = 0.0
self.central = central
self.batch_size = batch_size
self.num_workers = num_workers
self.analytic_sender_id = analytic_sender_id

    # Epoch counter
    self.epoch_of_start_time = 0
    self.epoch_global = 0

    # Initialize variables that will be created later
    self.local_model_file = None
    self.best_local_model_file = None
    self.writer = None
    self.device = None
    self.model = None
    self.optimizer = None
    self.loss_fn = None
    self.train_dataset = None
    self.valid_dataset = None
    self.train_loader = None
    self.valid_loader = None

def initialize(self):
    """Initialize the necessary variables and TensorFlow components."""
    self.info(
        f"Client {self.site_name} initialized at \n {self.app_root} \n with args: {self.args}",
    )

    self.local_model_file = os.path.join(self.app_root, "local_model.h5")
    self.best_local_model_file = os.path.join(self.app_root, "best_local_model.h5")

    # Replace TensorBoard with tf.summary.create_file_writer
    self.writer = tf.summary.create_file_writer(self.app_root)  # Use TensorFlow's summary writer

    # Set training-related parameters
    #self.device = "/GPU:0" if tf.config.list_physical_devices('GPU') else "/CPU:0"
    self.device = tf.config.experimental.list_physical_devices("GPU")
    for device in self.device:
        tf.config.experimental.set_memory_growth(device,True)

    # Use the custom model provided
    self.model = CustomNN(input_shape=(224, 224, 3)) 
    self.model.pretrained_model.trainable = False
    # Ensure only the custom layers are trainable
      # Ensure only the custom layers are trainable
    for layer in self.model.custom_layers.layers:
        layer.trainable = True

    for layer in self.model.layers:
         print(f"Layer {layer.name} trainable: {layer.trainable}")

    self.model.build(input_shape=(None,224, 224, 3))

    self.optimizer = optimizers.SGD(learning_rate=self.lr, momentum=0.9)
    self.loss_fn = losses.SparseCategoricalCrossentropy(from_logits=True)
    self.model.summary()

def _create_datasets(self):
    """To be called after the dataset has been loaded and split."""
    if self.train_dataset is None or self.train_loader is None:
        # Since we already have a function to load and split the dataset, use it here
        all_images, all_labels = load_dataset(self.dataset_path)  # Use your dataset loading function
        print(f"Loaded dataset with images: {all_images.shape}, labels: {all_labels.shape}")
    
        (train_images, train_labels), (test_images, test_labels) = data_split(all_images, all_labels)
        print(f"Train images: {train_images.shape}, Train labels: {train_labels.shape}")
        print(f"Test images: {test_images.shape}, Test labels: {test_labels.shape}")


        # Create the TensorFlow dataset loaders
        self.train_dataset = tf.data.Dataset.from_tensor_slices((train_images, train_labels))
        self.train_loader = self.train_dataset.batch(self.batch_size).shuffle(1000)

        self.valid_dataset = tf.data.Dataset.from_tensor_slices((test_images, test_labels))
        self.valid_loader = self.valid_dataset.batch(self.batch_size)



def finalize(self):
    # collect threads, close files here
    pass

def local_train(self, train_loader, model_global, val_freq: int = 0):
    """Train the model locally using TensorFlow."""
    for epoch in range(self.aggregation_epochs):
        self.model.trainable = True  # Ensure the model is trainable

        epoch_len = len(train_loader)
        self.epoch_global = self.epoch_of_start_time + epoch
        self.info(f"Local epoch {self.site_name}: {epoch + 1}/{self.aggregation_epochs} (lr={self.lr})")
        avg_loss = 0.0
        for inputs, labels in train_loader:
            with tf.GradientTape() as tape:
                outputs = self.model(inputs, training = True)
                loss = self.loss_fn(labels, outputs)
                # FedProx loss term
                if self.fedproxloss_mu > 0:
                    fed_prox_loss = self.calculate_fedprox_loss(self.model,model_global)
                    loss += fed_prox_loss

            grads = tape.gradient(loss, self.model.trainable_variables)
            self.optimizer.apply_gradients(zip(grads, self.model.trainable_variables))

            avg_loss += loss.numpy()

        with self.writer.as_default():  # Log the scalar using TensorFlow's writer
            tf.summary.scalar("train_loss", avg_loss / len(train_loader), step=self.epoch_global)


        if val_freq > 0 and epoch % val_freq == 0:
            acc = self.local_valid(self.valid_loader, tb_id="val_acc_local_model")
            if acc > self.best_acc:
                self.best_acc = acc
                self.save_model(is_best=True)



def calculate_fedprox_loss(self, model_global):
    """Calculate the FedProx loss (if necessary)."""
    fed_prox_loss = 0
    for local_param, global_param in zip(self.model.trainable_weights, model_global.trainable_weights):
        fed_prox_loss += tf.reduce_sum(tf.square(local_param - global_param))
    return self.fedproxloss_mu * fed_prox_loss


def save_model(self, is_best=False):
    """Save the model weights along with additional metadata."""
    
    # Save model weights (in TensorFlow/Keras format)
    model_weights = self.model.get_weights()  # This fetches the weights of the model
    save_dict = {"model_weights": model_weights, "epoch": self.epoch_global}

    # Save the best model weights (if is_best is True) or normal model weights
    if is_best:
        save_dict.update({"best_acc": self.best_acc})
        # Save the model weights using TensorFlow's save_weights method
        self.model.save_weights(self.best_local_model_file)
        self.info(f"Best model saved at {self.best_local_model_file} with accuracy {self.best_acc:.4f}")
    else:
        # Save the model weights using TensorFlow's save_weights method
        self.model.save_weights(self.local_model_file)
        self.info(f"Model saved at {self.local_model_file} for epoch {self.epoch_global}")


def train(self, model: FLModel) -> Union[str, FLModel]:
    self._create_datasets()

    self.info(f"Current/Total Round: {self.current_round + 1}/{self.total_rounds}")
    self.info(f"Client identity: {self.site_name}")

    global_weights = model.params
    weight_values = [global_weights[name] for name in sorted(global_weights.keys())]
    # local steps
    epoch_len = len(self.train_loader)
    self.info(f"Local steps per epoch: {epoch_len}")

    # Separate the weights: pretrained and custom layers
    pretrained_len = len(self.model.pretrained_model.weights)
    custom_layer_weights = weight_values[pretrained_len:]  # Skip pretrained weights

    print("__________________________________________")


    # Initialize flat_weights with the custom layer weights
    flat_weights = np.concatenate([w.flatten() for w in custom_layer_weights])

    # Iterate through layers and assign weights
    for layer_index, layer in enumerate(self.model.custom_layers.layers):
        # Skip layers that don't have weights
        if not layer.get_weights():
            print(f"Layer {layer.name} has no weights. Skipping.")
            continue

        if layer_index >= len(custom_layer_weights):
            break

        # Fetch expected weight shapes
        kernel_shape = layer.get_weights()[0].shape
        has_bias = len(layer.get_weights()) > 1
        bias_shape = layer.get_weights()[1].shape if has_bias else None

        # Extract and reshape kernel weights
        kernel_size = np.prod(kernel_shape)
        reshaped_kernel = flat_weights[:kernel_size].reshape(kernel_shape)

        # Handle bias if present
        if has_bias:
            reshaped_bias = flat_weights[kernel_size:kernel_size + np.prod(bias_shape)].reshape(bias_shape)
            layer.set_weights([reshaped_kernel, reshaped_bias])
        else:
            layer.set_weights([reshaped_kernel])

        # Update flat_weights by removing the part used
        flat_weights = flat_weights[kernel_size + (np.prod(bias_shape) if has_bias else 0):]


    model_global = copy.deepcopy(self.model)


    # Training
    self.local_train(train_loader=self.train_loader, model_global=model_global, val_freq=1 if self.central else 0)
    self.epoch_of_start_time += self.aggregation_epochs

    acc = self.local_valid(self.valid_loader, tb_id="val_acc_local_model")
    self.info(f"val_acc_local_model: {acc:.4f}")

    self.save_model(is_best=False)
    if acc > self.best_acc:
        self.best_acc = acc
        self.save_model(is_best=True)

    # Compute and return deltas only for custom layers
    local_weights = {layer.name: layer.get_weights() for layer in self.model.custom_layers.layers}
    model_diff = {}

    for name, global_weight in global_weights.items():
        if name not in local_weights:
            continue  # Skip layers that aren't part of the local model

        local_layer_weights = local_weights[name]
        for i, local_weight in enumerate(local_layer_weights):
            # Ensure global weights are compatible
            if local_weight.shape != global_weight.shape:
                # Reshape if total size matches
                if np.prod(local_weight.shape) == np.prod(global_weight.shape):
                    print(f"Fixing shape mismatch for {name}: reshaping global weights")
                    global_weight = global_weight.reshape(local_weight.shape)
                else:
                    print(f"Skipping {name} due to size mismatch: local size {local_weight.size}, global size {global_weight.size}")
                    continue

            # Compute the difference
            model_diff[f"{name}_{i}"] = np.subtract(local_weight, global_weight, dtype=np.float32)

            # Check for NaNs
            if np.any(np.isnan(model_diff[f"{name}_{i}"])):
                self.stop_task(f"{name} weights became NaN at index {i}...")
                return ReturnCode.EXECUTION_EXCEPTION


    fl_model =FLModel(
        params_type=ParamsType.DIFF,
        params=model_diff
    )

    FLModelUtils.set_meta_prop(fl_model, FLMetaKey.NUM_STEPS_CURRENT_ROUND, epoch_len)
    self.info("Local epochs finished. Returning FLModel")
    return fl_model



    # Training

    self.local_train(train_loader=self.train_loader, model_global=self.model, val_freq=1 if self.central else 0)

    acc = self.local_valid(self.valid_loader, tb_id="val_acc_local_model")
    self.save_model(is_best=(acc > self.best_acc))

    # Compute and return deltas only for custom layers
    local_custom_weights = [layer.get_weights()[0] for layer in self.model.custom_layers.layers if layer.get_weights()]

    model_diff = {}

    for i, local_weight in enumerate(local_custom_weights):
        global_weight = custom_layer_weights[i]
        print(f"Global weight shape: {global_weight.shape}, Local weight shape: {local_weight.shape}")

        if local_weight.shape != global_weight.shape:
            # Reshape if needed
            if global_weight.size == local_weight.size:
                global_weight = global_weight.reshape(local_weight.shape)
            else:
                print(f"Skipping reshape: global_weight.size={global_weight.size}, local_weight.size={local_weight.size}")

        model_diff[f"custom_layer_{i}"] = local_weight - global_weight

    return FLModel(params_type=ParamsType.DIFF, params=model_diff)


def get_model(self, model_name: str) -> Union[str, FLModel]:
    # Retrieve the best local model saved during training.
    if model_name == ModelName.BEST_MODEL:
        try:
            # Load the saved model weights
            model = CustomNN(input_shape=(224, 224, 3)) 
            model.build(input_shape=(None,224, 224, 3))
            model.load_weights(self.best_local_model_file)
        except Exception as e:
            raise ValueError("Unable to load best model") from e

        # Extract weights and convert to numpy for FLModel compatibility
        model_weights = {layer.name: layer.get_weights() for layer in model.layers if layer.weights}
        numpy_weights = {}
        for layer_name, weights in model_weights.items():
            numpy_weights[layer_name] = [w for w in weights]  # Already NumPy arrays in TensorFlow

        # Return as FLModel
        return FLModel(params_type=ParamsType.FULL, params=numpy_weights)
    else:
        raise ValueError(f"Unknown model_type: {model_name}")  # Raised errors are caught in LearnerExecutor class.




def local_valid(self, valid_loader, tb_id=None):
    correct, total = 0, 0

    for inputs, labels in valid_loader:
        inputs = tf.convert_to_tensor(inputs)
        labels = tf.convert_to_tensor(labels)

        # Perform inference
        outputs = self.model(inputs, training=False)

        # Get predicted labels
        pred_label = tf.argmax(outputs, axis=1)

        # Ensure both labels and pred_label are of the same dtype
        pred_label = tf.cast(pred_label, tf.int32)  # Cast to int32 to match labels
        labels = tf.cast(labels, tf.int32)

        # Update total and correct predictions
        total += tf.shape(inputs)[0].numpy()
        correct += tf.reduce_sum(tf.cast(pred_label == labels, tf.int32)).numpy()

    # Calculate the accuracy metric
    metric = correct / float(total)

    if tb_id:
    # Correct way to log scalar value
        with self.writer.as_default():
            tf.summary.scalar(tb_id, metric, step=self.epoch_global)

    return metric


def validate(self, model: FLModel) -> Union[str, FLModel]:
    self._create_datasets()
    self.info(f"Client identity: {self.site_name}")

    global_weights = model.params
    weight_values = [global_weights[name] for name in sorted(global_weights.keys())]

    # Only assign weights to custom layers
    pretrained_len = len(self.model.pretrained_model.weights)
    custom_layer_weights = weight_values[pretrained_len:]
    print("\nlen\n",len(custom_layer_weights) ,"\n")

    if len(custom_layer_weights) == len(self.model.custom_layers.weights):
        self.model.custom_layers.set_weights(custom_layer_weights)
    else:
        raise ValueError(f"Mismatch in custom layers: expected {len(self.model.custom_layers.weights)}, but got {len(custom_layer_weights)}")

    validate_type = FLModelUtils.get_meta_prop(
    model, FLMetaKey.VALIDATE_TYPE, ValidateType.MODEL_VALIDATE
    )  # TODO: enable model.get_meta_prop(...)
    model_owner = self.get_shareable_header(AppConstants.MODEL_OWNER)

    train_acc = self.local_valid(self.train_loader, 
                                 tb_id="train_acc_global_model" if validate_type == ValidateType.BEFORE_TRAIN_VALIDATE else None,
    )
    self.info(f"Training accuracy ({model_owner}): {train_acc:.4f}")

    val_acc = self.local_valid(self.valid_loader, 
                               tb_id="val_acc_global_model" if validate_type == ValidateType.BEFORE_TRAIN_VALIDATE else None,
    )
    self.info(f"Validation accuracy ({model_owner}): {val_acc:.4f}")
    self.info("Evaluation finished. Returning result")
    
    if val_acc > self.best_acc:
        self.best_acc = val_acc
        self.save_model(is_best=True)

    val_results = {"train_accuracy": train_acc, "val_accuracy": val_acc}
    return FLModel(metrics=val_results)




def calculate_fedprox_loss(self, model_global):
    """Calculate the FedProx loss (if necessary)."""
    fed_prox_loss = 0
    for local_param, global_param in zip(self.model.trainable_weights, model_global.trainable_weights):
        fed_prox_loss += tf.reduce_sum(tf.square(local_param - global_param))
    return self.fedproxloss_mu * fed_prox_loss

net.py

import tensorflow as tf
from tensorflow.keras import layers, models

class CustomNN(models.Model):
def init(self, input_shape=None):
super(CustomNN, self).init()

    # Use a different name to avoid conflict with TensorFlow's input_shape
    self._input_shape = input_shape if input_shape is not None else (224, 224, 3)

    # Define MobileNetV2 as the base model with the given input_shape
    self.pretrained_model = tf.keras.applications.MobileNetV2(
        weights='imagenet', 
        input_shape=self._input_shape, 
        include_top=False,
        pooling='avg'  # Global Average Pooling handled by MobileNetV2 itself
    )
    self.pretrained_model.trainable = False
    # Define the custom layers after the global pooling layer
    self.custom_layers = models.Sequential([
        layers.Dense(256, activation='relu'),
        layers.Dropout(0.2),
        layers.Dense(256, activation='relu'),
        layers.Dropout(0.2),
        layers.Dense(2, activation='softmax')  # Output layer for binary classification
    ])
    
def call(self, inputs):
    # Resize and rescale input images before passing to the pretrained model
    resize_and_rescale = models.Sequential([
        layers.experimental.preprocessing.Resizing(224, 224),
        layers.experimental.preprocessing.Rescaling(1./255),
    ])

    # Apply resizing and rescaling
    x = resize_and_rescale(inputs)
    
    # Pass input through the pretrained MobileNetV2 model
    x = self.pretrained_model(x)

    # Pass the output through the custom layers
    x = self.custom_layers(x)

    return x

`

config_fed_server.json

`
{
"format_version": 2,
"min_clients": 2,
"num_rounds": 20,
"alpha": 1.0,
"server": {
"heart_beat_timeout": 600
},
"task_data_filters": [],
"task_result_filters": [],
"components": [
{
"id": "fire-learner",
"path": "custom.learners.training.CustomModelLearner",
"args": {
"aggregation_epochs": 4,
"lr": 0.01
}
},
{
"id": "serialize_filter",
"path": "nvflare.app_opt.he.model_serialize_filter.HEModelSerializeFilter",
"args": {}
},
{
"id": "persistor",
"path": "nvflare.app_opt.tf.model_persistor.TFModelPersistor",
"args": {
"model": {
"path": "custom.networks.net.CustomNN",
"args": {}
},
"filter_id": "serialize_filter"
}
},
{
"id": "shareable_generator",
"path": "nvflare.app_common.shareablegenerators.full_model_shareable_generator.FullModelShareableGenerator",
"args": {}
},
{
"id": "aggregator",
"path": "nvflare.app_opt.he.intime_accumulate_model_aggregator.HEInTimeAccumulateWeightedAggregator",
"args": {
"weigh_by_local_iter": false
}
},
{
"id": "model_selector",
"path": "nvflare.app_common.widgets.intime_model_selector.IntimeModelSelector",
"args": {}
},
{
"id": "model_locator",
"path": "nvflare.app_common.model_locator.list_model_locator.ListModelLocator",
"args": {
"persistor_id": "persistor"
}
},
{
"id": "json_generator",
"path": "nvflare.app_common.widgets.validation_json_generator.ValidationJsonGenerator",
"args": {}
}
],
"workflows": [
{
"id": "scatter_gather_ctl",
"path": "nvflare.app_common.workflows.scatter_and_gather.ScatterAndGather",
"args": {
"min_clients": "{min_clients}",
"num_rounds": "{num_rounds}",
"start_round": 0,
"wait_time_after_min_received": 10,
"aggregator_id": "aggregator",
"persistor_id": "persistor",
"shareable_generator_id": "shareable_generator",
"train_task_name": "train",
"train_timeout": 0
}
},
{
"id": "cross_site_model_eval",
"path": "nvflare.app_opt.he.cross_site_model_eval.HECrossSiteModelEval",
"args": {
"model_locator_id": "model_locator",
"submit_model_timeout": 600,
"validation_timeout": 6000,
"cleanup_models": true
}
}
]

}
config_fed_client.json{
"format_version": 2,
"AGGREGATION_EPOCHS": 4,
"executors": [
{
"tasks": [
"train",
"submit_model",
"validate"
],
"executor": {
"id": "Executor",
"path": "nvflare.app_common.executors.model_learner_executor.ModelLearnerExecutor",
"args": {
"learner_id": "fire-learner"
}
}
}
],
"task_result_filters": [
{
"tasks": [
"train"
],
"filters": [
{
"path": "nvflare.app_opt.he.model_encryptor.HEModelEncryptor",
"args": {
"weigh_by_local_iter": true
}
}
]
},
{
"tasks": [
"submit_model"
],
"filters": [
{
"path": "nvflare.app_opt.he.model_encryptor.HEModelEncryptor",
"args": {
"weigh_by_local_iter": false
}
}
]
}
],
"task_data_filters": [
{
"tasks": [
"train",
"validate"
],
"filters": [
{
"path": "nvflare.app_opt.he.model_decryptor.HEModelDecryptor",
"args": {}
}
]
}
],
"components": [
{
"id": "fire-learner",
"path": "custom.learners.training.CustomModelLearner",
"args": {
"aggregation_epochs": "{AGGREGATION_EPOCHS}",
"lr": 0.01
}
}
]
}
`