text: # Copyright 2020 Adap GmbH. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """PyTorch CIFAR-10 image classification.
  
  The code is generally adapted from 'PyTorch: A 60 Minute Blitz'. Further
  explanations are given in the official PyTorch tutorial:
  
  https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html
  """
  
  
  # mypy: ignore-errors
  # pylint: disable=W0223
  
  
  from collections import OrderedDict
  from typing import Tuple
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import torchvision
  import torchvision.transforms as transforms
  from torch import Tensor
  
  import flwr as fl
  
  DATA_ROOT = "~/.flower/data/cifar-10"
  
  # pylint: disable-msg=unsubscriptable-object
  class Net(nn.Module):
      """Simple CNN adapted from 'PyTorch: A 60 Minute Blitz'."""
  
      def __init__(self) -> None:
          super(Net, self).__init__()
          self.conv1 = nn.Conv2d(3, 6, 5)
          self.pool = nn.MaxPool2d(2, 2)
          self.conv2 = nn.Conv2d(6, 16, 5)
          self.fc1 = nn.Linear(16 * 5 * 5, 120)
          self.fc2 = nn.Linear(120, 84)
          self.fc3 = nn.Linear(84, 10)
  
      # pylint: disable-msg=arguments-differ,invalid-name
      def forward(self, x: Tensor) -> Tensor:
          """Compute forward pass."""
          x = self.pool(F.relu(self.conv1(x)))
          x = self.pool(F.relu(self.conv2(x)))
          x = x.view(-1, 16 * 5 * 5)
          x = F.relu(self.fc1(x))
          x = F.relu(self.fc2(x))
          x = self.fc3(x)
          return x
  
      def get_weights(self) -> fl.common.Weights:
          """Get model weights as a list of NumPy ndarrays."""
          return [val.cpu().numpy() for _, val in self.state_dict().items()]
  
      def set_weights(self, weights: fl.common.Weights) -> None:
          """Set model weights from a list of NumPy ndarrays."""
          state_dict = OrderedDict(
              {k: torch.Tensor(v) for k, v in zip(self.state_dict().keys(), weights)}
          )
          self.load_state_dict(state_dict, strict=True)
  
  
  # pylint: disable-msg=unused-argument
  def load_data() -> Tuple[torch.utils.data.DataLoader, torch.utils.data.DataLoader]:
      """Load CIFAR-10 (training and test set)."""
      transform = transforms.Compose(
          [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
      )
  
      # Training set
      trainset = torchvision.datasets.CIFAR10(
          root=DATA_ROOT, train=True, download=True, transform=transform
      )
      trainloader = torch.utils.data.DataLoader(trainset, batch_size=32, shuffle=True)
  
      # Test set
      testset = torchvision.datasets.CIFAR10(
          root=DATA_ROOT, train=False, download=True, transform=transform
      )
      testloader = torch.utils.data.DataLoader(testset, batch_size=32, shuffle=False)
  
      return trainloader, testloader
  
  
  def train(
      net: Net,
      trainloader: torch.utils.data.DataLoader,
      epochs: int,
      device: torch.device,  # pylint: disable=no-member
  ) -> None:
      """Train the network."""
      # Define loss and optimizer
      criterion = nn.CrossEntropyLoss()
      optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
  
      print(f"Training {epochs} epoch(s) w/ {len(trainloader)} batches each")
  
      # Train the network
      for epoch in range(epochs):  # loop over the dataset multiple times
          running_loss = 0.0
          for i, data in enumerate(trainloader, 0):
              images, labels = data[0].to(device), data[1].to(device)
  
              # zero the parameter gradients
              optimizer.zero_grad()
  
              # forward + backward + optimize
              outputs = net(images)
              loss = criterion(outputs, labels)
              loss.backward()
              optimizer.step()
  
              # print statistics
              running_loss += loss.item()
              if i % 2000 == 1999:  # print every 2000 mini-batches
                  print("[%d, %5d] loss: %.3f" % (epoch + 1, i + 1, running_loss / 2000))
                  running_loss = 0.0
  
  
  def test(
      net: Net,
      testloader: torch.utils.data.DataLoader,
      device: torch.device,  # pylint: disable=no-member
  ) -> Tuple[float, float]:
      """Validate the network on the entire test set."""
      criterion = nn.CrossEntropyLoss()
      correct = 0
      total = 0
      loss = 0.0
      with torch.no_grad():
          for data in testloader:
              images, labels = data[0].to(device), data[1].to(device)
              outputs = net(images)
              loss += criterion(outputs, labels).item()
              _, predicted = torch.max(outputs.data, 1)  # pylint: disable-msg=no-member
              total += labels.size(0)
              correct += (predicted == labels).sum().item()
      accuracy = correct / total
      return loss, accuracy
  