<jupyter_start><jupyter_text>Preparation: start the websocket server workersEach worker is represented by two parts, a local handle (websocket client worker) and the remote instance that holds the data and performs the computations. The remote part is called a websocket server worker.So first, we need to create the remote workers. For this, you need to run in a terminal (not possible from the notebook):```bashpython start_websocket_servers.py```  Setting up the websocket client workers We first need to perform the imports and setup some arguments and variables.<jupyter_code>%load_ext autoreload
%autoreload 2
import sys
import syft as sy
from syft.workers import WebsocketClientWorker
import torch
from torchvision import datasets, transforms

from syft.frameworks.torch.federated import utils
import run_websocket_client as rwc
args = rwc.define_and_get_arguments(args=[])
use_cuda = args.cuda and torch.cuda.is_available()
torch.manual_seed(args.seed)
device = torch.device("cuda" if use_cuda else "cpu")
print(args)<jupyter_output>Namespace(batch_size=64, cuda=False, epochs=2, federate_after_n_batches=50, lr=0.01, save_model=False, seed=1, test_batch_size=1000, use_virtual=False, verbose=False)<jupyter_text>Now let's instantiate the websocket client workers, our local access point to the remote workers.Note that **this step will fail, if the websocket server workers are not running**.<jupyter_code>hook = sy.TorchHook(torch)

kwargs_websocket = {"host": "http://ec2-13-233-99-209.ap-south-1.compute.amazonaws.com", "hook": hook, "verbose": args.verbose}
#kwargs_websocket = {"host":"172.20.10.2", "hook": hook, "verbose": args.verbose}
alice = WebsocketClientWorker(id="alice", port=8777, **kwargs_websocket)
bob = WebsocketClientWorker(id="bob", port=8778, **kwargs_websocket)
charlie = WebsocketClientWorker(id="charlie", port=8779, **kwargs_websocket)

workers = [alice, bob, charlie]
print(workers)

federated_train_loader = sy.FederatedDataLoader(
    datasets.MNIST(
        "../data",
        train=True,
        download=True,
        transform=transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
        ),
    ).federate(tuple(workers)),
    batch_size=args.batch_size,
    shuffle=True,
    iter_per_worker=True
)

test_loader = torch.utils.data.DataLoader(
    datasets.MNIST(
        "../data",
        train=False,
        transform=transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
        ),
    ),
    batch_size=args.test_batch_size,
    shuffle=True
)

model = rwc.Net().to(device)
print(model)
def train(model, device, federated_train_loader, lr, federate_after_n_batches):
    model.train()

    nr_batches = federate_after_n_batches

    models = {}
    loss_values = {}

    iter(federated_train_loader)  # initialize iterators
    batches = rwc.get_next_batches(federated_train_loader, nr_batches)
    counter = 0

    while True:
        print("Starting training round, batches [{}, {}]".format(counter, counter + nr_batches))
        data_for_all_workers = True
        for worker in batches:
            curr_batches = batches[worker]
            if curr_batches:
                models[worker], loss_values[worker] = rwc.train_on_batches(
                    worker, curr_batches, model, device, lr
                )
            else:
                data_for_all_workers = False
        counter += nr_batches
        if not data_for_all_workers:
            logger.debug("At least one worker ran out of data, stopping.")
            break

        model = utils.federated_avg(models)
        batches = rwc.get_next_batches(federated_train_loader, nr_batches)
    return model
import logging
FORMAT = "%(asctime)s %(levelname)s %(filename)s(l:%(lineno)d) - %(message)s"
LOG_LEVEL = logging.DEBUG
logging.basicConfig(format=FORMAT, level=LOG_LEVEL)
logger = logging.getLogger("main")
for epoch in range(1, args.epochs + 1):
    print("Starting epoch {}/{}".format(epoch, args.epochs))
    model = train(model, device, federated_train_loader, args.lr, args.federate_after_n_batches)
    rwc.test(model, device, test_loader)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10872675250357654, 'fraction_numerical': 0.016452074391988557, 'mean_word_length': 4.0972053462940465, 'pattern_counts': {'":': 6, '<': 8, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 1, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '47857', 'license': 'MIT', 'n_tokens_mistral': 1392, 'n_tokens_neox': 1320, 'n_words': 321, 'path': 'PySyft/Federated learning with websockets and federated averaging.ipynb', 'repo_name': 'rohanbanerjee/morphic.ai', 'size': '4194'}