Workflow Interface: Federated PyTorch MNIST

This tutorial introduces the API to get up and running with your first horizontal federated learning workflow. This work has the following goals:

  • Simplify the federated workflow representation

  • Help users better understand the steps in federated learning (weight extraction, compression, etc.)

  • Designed to maintain data privacy

  • Aims for syntatic consistency with the Netflix MetaFlow project. Infrastructure reuse where possible.

See full notebook.

What is it? The workflow interface is a new way of composing federated learning experiments with OpenFL. It was borne through conversations with researchers and existing users who had novel use cases that didn’t quite fit the standard horizontal federated learning paradigm.

Getting Started First we start by installing the necessary dependencies for the workflow interface:

$ pip install git+https://github.com/intel/openfl.git
$ pip install -r requirements_workflow_interface.txt
$ pip install torch
$ pip install torchvision

We begin with the quintessential example of a small pytorch CNN model trained on the MNIST dataset. Let’s start define our dataloaders, model, optimizer, and some helper functions like we would for any other deep learning experiment

import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch
import torchvision
import numpy as np

n_epochs = 3
batch_size_train = 64
batch_size_test = 1000
learning_rate = 0.01
momentum = 0.5
log_interval = 10

random_seed = 1
torch.backends.cudnn.enabled = False
torch.manual_seed(random_seed)

mnist_train = torchvision.datasets.MNIST(
    "./files/",
    train=True,
    download=True,
    transform=torchvision.transforms.Compose(
        [
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize((0.1307,), (0.3081,)),
        ]
    ),
)

mnist_test = torchvision.datasets.MNIST(
    "./files/",
    train=False,
    download=True,
    transform=torchvision.transforms.Compose(
        [
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize((0.1307,), (0.3081,)),
        ]
    ),
)

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x)

def inference(network,test_loader):
    network.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
    for data, target in test_loader:
        output = network(data)
        test_loss += F.nll_loss(output, target, size_average=False).item()
        pred = output.data.max(1, keepdim=True)[1]
        correct += pred.eq(target.data.view_as(pred)).sum()
    test_loss /= len(test_loader.dataset)
    print('\nTest set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
    test_loss, correct, len(test_loader.dataset),
    100. * correct / len(test_loader.dataset)))
    accuracy = float(correct / len(test_loader.dataset))
    return accuracy

Next we import the FLSpec, LocalRuntime, and placement decorators.

  • FLSpec – Defines the flow specification. User defined flows are subclasses of this.

  • Runtime – Defines where the flow runs, infrastructure for task transitions (how information gets sent). The LocalRuntime runs the flow on a single node.

  • aggregator/collaborator - placement decorators that define where the task will be assigned

from copy import deepcopy

from openfl.experimental.interface import FLSpec, Aggregator, Collaborator
from openfl.experimental.runtime import LocalRuntime
from openfl.experimental.placement import aggregator, collaborator


def FedAvg(models, weights=None):
    new_model = models[0]
    state_dicts = [model.state_dict() for model in models]
    state_dict = new_model.state_dict()
    for key in models[1].state_dict():
        state_dict[key] = torch.from_numpy(np.average([state[key].numpy() for state in state_dicts],
                                                    axis=0,
                                                    weights=weights))
    new_model.load_state_dict(state_dict)
    return new_model

Now we come to the flow definition. The OpenFL Workflow Interface adopts the conventions set by Metaflow, that every workflow begins with start and concludes with the end task. The aggregator begins with an optionally passed in model and optimizer. The aggregator begins the flow with the start task, where the list of collaborators is extracted from the runtime (self.collaborators = self.runtime.collaborators) and is then used as the list of participants to run the task listed in self.next, aggregated_model_validation. The model, optimizer, and anything that is not explicitly excluded from the next function will be passed from the start function on the aggregator to the aggregated_model_validation task on the collaborator. Where the tasks run is determined by the placement decorator that precedes each task definition (@aggregator or @collaborator). Once each of the collaborators (defined in the runtime) complete the aggregated_model_validation task, they pass their current state onto the train task, from train to local_model_validation, and then finally to join at the aggregator. It is in join that an average is taken of the model weights, and the next round can begin.n”,

../../_images/workflow_interface.png 
class FederatedFlow(FLSpec):

    def __init__(self, model=None, optimizer=None, rounds=3, **kwargs):
        super().__init__(**kwargs)
        if model is not None:
            self.model = model
            self.optimizer = optimizer
        else:
            self.model = Net()
            self.optimizer = optim.SGD(self.model.parameters(), lr=learning_rate,
                                    momentum=momentum)
        self.rounds = rounds

    @aggregator
    def start(self):
        print(f'Performing initialization for model')
        self.collaborators = self.runtime.collaborators
        self.private = 10
        self.current_round = 0
        self.next(self.aggregated_model_validation, foreach='collaborators', exclude=['private'])

    @collaborator
    def aggregated_model_validation(self):
        print(f'Performing aggregated model validation for collaborator {self.input}')
        self.agg_validation_score = inference(self.model, self.test_loader)
        print(f'{self.input} value of {self.agg_validation_score}')
        self.next(self.train)

    @collaborator
    def train(self):
        self.model.train()
        self.optimizer = optim.SGD(self.model.parameters(), lr=learning_rate,
                                momentum=momentum)
        train_losses = []
        for batch_idx, (data, target) in enumerate(self.train_loader):
            self.optimizer.zero_grad()
            output = self.model(data)
            loss = F.nll_loss(output, target)
            loss.backward()
            self.optimizer.step()
            if batch_idx % log_interval == 0:
                print('Train Epoch: 1 [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                    batch_idx * len(data), len(self.train_loader.dataset),
                    100. * batch_idx / len(self.train_loader), loss.item()))
                self.loss = loss.item()
                torch.save(self.model.state_dict(), 'model.pth')
                torch.save(self.optimizer.state_dict(), 'optimizer.pth')
        self.training_completed = True
        self.next(self.local_model_validation)

    @collaborator
    def local_model_validation(self):
        self.local_validation_score = inference(self.model, self.test_loader)
        print(
            f'Doing local model validation for collaborator {self.input}: {self.local_validation_score}')
        self.next(self.join, exclude=['training_completed'])

    @aggregator
    def join(self, inputs):
        self.average_loss = sum(input.loss for input in inputs) / len(inputs)
        self.aggregated_model_accuracy = sum(
            input.agg_validation_score for input in inputs) / len(inputs)
        self.local_model_accuracy = sum(
            input.local_validation_score for input in inputs) / len(inputs)
        print(f'Average aggregated model validation values = {self.aggregated_model_accuracy}')
        print(f'Average training loss = {self.average_loss}')
        print(f'Average local model validation values = {self.local_model_accuracy}')
        self.model = FedAvg([input.model for input in inputs])
        self.optimizer = [input.optimizer for input in inputs][0]
        self.current_round += 1
        if self.current_round < self.rounds:
            self.next(self.aggregated_model_validation,
                    foreach='collaborators', exclude=['private'])
        else:
            self.next(self.end)

    @aggregator
    def end(self):
        print(f'This is the end of the flow')

You’ll notice in the FederatedFlow definition above that there were certain attributes that the flow was not initialized with, namely the train_loader and test_loader for each of the collaborators. These are private_attributes that are exposed only throught he runtime. Each participant has it’s own set of private attributes: a dictionary where the key is the attribute name, and the value is the object that will be made accessible through that participant’s task.

Below, we segment shards of the MNIST dataset for four collaborators: Portland, Seattle, Chandler, and Portland. Each has their own slice of the dataset that’s accessible via the train_loader or test_loader attribute. Note that the private attributes are flexible, and you can choose to pass in a completely different type of object to any of the collaborators or aggregator (with an arbitrary name). These private attributes will always be filtered out of the current state when transfering from collaborator to aggregator, or vice versa.

# Aggregator
aggregator_ = Aggregator()

collaborator_names = ["Portland", "Seattle", "Chandler", "Bangalore"]

def callable_to_initialize_collaborator_private_attributes(index, n_collaborators, batch_size, train_dataset, test_dataset):
    train = deepcopy(train_dataset)
    test = deepcopy(test_dataset)
    train.data = train_dataset.data[index::n_collaborators]
    train.targets = train_dataset.targets[index::n_collaborators]
    test.data = test_dataset.data[index::n_collaborators]
    test.targets = test_dataset.targets[index::n_collaborators]

    return {
        "train_loader": torch.utils.data.DataLoader(train, batch_size=batch_size, shuffle=True),
        "test_loader": torch.utils.data.DataLoader(test, batch_size=batch_size, shuffle=True),
    }

# Setup collaborators private attributes via callable function
collaborators = []
for idx, collaborator_name in enumerate(collaborator_names):
    collaborators.append(
        Collaborator(
            name=collaborator_name,
            private_attributes_callable=callable_to_initialize_collaborator_private_attributes,
            index=idx,
            n_collaborators=len(collaborator_names),
            train_dataset=mnist_train,
            test_dataset=mnist_test,
            batch_size=64
        )
    )

local_runtime = LocalRuntime(aggregator=aggregator_, collaborators=collaborators,
                            backend="ray")
print(f'Local runtime collaborators = {local_runtime.collaborators}')

Now that we have our flow and runtime defined, let’s run the experiment!

model = None
best_model = None
optimizer = None
flflow = FederatedFlow(model, optimizer, checkpoint=True)
flflow.runtime = local_runtime
flflow.run()

Now that the flow has completed, let’s get the final model and accuracy:

print(f'Sample of the final model weights: {flflow.model.state_dict()["conv1.weight"][0]}')

print(f'\nFinal aggregated model accuracy for {flflow.rounds} rounds of training: {flflow.aggregated_model_accuracy}')

We can get the final model, and all other aggregator attributes after the flow completes. But what if there’s an intermediate model task and its specific output that we want to look at in detail? This is where checkpointing and reuse of Metaflow tooling come in handy.

Let’s make a tweak to the flow object, and run the experiment one more time (we can even use our previous model / optimizer as a base for the experiment)

flflow2 = FederatedFlow(model=flflow.model, optimizer=flflow.optimizer, checkpoint=True)
flflow2.runtime = local_runtime
flflow2.run()

Now that the flow is complete, let’s dig into some of the information captured along the way

run_id = flflow2._run_id

import metaflow
from metaflow import Metaflow, Flow, Task, Step

m = Metaflow()
list(m)

For existing users of Metaflow, you’ll notice this is the same way you would examine a flow after completion. Let’s look at the latest run that generated some results:

f = Flow('FederatedFlow').latest_run
f

And its list of steps

list(f)

This matches the list of steps executed in the flow, so far so good…

s = Step(f'FederatedFlow/{run_id}/train')
s

list(s)

Now we see 12 steps: 4 collaborators each performed 3 rounds of model training

t = Task(f'FederatedFlow/{run_id}/train/9')
t

t.data

t.data.input

Now let’s look at its log output (stdout) and any error logs (stderr)

print(t.stdout)
print(t.stderr)

Congratulations! Now that you’ve completed your first workflow interface quickstart notebook,

see some of the more advanced things you can do in our other tutorials, including:

  • Using the LocalRuntime Ray Backend for dedicated GPU access

  • Vertical Federated Learning

  • Model Watermarking

  • Differential Privacy

  • And More!