code: import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from torchvision import transforms, datasets
  
  import logging
  import argparse
  import sys
  import asyncio
  import numpy as np
  
  import syft as sy
  from syft import workers
  from syft.frameworks.torch.federated import utils
  
  logger = logging.getLogger(__name__)
  
  LOG_INTERVAL = 25
  
  
  # Loss function
  @torch.jit.script
  def loss_fn(pred, target):
      return F.nll_loss(input=pred, target=target)
  
  
  # Model
  class Net(nn.Module):
      def __init__(self):
          super(Net, self).__init__()
          self.conv1 = nn.Conv2d(1, 20, 5, 1)
          self.conv2 = nn.Conv2d(20, 50, 5, 1)
          self.fc1 = nn.Linear(4 * 4 * 50, 500)
          self.fc2 = nn.Linear(500, 10)
  
      def forward(self, x):
          x = F.relu(self.conv1(x))
          x = F.max_pool2d(x, 2, 2)
          x = F.relu(self.conv2(x))
          x = F.max_pool2d(x, 2, 2)
          x = x.view(-1, 4 * 4 * 50)
          x = F.relu(self.fc1(x))
          x = self.fc2(x)
          return F.log_softmax(x, dim=1)
  
  
  def define_and_get_arguments(args=sys.argv[1:]):
      parser = argparse.ArgumentParser(
          description="Run federated learning using websocket client workers."
      )
      parser.add_argument("--batch_size", type=int, default=32, help="batch size of the training")
      parser.add_argument(
          "--test_batch_size", type=int, default=128, help="batch size used for the test data"
      )
      parser.add_argument(
          "--training_rounds", type=int, default=40, help="number of federated learning rounds"
      )
      parser.add_argument(
          "--federate_after_n_batches",
          type=int,
          default=10,
          help="number of training steps performed on each remote worker before averaging",
      )
      parser.add_argument("--lr", type=float, default=0.1, help="learning rate")
      parser.add_argument("--cuda", action="store_true", help="use cuda")
      parser.add_argument("--seed", type=int, default=1, help="seed used for randomization")
      parser.add_argument("--save_model", action="store_true", help="if set, model will be saved")
      parser.add_argument(
          "--verbose",
          "-v",
          action="store_true",
          help="if set, websocket client workers will be started in verbose mode",
      )
  
      args = parser.parse_args(args=args)
      return args
  
  
  async def fit_model_on_worker(
      worker: workers.WebsocketClientWorker,
      traced_model: torch.jit.ScriptModule,
      batch_size: int,
      curr_round: int,
      max_nr_batches: int,
      lr: float,
  ):
      """Send the model to the worker and fit the model on the worker's training data.
  
      Args:
          worker: Remote location, where the model shall be trained.
          traced_model: Model which shall be trained.
          batch_size: Batch size of each training step.
          curr_round: Index of the current training round (for logging purposes).
          max_nr_batches: If > 0, training on worker will stop at min(max_nr_batches, nr_available_batches).
          lr: Learning rate of each training step.
  
      Returns:
          A tuple containing:
              * worker_id: Union[int, str], id of the worker.
              * improved model: torch.jit.ScriptModule, model after training at the worker.
              * loss: Loss on last training batch, torch.tensor.
      """
      train_config = sy.TrainConfig(
          model=traced_model,
          loss_fn=loss_fn,
          batch_size=batch_size,
          shuffle=True,
          max_nr_batches=max_nr_batches,
          epochs=1,
          lr=lr,
      )
      train_config.send(worker)
      logger.info(
          "Training round %s, calling fit on worker: %s, lr = %s",
          curr_round,
          worker.id,
          "{:.3f}".format(train_config.lr),
      )
      loss = await worker.async_fit(dataset_key="mnist", return_ids=[0])
      logger.info("Training round: %s, worker: %s, avg_loss: %s", curr_round, worker.id, loss.mean())
      model = train_config.model_ptr.get().obj
      return worker.id, model, loss
  
  
  def evaluate_models_on_test_data(test_loader, results):
      np.set_printoptions(formatter={"float": "{: .0f}".format})
      for worker_id, worker_model, _ in results:
          evaluate_model(worker_id, worker_model, "cpu", test_loader, print_target_hist=False)
  
  
  def evaluate_model(worker_id, model, device, test_loader, print_target_hist=False):
      model.eval()
      test_loss = 0.0
      correct = 0
      hist_target = np.zeros(10)
      hist_pred = np.zeros(10)
  
      with torch.no_grad():
          for data, target in test_loader:
              data, target = data.to(device), target.to(device)
              hist, _ = np.histogram(target, bins=10, range=(0, 10))
              hist_target += hist
              output = model(data)
              test_loss += loss_fn(output, target).item()  # sum up batch loss
              pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
              hist, _ = np.histogram(pred, bins=10, range=(0, 10))
              hist_pred += hist
              correct += pred.eq(target.view_as(pred)).sum().item()
  
      test_loss /= len(test_loader.dataset)
  
      if print_target_hist:
          logger.info("Target histogram: %s", hist_target)
      logger.info("Prediction hist.: %s", hist_pred)
  
      logger.info(
          "%s: Test set: Average loss: %s, Accuracy: %s/%s (%s)",
          worker_id,
          "{:.4f}".format(test_loss),
          correct,
          len(test_loader.dataset),
          "{:.2f}".format(100.0 * correct / len(test_loader.dataset)),
      )
  
  
  async def main():
      args = define_and_get_arguments()
  
      hook = sy.TorchHook(torch)
  
      kwargs_websocket = {"host": "localhost", "hook": hook, "verbose": args.verbose}
      alice = workers.WebsocketClientWorker(id="alice", port=8777, **kwargs_websocket)
      bob = workers.WebsocketClientWorker(id="bob", port=8778, **kwargs_websocket)
      charlie = workers.WebsocketClientWorker(id="charlie", port=8779, **kwargs_websocket)
  
      worker_instances = [alice, bob, charlie]
  
      use_cuda = args.cuda and torch.cuda.is_available()
  
      torch.manual_seed(args.seed)
  
      device = torch.device("cuda" if use_cuda else "cpu")
  
      kwargs = {"num_workers": 1, "pin_memory": True} if use_cuda else {}
  
      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=False,
          drop_last=False,
          **kwargs,
      )
  
      model = Net().to(device)
  
      (data, target) = test_loader.__iter__().next()
      traced_model = torch.jit.trace(model, data)
      learning_rate = args.lr
  
      for curr_round in range(1, args.training_rounds + 1):
          logger.info("Starting training round %s/%s", curr_round, args.training_rounds)
  
          results = await asyncio.gather(
              *[
                  fit_model_on_worker(
                      worker=worker,
                      traced_model=traced_model,
                      batch_size=args.batch_size,
                      curr_round=curr_round,
                      max_nr_batches=args.federate_after_n_batches,
                      lr=learning_rate,
                  )
                  for worker in worker_instances
              ]
          )
          models = {}
          loss_values = {}
  
          test_models = curr_round % 10 == 1 or curr_round == args.training_rounds
          if test_models:
              evaluate_models_on_test_data(test_loader, results)
  
          for worker_id, worker_model, worker_loss in results:
              if worker_model is not None:
                  models[worker_id] = worker_model
                  loss_values[worker_id] = worker_loss
  
          traced_model = utils.federated_avg(models)
          if test_models:
              evaluate_model(
                  "Federated model", traced_model, "cpu", test_loader, print_target_hist=True
              )
  
          # decay learning rate
          learning_rate = max(0.98 * learning_rate, args.lr * 0.01)
  
      if args.save_model:
          torch.save(model.state_dict(), "mnist_cnn.pt")
  
  
  if __name__ == "__main__":
      # Logging setup
      logger = logging.getLogger("run_websocket_server")
      FORMAT = "%(asctime)s %(levelname)s %(filename)s(l:%(lineno)d, p:%(process)d) - %(message)s"
      logging.basicConfig(format=FORMAT)
      logger.setLevel(level=logging.DEBUG)
  
      # Websockets setup
      websockets_logger = logging.getLogger("websockets")
      websockets_logger.setLevel(logging.INFO)
      websockets_logger.addHandler(logging.StreamHandler())
  
      # Run main
      asyncio.get_event_loop().run_until_complete(main())
   import torch as th
  from torch.utils.data import BatchSampler, RandomSampler, SequentialSampler
  
  from syft.generic import ObjectStorage
  from syft.federated.train_config import TrainConfig
  
  
  class FederatedClient(ObjectStorage):
      """A Client able to execute federated learning in local datasets."""
  
      def __init__(self, datasets=None):
          super().__init__()
          self.datasets = datasets if datasets is not None else dict()
          self.optimizer = None
          self.train_config = None
  
      def add_dataset(self, dataset, key: str):
          self.datasets[key] = dataset
  
      def remove_dataset(self, key: str):
          if key in self.datasets:
              del self.datasets[key]
  
      def set_obj(self, obj: object):
          """Registers objects checking if which objects it should cache.
  
          Args:
              obj: An object to be registered.
          """
          if isinstance(obj, TrainConfig):
              self.train_config = obj
              self.optimizer = None
          else:
              super().set_obj(obj)
  
      def _build_optimizer(
          self, optimizer_name: str, model, optimizer_args: dict
      ) -> th.optim.Optimizer:
          """Build an optimizer if needed.
  
          Args:
              optimizer_name: A string indicating the optimizer name.
              optimizer_args: A dict containing the args used to initialize the optimizer.
          Returns:
              A Torch Optimizer.
          """
          if self.optimizer is not None:
              return self.optimizer
  
          if optimizer_name in dir(th.optim):
              optimizer = getattr(th.optim, optimizer_name)
              self.optimizer = optimizer(model.parameters(), **optimizer_args)
          else:
              raise ValueError("Unknown optimizer: {}".format(optimizer_name))
          return self.optimizer
  
      def fit(self, dataset_key: str, **kwargs):
          """Fits a model on the local dataset as specified in the local TrainConfig object.
  
          Args:
              dataset_key: Identifier of the local dataset that shall be used for training.
              **kwargs: Unused.
  
          Returns:
              loss: Training loss on the last batch of training data.
          """
          if self.train_config is None:
              raise ValueError("TrainConfig not defined.")
  
          if dataset_key not in self.datasets:
              raise ValueError("Dataset {} unknown.".format(dataset_key))
  
          model = self.get_obj(self.train_config._model_id).obj
          loss_fn = self.get_obj(self.train_config._loss_fn_id).obj
  
          self._build_optimizer(
              self.train_config.optimizer, model, optimizer_args=self.train_config.optimizer_args
          )
  
          return self._fit(model=model, dataset_key=dataset_key, loss_fn=loss_fn)
  
      def _create_data_loader(self, dataset_key: str, shuffle: bool = False):
          data_range = range(len(self.datasets[dataset_key]))
          if shuffle:
              sampler = RandomSampler(data_range)
          else:
              sampler = SequentialSampler(data_range)
          data_loader = th.utils.data.DataLoader(
              self.datasets[dataset_key],
              batch_size=self.train_config.batch_size,
              sampler=sampler,
              num_workers=0,
          )
          return data_loader
  
      def _fit(self, model, dataset_key, loss_fn):
          model.train()
          data_loader = self._create_data_loader(
              dataset_key=dataset_key, shuffle=self.train_config.shuffle
          )
  
          loss = None
          iteration_count = 0
  
          for _ in range(self.train_config.epochs):
              for (data, target) in data_loader:
                  # Set gradients to zero
                  self.optimizer.zero_grad()
  
                  # Update model
                  output = model(data)
                  loss = loss_fn(target=target, pred=output)
                  loss.backward()
                  self.optimizer.step()
  
                  # Update and check interation count
                  iteration_count += 1
                  if iteration_count >= self.train_config.max_nr_batches >= 0:
                      break
  
          return loss
  
apis: torch.nn.functional.log_softmax torch.nn.Linear numpy.zeros torch.nn.functional.max_pool2d numpy.histogram torch.nn.functional.nll_loss torch.manual_seed torch.no_grad numpy.set_printoptions torch.cuda.is_available torch.nn.Conv2d torch.device torch.jit.trace
apis: torch.utils.data.DataLoader torch.utils.data.RandomSampler torch.utils.data.SequentialSampler