{'fraction_non_alphanumeric': 0.028302958042355077, 'fraction_numerical': 0.008062226764321808, 'mean_word_length': 4.908587722240859, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
TorchFL: A Performant Library for Bootstrapping Federated Learning Experiments

With the increased legislation around data privacy, federated learning (FL) has emerged as a promising technique that allows the clients (end-user) to collaboratively train deep learning (DL) models without transferring and storing the data in a centralized, third-party server. Despite the theoretical success, FL is yet to be adopted in real-world systems due to the hardware, computing, and various infrastructure constraints presented by the edge and mobile devices of the clients. As a result, simulated datasets, models, and experiments are heavily used by the FL research community to validate their theories and findings. We introduce TorchFL, a performant library for (i) bootstrapping the FL experiments, (ii) executing them using various hardware accelerators, (iii) profiling the performance, and (iv) logging the overall and agent-specific results on the go. Being built on a bottom-up design using PyTorch and Lightning, TorchFL provides ready-to-use abstractions for models, datasets, and FL algorithms, while allowing the developers to customize them as and when required.


Introduction
With the rapid advancement of sensing and computing capabilities, the amount of data generated from mobile (client) devices has exponentially increased in recent years [24]. The increased volume of data has enabled deep learning (DL) [21] to become a widely adopted technique to train the computational models on the users' data and actively learn from their browsing patterns. In addition to the availability of the data, the recent research advances in specialized hardware [20] have enabled us to achieve astonishing results in ad targeting, language translation, image generation, content recommendations, and a lot more problems that were difficult to solve without a neural network (NN) [8].
Given the magnitude of data and computational resources required to train these models, the user data is often sent and collected in a centralized server, and the trained model is deployed on the device. While this technique has been effective, the stricter legislation and norms around data privacy make it difficult to use the data for training the DL models [33]. As a result, federated learning (FL) has emerged as a promising technique for training models across multiple clients (e.g., edge or mobile devices) without requiring the exchange of locally stored data but only that of the parameters [25]. Despite being applied and having displayed initial success with the Google Keyboard [23], FL faces serious hardware and infrastructure challenges before it can provide the same results as the traditional DL techniques [16]. While edge and mobile devices of the clients are becoming more computationally powerful, there's a long way to go before we can train state-of-the-art DL models on these devices, or even run a training process in the background without impacting the performance, battery, or daily user activity [29].
While part of the research community is dealing with these hardware and infrastructure challenges, most of the community focuses on designing new sampling algorithms, incentive mechanisms, aggregation protocols, defense mechanisms, and all of the other components that are integral to the FL pipeline [18]. While working on these problems, it is often a common practice to simulate the FL experiments where the clients, datasets, and FL models are spawned on a single machine and the empirical results are collected. In this paper, we present TorchFL, a plug-and-play and performant library for bootstrapping simulated FL experiments.
In addition, the DL models are getting more complex, and constant innovations are being made on the hardware accelerators as well as the tools around the data collection and visualizations. TorchFL can be used to minimize the implementation overhead for the FL community by abstracting the hardware, infrastructure, data, and DL implementations to set up the experiments. Given that Python and PyTorch are the most popular languages and frameworks used for DL respectively [32], TorchFL is a Python library that is built on PyTorch and can be used by a developer to run an end-to-end experiment by providing the following features: 1. Wrappers for state-of-the-art DL models which can be trained in federated or non-federated settings, 2. Wrappers for the most commonly used state-of-the-art datasets and the ability to automatically create the data shards based on the FL configuration, 3. Added support for the fine-tuning or feature extraction from the pre-trained DL models, allowing for faster training using federated transfer learning, 4. Customizable FL layer with the ready-to-use implementation of FL clients, samplers, and aggregators, which can be used to quickly spawn the experiments using the configuration files, 5. Backward compatibility with the PyTorch Lightning loggers, profilers, hardware accelerators, and the latest DevOps tools to help avoid the implementation and performance overhead for recording and collecting the experimental results.
In addition to the aforementioned features, abstractions, and ease-of-use, TorchFL is implemented using a bottom-up approach as shown in Figure 2, which allows the developers to customize every layer and build on top of the library to validate their research and hypothesis.
The rest of the paper is organized as follows. In Section 2, we present background material on FL. Section 3 discusses the detailed explanation and design decisions for the TorchFL architecture. In Section 4, we run FL experiments using TorchFL and provide examples to demonstrate the abstraction, effectiveness, and customizability of the framework. We wrap up by reviewing the existing FL toolkits (Section 5) and discussing future work (Section 6).

Background on FL
FL is a machine-learning setting where many clients collaboratively train a model under the orchestration of a central server [15]. In this work, we focus on cross-device FL, where clients are a very large number of mobile or edge devices with private and locally stored training data. Given a learning objective, the goal of the server is to train a model optimizing for the aforementioned objective by performing updates and aggregation using clients' data without transferring or exchanging any data from the clients. This model can then be deployed on clients as well as new users. The lifecycle of a cross-device FL systems can be better understood using the visual representation provided in Figure 1 [15].
Mathematically, the goal in cross-device FL is to optimize a loss function L using a model M parameterized by a vector W M ∈ R d in T rounds via a training dataset D, which is distributed and privately stored among a set of K devices or agents A (i.e., D = D 1 ∪ D 2 ∪ ... ∪ D K ). At each round t ∈ {1, . . . , T }, a subset of agents, A t ⊆ A, is chosen by the server for training. This process is called sampling. The goal of each sampled agent i ∈ A T is to minimize the loss function L over its own privately stored data min-batch of data D t i at time t. The agents minimize the loss L by starting from the global weight vector W t M at time t and Figure 1: The lifecycle of a cross-device FL system with multiple client devices and a single server. The figure is adopted from [15].
running an algorithm such as stochastic gradient descent (SGD) [26]. At the end of the update, the agent i obtains a local weight vector W t+1 Mi using only their privately stored mini-batch of data samples, computes and sends it to the server. To update the global weight vector W t+1 M for the next round, an aggregation mechanism is used over all the collected parameter updates ∆ t+1 = {δ t+1 i } i∈A t . A commonly used aggregation mechanisms is weighted averaging [25]: where Γ t i is the non-negative weight assigned to agent i ∈ A t at time t and Σ i∈A t Γ t i = 1. In summary, the FL process consists of multiple components including but not limited to (i) data distribution among the agents, (i) global model selection, (iii) agent sampling, (iv) local model training and (v) aggregation mechanism to update the global model. A practical implementation of a FL system should allow different choices for each of these components. See Section 3 for more details.
Remark 1. The type of data distribution heavily impacts the convergence rate and the performance of FL systems [26]. There are two common data distributions: (1) Independent and Identically Distributed (IID) and (2) Non-Independent and Identically Distributed (Non-IID). Having IID data for the agents means that each batch of data used for a client's local update is statistically identical to a uniformly drawn sample (with replacement) from the entire training dataset D (which is the union of all local datasets for the K agents). In practice, it is unreasonable to assume that every agent has IID data, and hence, throughout the discussion, we heavily focus on the availability of the non-IID datasets.

Architecture of TorchFL
Cross-device FL is a multi-step learning process with a lot of implementation, hardware, communication, and systems overhead. Despite the tremendous theoretical success [26], research and work are being done to use FL for real-world applications on actual client devices. Yet, as constant research is being done to explore the non-systemic challenges in FL, TorchFL aims to provide an end-to-end toolkit for the researchers and developers to quickly bootstrap an FL experiment, orchestrate it on multiple devices, and execute it using the hardware accelerators. As shown in Figure 2, TorchFL is built using a bottom-up approach to allow the customizability of datasets, models, and the entire FL layer. TorchFL is primarily built using Python, PyTorch, and Lightning frameworks and is backward compatible with PyTorch Lightning loggers, profilers, and accelerators. In this section, we explain the motivation behind the design of TorchFL and describe each of the layers in TorchFL's architecture.

Datasets and Models Libraries
As shown in Equation 1, the local training updates computed by the selected agents are shared with the server to compute the global model parameters as shown in Equation 2. In this section, we dig deeper into the design, features, and abstractions offered by the datasets and DL model offerings in TorchFL. As these modules are the first step in setting up an FL experiment, they appear at the bottommost layer in Figure 2.  Table 1 and provides an endto-end pipeline for retrieval and distribution of the data. All of these data-related features have been wrapped under the datamodules module in the codebase and an easy-to-use abstract base class has also been provided for the developers to easily add new datasets. The design of the datamodules interface is explained using a UML diagram in Figure 3. Notably, the abstract base class is inherited and made backward compatible with PyTorch Lightning Datamodule to make full use of their infrastructure and loggers offerings. This design allows us to collect the metadata related to the data distribution alongside all the model-related information which we will see in future sections. As shown in Figure 3, the datamodules also provides methods for creating the IID and non-IID distributions for FL. In fact, all of the logic for these methods have been written such that it works out of the box for any new dataset. The non-IID dataloader also provides a parameter named niid factor which allows the developers to control the intensity of the distribution between the agents. Higher niid factor would mean a higher imbalance in the labels split between the agents. In summary, any developer willing to add a custom dataset must create a new class, inherit from the BaseDatamodule, and only override the relevant methods as required. Table 2: Collection of models currently supported by TorchFL and the availability of feature extraction and finetuning for every group.

Models Variants Feature Extraction
Fine Tuning 2. DL Models: Once the global dataset D has been set up and distributed among the set of K agents A, the next step in the FL process is to initialize a global model M which is maintained by a service provider, and assign individual models M i to every agent i. To serve both of these purposes, T orchF L contains a models module which contains an extensive library of state-of-the-art DL models as shown in Table 2. Every major model has various architectures, a popular example being ResNet [28] with variants like ResNet18, ResNet34, etc. Although T orchF L currently only supports the computer vision models, Figure 4 explains how the module has been designed with multiple layers of abstraction to prioritize extensibility and ease of adding new models in the future. When trying to add a new model, the developer needs to take the following steps: (i) define the core PyTorch logic at the bottommost layer, (ii) use the templating-based code generator to generate the boilerplate code to ensure compatibility with the datasets, (iii) use the entry points to pass the model and optimizer hyperparameters to generate a trainable model object, and (iv) use the generated models and existing datamodules to initialize a Lightning, compatible trainer. In addition to the ease of selecting and using state-of-the-art models, TorchFL also supports federated transfer learning [14]. Given the millions of trainable parameters, and the amount of time and hardware resources required to achieve convergence, it is unreasonable to expect mobile devices to be able to train these models from scratch [34]. As such, transfer learning [38] within FL systems is heavily being studied. The primary idea of transfer learning is to reuse the DL model trained on an existing task for a similar new task. It is analogous to expecting a human proficient at riding bikes to be able to apply or reuse balance and navigation skills while learning to ride motorbikes. TorchFL supports the following transfer learning paradigms: (a) Finetuning [2]: This approach requires the agents to start with the model parameters pre-trained on an existing task, and retrain them for the new task. Despite the complete retraining, we expect the overall training time to be lower as the initial parameters are not randomly initialized.
(b) Feature Extraction: [38]: This approach requires the agents to start with the pre-trained model parameters and only retrain the final classification layers instead of the entire model. This technique results in a substantial decrease in the number of trainable parameters and hence accounts for lower training time. Table 2 contains an archive of all the models available in TorchFL that support finetuning and feature extraction. Relevant hyperparameters have been provided in the model entry points to allow the developers to retrieve a finetuned or feature-extracted model on the ImageNet dataset [6]. In Section 4, we demonstrate how transfer learning in FL can help drastically reduce the trainable parameters and training time while achieving optimal performance.

Federated Module
In this section, we discuss how to FL modules are designed to work with the datasets and models while being backward compatible with various PyTorch Lightning utilities.

Agents:
The primary entity of any FL system is an agent, which is synonymous with a client, mobile, or edge device in the real world. In most of the simulated FL experiments, agents are usually treated as a collection of uniquely identifiable integers which are subject to random sampling and selection for every training round. However, as extensive research is currently being done in areas like reputationbased sampling protocols [35], robust defense mechanisms against model poisoning [18], and gametheoretic incentive systems [36], we saw value decoupling the agent as a separate entity in TorchFL.
Currently, a TorchFL agent is identified by a unique identifier, and a shard of the federated dataset when initialized. However, it is designed to be extendable to store more metadata as required. For example, if an incentive mechanism computes the reputation of an agent and uses it to incentivize them for every training round, this can be supported by the agent object by adding the relevant fields with relevant datatypes. Lastly, the agent object is wrapped under a PyTorch Lightning wrapper which allows us to log and collect all the agent-related metadata, individual agent training performance, and statistics, as and when required.
2. Sampler: Once the agents are initialized and assigned their shards of the federated data, the sampler module is responsible to take the collection of agent as an input, and outputting the ones that are selected for training. TorchFL currently supports random sampling as the baseline algorithm, however, extensive work is being done in exploring the sampling approaches to optimize the learning process [35]. In order to add a custom sampling mechanism, the developers need to adhere to the base sampler interface and populate the relevant methods with their custom code. Once a new sampler is defined, it can easily be used in any FL experiments via config files to the entry points. , but an interface has been provided for the developers to implement their own custom aggregators.

Entrypoint:
Once we have the models, dataset, agent, the sampler, and an aggregator, the last step is to wrap them all to produce a complete FL experiment which can execute on various hardware accelerators, and log the metadata and performance as required. TorchFL provides an entrypoint module which handles all of these parts and the users are only required to supply their hyperparameters via a config file when initializing the entrpoint object. Some good examples of the FL hyperparameters are (i) number of agents, (ii) number of global training epochs, (iii) local training epochs, (iv) performance benchmark, (iv) IID or non-IID dataset split, (v) type of sampling mechanism, (vi) type of aggregation protocol, (vii) type of hardware accelerator (TPU, GPU, CPU, or more), (viii) type of logger (CSV, TensorBoard, etc.) and more. As the entrypoint module in TorchFL automatically abstracts all of these hyperparameters, the developers can only focus on running the required set of experiments. See Figure 5 for more details.

Backward Compatibility with PyTorch Lightning
While discussing the design and architecture of various modules in the previous subsections, we have repeatedly emphasized making modules backward compatible with PyTorch Lightning [9]. For context, Lightning is a DL framework and wrapper which is built on PyTorch to provide infrastructure support and features without sacrificing the training performance at scale. By making the TorchFL modules backward compatible with Lightning, we are able to leverage the following features.

Experiments
In this section, we present experimental results using TorchFL to demonstrate its features. All experiments are performed on the following devices: (i) a virtual machine running on a node with 60GB of RAM, 30 CPU cores, virtualized with the hypervisor running on AMD Epyc cores, and no hardware accelerator, (ii) a virtual machine running on a node with 32GB of RAM, 8 CPU cores, virtualized with the hypervisor running on AMD Epyc cores, and NVIDIA Tesla T4 GPU. For the rest of the section, we will refer to these devices as CPU and GPU, respectively.

Usage & Demo
Before we dig deeper into the advanced features, we start by demonstrating the basic techniques to use the toolkit and quickly bootstrap the datasets, models, and FL experiments.

Datasets
We discussed the design of the datamodules in Section 3.1 and listed all of our dataset offerings in Table 1. In this section, we demonstrate the usage of TorchFL to promptly set up the CIFAR-10 dataset [17] and present some visualizations to explain the results. For background, CIFAR-10 is one of the most extensively used datasets in academia given the dimensions and number of total images [27]. It consists of 60000 32×32 color images in 10 classes containing 6000 images each, which are evenly split into 5000 training images and 1000 testing images per class. This means that the dataset consists of 50000 training images and 10000 testing images in total. Each label in the CIFAR-10 dataset can be identified using a unique integer in the range [0,9]. For the first experiment, we split the CIFAR-10 training data between 5 agents, split them using IID and various non-IID configurations provided by TorchFL's datamodule, and visually represent the results in Figure 6. All of the datasets-related results shown on the CIFAR-10 dataset with 5 agents would apply to systems consisting of more agents, or training larger or more complex datasets with more unique labels. Good examples of such datasets are CIFAR-100 [17] with 100 labels, or even ImageNet [6] with 1000 labels.
We chose to first demonstrate the results on a smaller experiment configuration to provide more granular visualizations.

Models
We discussed the design of TorchFL's models in Section 3.1 and listed all of our DL model offerings in Table 2. We also discussed the concepts of feature extraction and finetuning in transfer learning to help the models quickly achieve global convergence. In this subsection, we utilize TorchFl model module to quickly bootstrap various state-of-the-art DL models and provide results to demonstrate the benefits of transfer learning paradigms.
For this set of experiments, we train ResNet152, which is one of the variants of the deep residual networks for image recognition [12]. In summary, ResNets are a variant of convolutional neural networks (CNN) that democratized the concepts of residual learning and skip connections to enable the training of much deeper models. The model will be trained on the CIFAR-10 dataset, and we experiment on the following settings -(i) training the model from scratch using randomly initialized weights, (ii) finetuning the model which is initialized using the weights pre-trained on the ImageNet dataset, and (iii) feature-extraction, i.e. retraining the classification layers of the model while reusing the pre-trained weights for other layers. Again, no additional code was written to set up each of these experiments. In fact, we were able to pass the hyperparameters configuration via TorchFL's model object and the setup was abstracted under the hood. Before we look at the experimental results, it is worth understanding the primary reason behind the differences in training time and resource utilization between each of these settings. ResNet152 being a massive model, Table 3 shows the number of the total, trainable, non-trainable parameters, and training time in seconds (per epoch) when we use either of these training paradigms. We can clearly see how the training time for the feature-extracted model drastically reduces as the number of trainable parameters decrease. As explained in Section 3.1, the finetuning technique still requires retraining of all the parameters, which means, the training time per epoch is the same as training from scratch. However, as the weights from pre-trained models are used (instead of random initialization), we expect to achieve convergence with a lesser number of epochs, which leads to a decrease in overall training time. The plot in Figure 7 shows how the validation accuracy and CrossEntropy loss for ResNet152 vary when trained for 10 epochs using different experimental settings. The plot clearly shows how finetuned and feature-extracted models start with a lower loss because of the pre-trained parameters. Notably, we only trained ResNet152 for 10 epochs as training an entire network on the ImageNet-1K dataset takes roughly more than 3 weeks [12]. For this work, our goal is to show how TorchFL is capable of bootstrapping a model as complicated as ResNet152 and also demonstrate its correctness through the reduction in CrossEntropy [37] loss as shown in Figure 7.

Federated Learning
We discussed the design of federated module in TorchFL in Section 3.2 and saw how decoupled implementation of agent, sampler, and aggregator are used to provide an entrypoint object which is eventually used by the developer. In the previous subsections, we saw how TorchFL provides various abstractions around datasets and models. In this section, we will demonstrate an end-to-end FL experiment which utilizes the datamodule, model, sampler, and aggregator to train a global model.
For this set of experiments, we will use an MNIST [7], a database of handwritten digits. It has a training set of 60000 images and a test set of 10000 images. The images are single-channel, B&W images, with the digits being size-normalized and centered. We choose MNIST over any other dataset for this experiment, as it is relatively easier to achieve convergence [4] using a simple CNN like LeNet-5 [19]. We used TorchFL's datamodules, models, and FL modules to generate an entrypoint object, which was used to abstract and generate all the experiments. We primarily aimed to demonstrate FL from scratch and federated transfer learning using different FL configurations. As seen in Figure 8, the trend in the accuracy-loss curves clearly demonstrates learning. Further, we also clearly notice the impact of non-IID data on model convergence. The second plot in Figure 8 shows the accuracy-loss trend for the federated transfer learning setup mentioned before. Again, there are multiple combinations of FL parameters that can be used for these experiments and would give us different results. However, our goal is to maximize the number of features used in the demo and hence we limit our experiments such that we can demonstrate FL training from scratch and also via transfer learning respectively.

Leveraging PyTorch Lightning Features
In Section 3.3 we discussed various benefits provided by the PyTorch Lightning compatibility through various modules in the library. In this section, we present various ways in which we can use the Lightning profilers and loggers to generate meaningful metrics without any implementation or performance overhead.

Granular Metrics for Individual Agents
Collecting granular training and hardware metrics for individual agents, over multiple local and glocal epochs can be tedious and compute inefficient especially while running a massive FL experiment. TorchFL's federated modules compatibility with Lightning loggers automatically logs the training metrics, sampling counts, parameters, device stats, etc. to the configured loggers. Figure 9 shows how TorchFL was able to automatically collect the local training metrics of a randomly selected agent, every time they were selected for training.

Pinpointing the Bottlenecks Using Profilers
The modules in TorchFL being backward compatible with the Lightning profilers, Table 4 presents a sample profiling result that was generated while training LeNet-5 on the MNIST dataset. Note that the results were generated using a SimpleProfiler object which only monitors the time for core components. Granular details about the system calls and the bottlenecks can be found using an advanced cProfiler.

Monitoring the Core Accelerator Resources
State-of-the-art DL models by themselves require significant memory usage, threads, and pools while training on a hardware accelerator (i.e. GPU). Especially, with FL experiments allocating multiple copies of models for individual agents, this usage is only expected to shoot upwards. As a result, TorchFL's ability to monitor the hardware used for individual agent models provides significant value to the developers. Figure 10 is an example of how TorchFL can monitor the memory usage on the hardware accelerator through the training process. Some other useful metrics recorded by TorchFL include thread pool allocation and thread usage within the individual pools.

Related Work
In this section, we discuss, compare, and analyze various open-source toolkits in the FL space. The goal of this section is to provide a broad, high-level overview of tools available to the best of our knowledge, and analyze how they compare to TorchFL, or can even be used in tandem with TorchFL.
Lately, a few open-source offerings have emerged in the FL space with one of the most promising ones being FedML [11]. With a multi-layer architecture, and beta offerings for iOS, Android, and IoT devices, FedML aims to provide a complete ecosystem to implement an end-to-end, real-world FL system. However, due to a lack of a standardized underlying framework (eg. PyTorch), setting up FL experiments on FedML still require overhead costs of implementing the models, datasets, and FL utilities that are compatible with the core API offering. On the other hand, with the features like models and dataset offerings, granular metric collection, and customizable interfaces, TorchFL aims to solely optimize for the ease of bootstrapping the FL experiments and integrate state-of-the-art research (samplers, aggregators, etc.) with minimum overhead.
LEAF [3] is another promising framework that aims to open-source and benchmark suitable datasets for FL settings. A few of their notable dataset offerings include Reddit data for language modeling, Shakespeare manuscripts for next-character prediction, and more. We believe these datasets can be readily ported into a TorchFL datamodule interface and can readily be trained and tested with our DL model offerings.
Moreover, Tensorflow-Federated (TFF) [1], PySyft [30], and FATE [22] are among the few other notable toolkits that we believe exist in the FL space. While TFF is backward compatible with Tensorflow [1], another popular DL framework, it only supports aggregation based on FedAvg and has no model offerings, monitoring, or logging support. On the other hand, while we believe PySyft to be an actively-maintained and robust offering, their major contributions are ready-to-use algorithms for privacy-preserving techniques like differential privacy [10] and encrypted computation [31]. PySyft's FL offering is the infrastructure to being able to set up privacy-aware agents and service providers as separate entities and able to develop custom PyTorch models, and aggregation protocols on top of it while ensuring private exchange of the computational structures, i.e. tensors. Lastly, FATE is a group of multiple FL frameworks that aim to let developers deploy their FL workflows on industrial-grade infrastructure and allow multiple organizations to potentially engage in the FL training process. In addition to their primary FL offering, they offer multiple tools for serving FL models on various environments and utilities for exchanging the data over the network.

Future Work
This section contains a high-level overview of the plans of maintaining, improving, and refactoring TorchFL in the future. As there are multiple directions in which this work can be extended, we restrict the scope of this discussion and attempt to summarize it in the following manner.

Exploring Performance Enhancements
As mentioned earlier in Section 3.3, Lightning distributed training strategies can be used in TorchFL to reduce the training time if developers have access to multiple hardware accelerators. In addition, the recent release of PyTorch C++ API and the related tools (TorchScript, ATen, etc.) now allows developers to leverage Python and PyTorch to write custom kernels using CUDA [5], that is backward compatible with the original PyTorch implementations. Once we ensure the stability of the current features and offerings, we are inclined to explore the possibilities of high-performance frontiers on TorchFL.

Added Support for More Models and Datasets
Currently, we support major datasets that are oriented toward solving image recognition, classification, or computer vision tasks in general. The models we support are also state-of-the-art models which are used to solve image recognition and segmentation tasks. However, the design and interfaces that we discussed in Section 3.1, are backward compatible with a more diverse range of DL tasks including but not limited to natural language processing (NLP) or reinforcement learning (RL). As we plan to actively maintain TorchFL, we are considering adding a diverse range of datasets and model implementations in the future. Again, it would only require the developers to make the new offerings backward compatible with the TorchFL modules and it would automatically work with all the FL components.

Broader Range of FL Components
One of the primary motivations to design the federated modules as explained in Section 3.2 was to provide an easy-to-use interface to quickly prototype and validate the latest FL components. As the FL community is constantly innovating approaches toward sampling, agent incentivization, defense mechanisms, gradient/parameter encryption, and more, we plan to add those modules in addition to the samplers and aggregators that we currently support. In addition, we also aim to add more samplers and aggregator offerings which will serve as an example for the users who are willing to customize TorchFL for their experiments.

A Code Snippets
This section will provide instructions and example code snippets for the users to get quickly started with TorchFL. First, it will provide an example to use the datamodules and models, and later, it will demonstrate how to build on top of those to bootstrap an FL experiment.

A.1 Using Datasets & Models
The following steps should be followed to bootstrap an experiment with EMNIST (MNIST) dataset and DenseNet121 model. 1. Import the relevant modules as done in Figure 11.   4. The corresponding files for the experiment (checkpoints, metadata, etc.) will be stored at the default root dir argument given to the PyTorch Lightning Trainer object in Step 2. For this experiment, we use the Tensorboard logger. To view the logs (and related plots and metrics), go to the Figure 13: Prepare the dataset and model using TorchFL's datamodules and models wrappers.
default root dir path and find the Tensorboard log files. Upload the files to the Tensorboard Development portal following the instructions here. Note that, TorchFL is compatible with all the loggers supported by PyTorch Lightning. More information about the Lightning loggers can be found here.
5. More example scripts with various models and datasets can be found here.

A.2 Federated Learning
Using the previously set up dataset and models, an FL experiment can be set up in the following manner.
1. Use the dataset generated by TorchFL's datamodules to create federated data shards with IID or non-IID distribution as shown in Figure 14. Figure 14: Create the data shards for the agents using the previously initialized dataset and model.
2. Initialize a global model, and agents, and distribute the global model parameters to every agent as shown in Figure 15.
3. Initialize a TorchFL FLParam object with the desired FL hyperparameters and pass it on to the Entrypoint object which will abstract the training as shown in Figure 16.
4. More federated learning example scripts can be found here.