code: import copy
  from dataclasses import dataclass
  from typing import List, Optional
  
  import torch
  from torch.nn import CrossEntropyLoss, Module
  from torch.utils.data import DataLoader
  
  
  def federated_averaging(models: List[Module]) -> Module:
      global_model = copy.deepcopy(models[0])
      global_weights = global_model.state_dict()
  
      local_weights = [m.state_dict() for m in models]
  
      for k in global_weights.keys():
          for i in range(1, len(local_weights)):
              global_weights[k] += local_weights[i][k]
          global_weights[k] = torch.div(global_weights[k], len(local_weights))
  
      global_model.load_state_dict(global_weights)
      return global_model
  
  
  class ModelAccumulator:
      def __init__(self):
          self.model_counter: int = 0
          self.global_model = None
          self.global_weights = None
  
      def update(self, model):
          local_weights = model.state_dict()
  
          if self.global_model is None:
              self.global_model = model
              self.global_weights = local_weights
              self.model_counter += 1
          else:
              for k in self.global_weights.keys():
                  self.global_weights[k] += local_weights[k]
              self.model_counter += 1
  
      def get(self):
          for k in self.global_weights.keys():
              self.global_weights[k] = torch.div(
                  self.global_weights[k], self.model_counter
              )
  
          self.global_model.load_state_dict(self.global_weights)
          return self.global_model
  
      def reset(self):
          self.global_model = None
          self.global_weights = None
          self.model_counter = 0
  
  
  @dataclass
  class EdgeDeviceSettings:
      batch_size: int
      epochs: int
      learning_rate: float
      learning_rate_decay: float
      device: str
  
  
  @dataclass
  class TrainingResult:
      loss: float
      steps: int
      learning_rate: float
  
  
  class EdgeDevice:
      def __init__(
          self, device_id: int, settings: EdgeDeviceSettings, data_loader: DataLoader
      ):
          self.device_id = device_id
          self._data_loader = data_loader
          self.setting = copy.deepcopy(settings)
          self._loss_func = CrossEntropyLoss()
          self._model: Optional[Module] = None
  
      def download(self, model: Module):
          self._model = copy.deepcopy(model)
  
      def upload(self) -> Module:
          if self._model is not None:
              return copy.deepcopy(self._model)
          else:
              raise ValueError("Model not found on this device!")
  
      def train(self) -> TrainingResult:
          if self._data_loader is None:
              raise ValueError("Dataset not found on this device!")
  
          self._model.train()
          self.setting.learning_rate = (
              self.setting.learning_rate * self.setting.learning_rate_decay
          )
          optimizer = torch.optim.SGD(
              params=self._model.parameters(), lr=self.setting.learning_rate
          )
          epoch_loss = []
          local_steps: int = 0
          for _ in range(self.setting.epochs):
              batch_loss = []
              for i_batch, (images, labels) in enumerate(self._data_loader):
                  self._model.zero_grad()
                  images = images.to(self.setting.device)
                  labels = labels.to(self.setting.device)
  
                  logits = self._model(images)
                  loss = self._loss_func(logits, labels)
                  loss.backward()
                  optimizer.step()
                  local_steps += 1
                  batch_loss.append(loss.item())
              epoch_loss.append(sum(batch_loss) / len(batch_loss))
  
          mean_loss = sum(epoch_loss) / len(epoch_loss)
          return TrainingResult(
              loss=mean_loss, steps=local_steps, learning_rate=self.setting.learning_rate
          )
   import argparse
  import glob
  import os
  import time
  
  import pandas as pd
  from tqdm import tqdm
  
  
  def generate_csv_file(data_dir: str, output_file: str):
      """Generates a csv file containing the image paths of the VGGFace2 dataset for use in triplet selection in
      triplet loss training.
  
      Args:
          dataroot (str): absolute path to the training dataset.
          csv_name (str): name of the resulting csv file.
      """
      print("\nLoading image paths ...")
      files = glob.glob(data_dir + "/*/*")
      start_time = time.time()
      list_rows = []
  
      print(data_dir)
      print(os.listdir(data_dir))
      print(files)
  
      print("Number of files: {}".format(len(files)))
      print("\nGenerating csv file ...")
  
      for file_index, file in enumerate(tqdm(files)):
          face_id = os.path.basename(file).split(".")[0]
          face_label = os.path.basename(os.path.dirname(file))
  
          # Better alternative than dataframe.append()
          row = {"id": face_id, "name": face_label}
          list_rows.append(row)
  
      df = pd.DataFrame(list_rows)
      df = df.sort_values(by=["name", "id"]).reset_index(drop=True)
  
      # Encode names as categorical classes
      df["class"] = pd.factorize(df["name"])[0]
      df.to_csv(path_or_buf=output_file, index=False)
  
      elapsed_time = time.time() - start_time
      print("\nDone! Elapsed time: {:.2f} minutes.".format(elapsed_time / 60))
  
  
  def parse_args():
      parser = argparse.ArgumentParser(
          description="Generating csv file for triplet loss!"
      )
      parser.add_argument(
          "--dataroot",
          "-d",
          type=str,
          required=True,
          help="(REQUIRED) Absolute path to the dataset folder to generate a csv "
          "file containing the paths of the images for triplet loss. ",
      )
  
      parser.add_argument(
          "--csv_name",
          type=str,
          help="Required name of the csv file to be generated. (default: "
          "'vggface2.csv') ",
      )
  
      args = parser.parse_args()
      return args
  
  
  def main():
      args = parse_args()
      generate_csv_file(data_dir=args.dataroot, output_file=args.csv_name)
  
  
  if __name__ == "__main__":
      main()
  
apis: torch.nn.CrossEntropyLoss torch.div
apis: pandas.DataFrame pandas.factorize