text: import os
  import json
  import urllib, tempfile, shutil
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import random
  import math
  
  from torch.optim.lr_scheduler import _LRScheduler
  
  
  def get_net_config(config_path='NetConfig.json'):
      assert os.path.exists(config_path)
      with open(config_path, 'r') as f:
          return json.loads(f.read())
  
  
  def get_num_workers(server_name, name_list_path='ServerList.json'):
      assert os.path.exists(name_list_path)
      with open(name_list, 'r') as f:
          name_list = json.loads(f.read())
  
      assert server_name in name_list
      return name_list[server_name]
  
  
  def load_adapted_params(base_model_name='mobilenet_v2'):
      assert base_model_name in ['mobilenet_v2', 'mnasnet_b1', 'fbnet_c']
  
      url = "https://raw.githubusercontent.com/beomwookang/het_pretrained/main/"+base_model_name+"_adapted.pth"
      req = urllib.request.Request(url)
  
      print("Loading %s_adapted params...\n" %base_model_name)
      with urllib.request.urlopen(req) as response:
          with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
              shutil.copyfileobj(response, tmp_file)
      return torch.load(tmp_file.name)
  
  
  # source code: https://github.com/clovaai/CutMix-PyTorch/blob/master/utils.py
  #__all__ = ["Compose", "Lighting", "ColorJitter"]
  
  
  class Compose(object):
      """Composes several transforms together.
      Args:
          transforms (list of ``Transform`` objects): list of transforms to compose.
      Example:
          >>> transforms.Compose([
          >>>     transforms.CenterCrop(10),
          >>>     transforms.ToTensor(),
          >>> ])
      """
  
      def __init__(self, transforms):
          self.transforms = transforms
  
      def __call__(self, img):
          for t in self.transforms:
              img = t(img)
          return img
  
      def __repr__(self):
          format_string = self.__class__.__name__ + '('
          for t in self.transforms:
              format_string += '\n'
              format_string += '    {0}'.format(t)
          format_string += '\n)'
          return format_string
  
  
  class Lighting(object):
      """Lighting noise(AlexNet - style PCA - based noise)"""
  
      def __init__(self, alphastd, eigval, eigvec):
          self.alphastd = alphastd
          self.eigval = torch.Tensor(eigval)
          self.eigvec = torch.Tensor(eigvec)
  
      def __call__(self, img):
          if self.alphastd == 0:
              return img
  
          alpha = img.new().resize_(3).normal_(0, self.alphastd)
          rgb = self.eigvec.type_as(img).clone() \
              .mul(alpha.view(1, 3).expand(3, 3)) \
              .mul(self.eigval.view(1, 3).expand(3, 3)) \
              .sum(1).squeeze()
  
          return img.add(rgb.view(3, 1, 1).expand_as(img))
  
  
  class Grayscale(object):
  
      def __call__(self, img):
          gs = img.clone()
          gs[0].mul_(0.299).add_(0.587, gs[1]).add_(0.114, gs[2])
          gs[1].copy_(gs[0])
          gs[2].copy_(gs[0])
          return gs
  
  
  class Saturation(object):
  
      def __init__(self, var):
          self.var = var
  
      def __call__(self, img):
          gs = Grayscale()(img)
          alpha = random.uniform(-self.var, self.var)
          return img.lerp(gs, alpha)
  
  
  class Brightness(object):
  
      def __init__(self, var):
          self.var = var
  
      def __call__(self, img):
          gs = img.new().resize_as_(img).zero_()
          alpha = random.uniform(-self.var, self.var)
          return img.lerp(gs, alpha)
  
  
  class Contrast(object):
  
      def __init__(self, var):
          self.var = var
  
      def __call__(self, img):
          gs = Grayscale()(img)
          gs.fill_(gs.mean())
          alpha = random.uniform(-self.var, self.var)
          return img.lerp(gs, alpha)
  
  
  class ColorJitter(object):
  
      def __init__(self, brightness=0.4, contrast=0.4, saturation=0.4):
          self.brightness = brightness
          self.contrast = contrast
          self.saturation = saturation
  
      def __call__(self, img):
          self.transforms = []
          if self.brightness != 0:
              self.transforms.append(Brightness(self.brightness))
          if self.contrast != 0:
              self.transforms.append(Contrast(self.contrast))
          if self.saturation != 0:
              self.transforms.append(Saturation(self.saturation))
  
          random.shuffle(self.transforms)
          transform = Compose(self.transforms)
          # print(transform)
          return transform(img)
  
  
  class LabelSmoothingCrossEntropy(nn.Module):
      def __init__(self):
          super(LabelSmoothingCrossEntropy, self).__init__()
      def forward(self, x, target, smoothing=0.1):
          confidence = 1. - smoothing
          logprobs = F.log_softmax(x, dim=-1)
          nll_loss = -logprobs.gather(dim=-1, index=target.unsqueeze(1))
          nll_loss = nll_loss.squeeze(1)
          smooth_loss = -logprobs.mean(dim=-1)
          loss = confidence * nll_loss + smoothing * smooth_loss
          return loss.mean()
  
  
  
  def initialize_weights(model):
      for m in model.modules():
          if isinstance(m, nn.Conv2d):
              nn.init.kaiming_normal_(m.weight, mode="fan_out",
                                      nonlinearity="relu")
              if m.bias is not None:
                  nn.init.zeros_(m.bias)
          elif isinstance(m, nn.BatchNorm2d):
              nn.init.ones_(m.weight)
              nn.init.zeros_(m.bias)
          elif isinstance(m, nn.Linear):
              nn.init.kaiming_uniform_(m.weight, mode="fan_out",
                                       nonlinearity="sigmoid")
              nn.init.zeros_(m.bias)
  
  
  
  ###Custom CosineAnnealing with Warm-up Restart
  ###Original Code: https://github.com/katsura-jp/pytorch-cosine-annealing-with-warmup/blob/master/cosine_annearing_with_warmup.py
  class CosineAnnealingWarmupRestarts(_LRScheduler):
      """
          optimizer (Optimizer): Wrapped optimizer.
          first_cycle_steps (int): First cycle step size.
          cycle_mult(float): Cycle steps magnification. Default: -1.
          max_lr(float): First cycle's max learning rate. Default: 0.1.
          min_lr(float): Min learning rate. Default: 0.001.
          warmup_steps(int): Linear warmup step size. Default: 0.
          gamma(float): Decrease rate of max learning rate by cycle. Default: 1.
          last_epoch (int): The index of last epoch. Default: -1.
      """
      
      def __init__(self,
                   optimizer : torch.optim.Optimizer,
                   first_cycle_steps : int,
                   cycle_mult : float = 1.,
                   max_lr : float = 0.1,
                   min_lr : float = 0.001,
                   warmup_steps : int = 0,
                   gamma : float = 1.,
                   last_epoch : int = -1
          ):
          assert warmup_steps < first_cycle_steps
          
          self.first_cycle_steps = first_cycle_steps # first cycle step size
          self.cycle_mult = cycle_mult # cycle steps magnification
          self.base_max_lr = max_lr # first max learning rate
          self.max_lr = max_lr # max learning rate in the current cycle
          self.min_lr = min_lr # min learning rate
          self.warmup_steps = warmup_steps # warmup step size
          self.gamma = gamma # decrease rate of max learning rate by cycle
          
          self.cur_cycle_steps = first_cycle_steps # first cycle step size
          self.cycle = 0 # cycle count
          self.step_in_cycle = last_epoch # step size of the current cycle
          
          super(CosineAnnealingWarmupRestarts, self).__init__(optimizer, last_epoch)
          
          # set learning rate min_lr
          self.init_lr()
      
      def init_lr(self):
          self.base_lrs = []
          for param_group in self.optimizer.param_groups:
              param_group['lr'] = self.min_lr
              self.base_lrs.append(self.min_lr)
      
      def get_lr(self):
          if self.step_in_cycle == -1:
              return self.base_lrs
          elif self.step_in_cycle < self.warmup_steps:
              return [(self.max_lr - base_lr)*self.step_in_cycle / self.warmup_steps + base_lr for base_lr in self.base_lrs]
          else:
              return [base_lr + (self.max_lr - base_lr) \
                      * (1 + math.cos(math.pi * (self.step_in_cycle-self.warmup_steps) \
                                      / (self.cur_cycle_steps - self.warmup_steps))) / 2
                      for base_lr in self.base_lrs]
  
      def step(self, epoch=None):
          if epoch is None:
              epoch = self.last_epoch + 1
              self.step_in_cycle = self.step_in_cycle + 1
              if self.step_in_cycle >= self.cur_cycle_steps:
                  self.cycle += 1
                  self.step_in_cycle = self.step_in_cycle - self.cur_cycle_steps
                  self.cur_cycle_steps = int((self.cur_cycle_steps - self.warmup_steps) * self.cycle_mult) + self.warmup_steps
          else:
              if epoch >= self.first_cycle_steps:
                  if self.cycle_mult == 1.:
                      self.step_in_cycle = epoch % self.first_cycle_steps
                      self.cycle = epoch // self.first_cycle_steps
                  else:
                      n = int(math.log((epoch / self.first_cycle_steps * (self.cycle_mult - 1) + 1), self.cycle_mult))
                      self.cycle = n
                      self.step_in_cycle = epoch - int(self.first_cycle_steps * (self.cycle_mult ** n - 1) / (self.cycle_mult - 1))
                      self.cur_cycle_steps = self.first_cycle_steps * self.cycle_mult ** (n)
              else:
                  self.cur_cycle_steps = self.first_cycle_steps
                  self.step_in_cycle = epoch
                  
          self.max_lr = self.base_max_lr * (self.gamma**self.cycle)
          self.last_epoch = math.floor(epoch)
          for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()):
              param_group['lr'] = lr
  