How to Create Your Custom PyTorch Training Plan¶

Fed-BioMed allows you to perform model training without changing your PyTorch training plan class completely. Integrating your PyTorch model to Fed-BioMed only requires to add extra attributes and methods to train your model based on a federated approach. In this tutorial, you will learn how to write/define your TrainingPlan (wrapping your model) using the Fedbiomed class CustomDataset in Fed-BioMed for PyTorch framework.

Note: Before starting this tutorial we highly recommend you to follow the previous tutorials to understand the basics of Fed-BioMed.

In this tutorial, we will be using Celaba (CelebaFaces) dataset to train the model. You can see details of the dataset here. In the following sections, you will have the instructions for downloading and configuring Celeba dataset for Fed-BioMed framework.

1. Fed-BioMed Training Plan and Custom Dataset¶

In this section, you will learn how to write your custom training plan.

What is Training Plan?¶

The training plan is the class where all the methods and attributes are defined to train your model on the nodes. Each training plan should inherit the base training plan class of the belonging ML framework that is provided by Fed-BioMed. For more details, you can visit documentation for training plan. The following code snippet shows a basic training plan that can be defined in Fed-BioMed for PyTorch framework.

from fedbiomed.common.training_plans import TorchTrainingPlan


class CustomTrainingPlan(TorchTrainingPlan):
    def init_model(self, model_args):
        # Define here your model
        # ...
        return

    def init_dependencies(self):
        # Add here the dependencies / third party libraries to be loaded
        #...
        return
    
    def init_optimizer(self, optimizer_args):
        # Define here your optimizer
        #...
        return

    def training_data(self,  batch_size = 48):
        # Define here how data are processed before feeding it to the model
        # ...
        return
    
    def training_step(self, data, target):
        # Define here the loss function
        # ...
        return

What is Custom Dataset?¶

Custom Dataset is a Fedbiomed class that allows you to register any file or folder to your node as a dataset. As a user, you define your own custom read and item retrieval logic.

As a user, you:

• Implement:

  • read(self): load and preprocess data once (e.g., read CSV, build arrays) and store into instance fields.
  • __len__(self): number of samples.
  • get_item(self, index): return exactly a tuple (data, target) for the given index. and

• Do not override __init__ or __getitem__. These are managed by Fed‑BioMed so it can apply format checks and safety validations.

• target can be None in unsupervised settings. In supervised tasks (like this one), return a proper target (float, int, tensor, etc.).

Good practice:

• Prepare features as float arrays/tensors
• Keep get_item focused and fast; heavy work belongs in read.

For more information, please refer to Custom Dataset for documentation.

init_model Method of Training Plan¶

init_model method of the training plan is where you initialize your neural network module as in classical PyTorch model class. The network should be defined inside the training plan class and init_model should instantiate this network (Module), and return it.

In this tutorial, we will be training a classification model for CelebA image dataset that will be able to predict whether the given face is smiling.

def init_model(self, model_args: dict = {}):
    return self.Net(model_args)

class Net(nn.Module):

    def __init__(model_args):
        super().__init__()
        # Convolutional layers
        self.conv1 = nn.Conv2d(3, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 32, 3, 1)
        self.conv3 = nn.Conv2d(32, 32, 3, 1)
        self.conv4 = nn.Conv2d(32, 32, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        # Classifier
        self.fc1 = nn.Linear(3168, 128)
        self.fc2 = nn.Linear(128, 2)

    def forward(self, x):

        x = self.conv1(x)
        x = F.max_pool2d(x, 2)
        x = F.relu(x)

        x = self.conv2(x)
        x = F.max_pool2d(x, 2)
        x = F.relu(x)

        x = self.conv3(x)
        x = F.max_pool2d(x, 2)
        x = F.relu(x)

        x = self.conv4(x)
        x = F.max_pool2d(x, 2)
        x = F.relu(x)

        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)

        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output

init_dependencies Method¶

Next, you should define the init_dependencies to declare the modules that are used in the training plan. The modules should be supported by the Fed-BioMed.

def init_depedencies(self)
    # Here we define the custom dependencies that will be needed by our custom Dataloader
    deps = ["from torch.utils.data import Dataset, DataLoader",
            "from torchvision import transforms",
            "import pandas as pd",
            "from PIL import Image",
            "import os",
            "import numpy as np"]
    return deps

init_optimizer Method¶

To optimize your model, you will need an optimizer. This is where init_optimizer method comes into play. In this method, you may change the optimizer you want to use, add a PyTorch learning rate Scheduler or provide your custom optimizer.

init_optimizer takes optimizer_args as argument, an entry from training_args (more details later) which is a dictionary containing parameters that may be needed for optimizer initialization (such as learning rate, Adagrad weights decay, Adam beta parameters, ...). init_optimizer method should return the initialized optimizer, that will be used to optimize model.

Defining an optimizer in Fed-BioMed is pretty similar to PyTorch, as shown in the example below (using PyTorch's SGD optimizer):

def init_optimizer(self, optimizer_args):
    return torch.optim.SGD(self.model().parameters(), lr=optimizer_args['lr'])

By default (if this method is not specified in the TrainingPlan), model will be optimized using default Adam optimizer.

training_data() and Custom Dataset¶

training_data() is the method responsible for loading the data on the node side during each round of federated training. When a training round starts, each participating node automatically calls this method to:

• Build the model,
• Load the dataset through the training_data() method, and
• Perform the training_step() using the loaded data.

In cases where you need to load data from custom formats or folder structures, you can create your own dataset class by subclassing CustomDataset. This base class provides the interface for defining how data should be read and accessed on each node. To make it compatible with the federated workflow, your subclass must implement three methods:

Method Purpose read() Loads your dataset into memory and prepares it for training (e.g., by reading CSVs, parsing image folders, etc.). get_item() Returns a single (data, target) pair given an index. __len__() Returns the total number of samples in the dataset.

ℹ️ Info: self.path Variable

In Fed-BioMed, when you add a dataset to a node (via fedbiomed node dataset add), the node saves the dataset’s root folder in its local database, through the path variable you enter. Later, when a training job needs the dataset, the node resolves that stored location and passes it into your CustomDataset through the self.path variable.

For example, when working with an image dataset such as CelebA, you could implement a CustomDataset subclass like this:

class CelebaDataset(CustomDataset):
    """Custom Dataset for loading CelebA face images"""

    def read(self):
        self.txt_path = self.path + '/' + TEXT_FILE_NAME
        self.img_dir = self.path + '/' + IMAGE_DIRECTORY_NAME

        # Read the CSV file that contains labels for each image
        df = pd.read_csv(self.txt_path, sep="\t", index_col=0)
        self.img_names = df.index.values
        self.y = df['Smiling'].values

    def get_item(self, index):
        img = np.asarray(Image.open(os.path.join(self.img_dir, self.img_names[index])))
        img = transforms.ToTensor()(img)
        label = self.y[index]
        return img, label

    def __len__(self):
        return self.y.shape[0]

This approach allows you to flexibly define how your dataset is structured and loaded, while Fed-BioMed ensures that your implementation integrates seamlessly with the federated workflow. When the node participates in training, the dataset you defined is automatically initialized and used in the training_data() method.

Later in the tutorial, we will explain how datasets like this are added to nodes using the CLI command fedbiomed node dataset add, which takes care of linking the local data path to the node configuration.

Now, you need to define a training_data method that will create a Fed-BioMed DataManager using custom CelebaDataset class.

    def training_data(self):
        # The training_data creates the dataset and returns DataManager to be used for training in the general class Torchnn of Fed-BioMed
        dataset = self.CelebaDataset(
            os.path.join(self.dataset_path, "target.csv"), os.path.join(self.dataset_path, "data")
            )
        loader_arguments = { 'shuffle': True}
        return DataManager(dataset, **loader_arguments)

training_step()¶

The last method that needs to be defined is the training_step. This method is responsible for executing the forward method and calculating the loss value for the backward process of the network. To access the forward method of the torch.nn.Module that is defined in the init_model, the getter method model() of training plan class should be used.

def training_step(self, data, target): 
    output = self.model().forward(data)
    loss   = torch.nn.functional.nll_loss(output, target)
    return loss

You are now ready to create your training plan class. All you need to do is to locate every method that has been explained in the previous sections in your training plan class. In the next steps we will;

1. launch the Researcher: fedbiomed researcher start
2. download the CelebA dataset and deploy it on the nodes
3. define our complete training
4. create an experiment and run it
5. evaluate our model using a testing dataset

2.Configuring Nodes¶

We will be working with CelebA (CelebFaces) dataset. Therefore, please visit here and download the files img/img_align_celeba.zip and Anno/list_attr_celeba.txt. After the download operation is completed;

• Please go to fbm-researcher/notebooks/data/Celeba from the folder your are running the Fed-BioMed Researcher.
• Create Celeba_raw/raw directory and copy the list_attr_celeba.txt file.
• Extract the zip file img_align_celeba.zip

Your folder should be same as the tree below;

Celeba
    README.md
    create_node_data.py    
    .gitignore
 
    Celeba_raw
        raw
            list_attr_celeba.txt
            img_align_celeba.zip
            img_align_celeba
              lots of images 

The dataset has to be processed and split to create three distinct datasets for Node 1, Node 2, and Node 3. You can do it easily by running the following script in your notebook. If you are working in a different directory than the fbm-researcher/notebooks, please make sure that you define/modify the correct paths in the following example.

Running the following scripts might take some time, please be patient.

Now if you go to the fdm-researcher/notebooks/data/Celaba directory you can see the folder called celeba_preprocessed. There will be three different folders that contain an image dataset for 3 nodes. The next step will be configuring the nodes and deploying the datasets. In the next steps, we will be configuring only two nodes. The dataset for the third node is going to be used for the testing.

Create 2 nodes for training :

• fedbiomed component create -c node --path ./my-first-node
• fedbiomed component create -c node --path ./my-second-node

Note: You may want to copy the dataset created into Nodes, in the folder my-first-node/data or ./my-second-node/data . This is not required for Fed-BioMed to work, since you can access dataset from outside Node component, but enables you to run the tutorials more easily. Commands for that will be:

• cp -r fbm-researcher/notebooks/data/Celeba/celeba_preprocessed/data_node_1/ my-first-node/data/
• cp -r fbm-researcher/notebooks/data/Celeba/celeba_preprocessed/data_node_2/ my-second-node/data/

Add data to each node :

• fedbiomed node --path ./my-first-node dataset add
• fedbiomed node --path ./my-second-node dataset add

2.1. Configuration Steps¶

It is necessary to previously configure at least a node:

1. fedbiomed node --path ./my-first-node dataset add or fedbiomed node --path ./my-second-node dataset add

• For each Node, select option 6 (Custom) to add your custom dataset to the node
• Add a name and the tag for the dataset (tag should contain '#celeba' as it is the tag used for this training). Add the description
• Enter the path for your dataset location. You can choose from the 3 generated datasets inside fbm-researcher/notebooks/data/Celeba/celeba_preprocessed (eg: data_node_1) or directly in the Node data folder. Path can be relative or absolute.
• Data must have been added (if you get a warning saying that data must be unique is because it's been already added)

2. Check that your data has been added by executing fedbiomed node --path <component-directory> dataset list

3. Run the node using fedbiomed node --path <component-directory> start. Wait until you get Starting task manager. it means you are online.

After the steps above are completed, you will be ready to train your classification model on two different nodes.

3. Defining Custom PyTorch Model and Training Plan¶

Next step is to create our Net class based on the methods that have been explained in the previous sections. This class is part of the training plan that will be passed to the Experiment. Afterwards, the nodes will receive the training plan and perform the training by retrieving training data and passing it to the training_step.

This group of arguments corresponds respectively to:

• model_args: a dictionary with the arguments related to the model (e.g. number of layers, features, etc.). This will be passed to the model class on the node-side.
• training_args: a dictionary containing the arguments for the training routine (e.g. batch size, learning rate, epochs, etc.). This will be passed to the routine on the node-side.

Note: Typos and/or lack of positional (required) arguments might raise an error.

4. Training Federated Model¶

To provide training orchestration over two nodes we need to define an experiment which:

• searches nodes serving data for the tags,
• defines the local training on nodes with the training plan saved in training_plan_path, and federates all local updates at each round with aggregator
• runs training for round_limit.

You can visit user guide to know much more about experiment.

Let's start the experiment.

By default, this function doesn't stop until all the round_limit rounds are done for all the nodes. While the experiment runs you can open the terminals where you have started the nodes and see the training progress. However, the loss values obtained from each node during the training will be printed as output in real time. Since we are working on an image dataset, training might take some time.

Save trained model to file

Loading Training Parameters¶

After all the rounds have been completed, you retrieve the aggregated parameters from the last round and load them.

5. Testing Federated Model¶

We will define a testing routine to extract the accuracy metrics on the testing dataset. We will use the dataset that has been extracted into data_node_3.

We also need to define a custom Dataset class for the test dataset in order to load it using PyTorch's DataLoader. This will be the same class that has been already defined in the training plan.

We will also show the alternative way to define a CustomDataset. Without a TrainingPlan and a DataManager, you can directly use complete_initialization by giving root to controller_kwargs to give a path and initialize your dataset.

Conclusions¶

In this tutorial, running a model with a Custom Dataset on Fed-BioMed for the PyTorch framework has been explained. Because the examples are designed for the development environment, we have been running nodes in the same host machine. In production, the nodes that you need to use to train your model will serve in remote servers. Please check out how to deploy Nodes in a production environment.