text: import argparse
  import logging
  import os
  import random
  import sys
  import time
  
  import numpy as np
  import torch
  import wandb
  
  sys.path.insert(0, os.path.abspath(os.path.join(os.getcwd(), "../../../")))
  
  import singleMachine.myUtils as mu
  
  ## data
  from fedml_api.data_preprocessing.spectrum.data_loader import load_partition_data_spectrum
  
  ## model
  from fedml_api.model.spectrum.fnn import FeedforwardNN
  
  ## system
  from fedml_api.standalone.fedavg.fedavg_api import FedAvgAPI
  from fedml_api.standalone.fedavg.my_model_trainer_value_prediction import MyModelTrainer
  
  
  # arg parser
  def add_args(parser):
      """
      parser : argparse.ArgumentParser
      return a parser added with args required by fit
      """
      # Training settings
  
      parser.add_argument('--alpha', type=float, default=1, metavar='N',
                           help='alpha parameter for mixing global weight between gens')
  
      parser.add_argument('--model', type=str, default='fnn_spectrum', metavar='N',
                           help='neural network used in training')
  
      parser.add_argument('--dataset', type=str, default='spectrum_gps', metavar='N',
                          help='dataset used for training')
  
      parser.add_argument('--data_dir', type=str, default='./../../../data/spectrum',
                          help='data directory')
  
      # currently we just divide dataset uniformly to edge devices
      parser.add_argument('--partition_method', type=str, default='hetero', metavar='N',
                          help='how to partition the dataset on local workers')
  
      parser.add_argument('--partition_alpha', type=float, default=0.5, metavar='PA',
                          help='partition alpha (default: 0.5)')
  
      parser.add_argument('--batch_size', type=int, default=10, metavar='N',
                          help='input batch size for training (default: 64)') # -1 means fullbatch
  
      parser.add_argument('--client_optimizer', type=str, default='adam',
                          help='SGD with momentum; adam')
  
      parser.add_argument('--lr', type=float, default=0.001, metavar='LR',
                          help='learning rate (default: 0.001)')
  
      parser.add_argument('--wd', help='weight decay parameter;', type=float, default=0.001)
  
      parser.add_argument('--epochs', type=int, default=5, metavar='EP',
                          help='how many epochs will be trained locally')
  
      parser.add_argument('--client_num_in_total', type=int, default=10, metavar='NN',
                          help='number of workers in a distributed cluster')
  
      parser.add_argument('--client_num_per_round', type=int, default=10, metavar='NN',
                          help='number of workers')
  
      parser.add_argument('--comm_round', type=int, default=10,
                          help='how many round of communications we shoud use')
  
      parser.add_argument('--frequency_of_the_test', type=int, default=1,
                          help='the frequency of the algorithms')
  
      parser.add_argument('--gpu', type=int, default=0,
                          help='gpu')
  
      parser.add_argument('--ci', type=int, default=0,
                          help='CI')
  
      parser.add_argument('--rand', type=int, default=0,
                          help='real random?')
      
      return parser
  
  # read json dataset 
  def load_data(args, dataset_name):
  
      # check if the full-batch training is enabled
      args_batch_size = args.batch_size
      if args.batch_size <= 0:
          full_batch = True
          args.batch_size = 128  # temporary batch size
      else:
          full_batch = False
  
      if dataset_name == "spectrum_gps":
          logging.info("load_data. dataset_name = %s" % dataset_name)
          client_num, train_data_num, test_data_num, train_data_global, test_data_global, \
          train_data_local_num_dict, train_data_local_dict, test_data_local_dict, \
          class_num = load_partition_data_spectrum(args.batch_size)
  
      # check if the centralized training is enabled
      centralized = True if args.client_num_in_total == 1 else False
      if centralized:
          train_data_local_num_dict = {
              0: sum(user_train_data_num for user_train_data_num in train_data_local_num_dict.values())}
          train_data_local_dict = {
              0: [batch for cid in sorted(train_data_local_dict.keys()) for batch in train_data_local_dict[cid]]}
          test_data_local_dict = {
              0: [batch for cid in sorted(test_data_local_dict.keys()) for batch in test_data_local_dict[cid]]}
          args.client_num_in_total = 1
  
      if full_batch:
          train_data_global = combine_batches(train_data_global)
          test_data_global = combine_batches(test_data_global)
          train_data_local_dict = {cid: combine_batches(train_data_local_dict[cid]) for cid in
                                   train_data_local_dict.keys()}
          test_data_local_dict = {cid: combine_batches(test_data_local_dict[cid]) for cid in test_data_local_dict.keys()}
          args.batch_size = args_batch_size
  
      dataset = [train_data_num, test_data_num, train_data_global, test_data_global,
                 train_data_local_num_dict, train_data_local_dict, test_data_local_dict, class_num]
      return dataset
  
  # used for full batch, but i may not use it. just leave it here
  def combine_batches(batches): 
      full_x = torch.from_numpy(np.asarray([])).float()
      full_y = torch.from_numpy(np.asarray([])).long()
      for (batched_x, batched_y) in batches:
          full_x = torch.cat((full_x, batched_x), 0)
          full_y = torch.cat((full_y, batched_y), 0)
      return [(full_x, full_y)]
  
  if __name__ == "__main__":
      logging.basicConfig()
      logger = logging.getLogger()
      logger.setLevel(logging.DEBUG)
      mpl_logger = logging.getLogger('matplotlib')
      mpl_logger.setLevel(logging.WARNING)
  
      parser = add_args(argparse.ArgumentParser(description='FedAvg-standalone'))
      args = parser.parse_args()
      logger.info(args)
      device = torch.device("cuda:" + str(args.gpu) if torch.cuda.is_available() else "cpu")
      logger.info(device)
  
      wandb.init(
          project="fedml_spectrum",
          name="spectrum-r" + str(args.comm_round) + "-e" + str(args.epochs) + "-lr" + str(args.lr) + "-cli" + str(args.client_num_in_total) + "-alpha" + str(args.alpha),
          config=args
      )
  
      # Set the random seed. The np.random seed determines the dataset partition.
      # The torch_manual_seed determines the initial weight.
      if args.rand == 1:
          rseed = round(time.time())
      else:
          rseed = 524287
      random.seed(rseed)
      np.random.seed(rseed)
      torch.manual_seed(rseed)
      if torch.cuda.is_available():
          torch.cuda.manual_seed_all(rseed)
  
      # load data
      dataset = load_data(args, args.dataset)
  
      # create model.
      # Note if the model is DNN (e.g., ResNet), the training will be very slow.
      # In this case, please use our FedML distributed version (./fedml_experiments/distributed_fedavg)
      model = FeedforwardNN()
      logging.info(model)
      fedavgAPI = FedAvgAPI(dataset, device, args, MyModelTrainer(model))
      fedavgAPI.train()
      mu.visFNN(model)
      mu.visFNN_small(model)