text: '''Some helper functions for PyTorch, including:
      - get_mean_and_std: calculate the mean and std value of dataset.
      - msr_init: net parameter initialization.
      - progress_bar: progress bar mimic xlua.progress.
  '''
  import os
  import sys
  import time
  import math
  
  import torch.nn as nn
  import torch.nn.functional as F
  import torch.nn.init as init
  import torch 
  import numpy as np
  
  import torchvision
  import torchvision.transforms as transforms
  
  class Cutout(object):
      """Randomly mask out one or more patches from an image.
      Args:
          n_holes (int): Number of patches to cut out of each image.
          length (int): The length (in pixels) of each square patch.
      """
      def __init__(self, n_holes, length):
          self.n_holes = n_holes
          self.length = length
  
      def __call__(self, img):
          """
          Args:
              img (Tensor): Tensor image of size (C, H, W).
          Returns:
              Tensor: Image with n_holes of dimension length x length cut out of it.
          """
          h = img.size(1)
          w = img.size(2)
  
          mask = np.ones((h, w), np.float32)
  
          for n in range(self.n_holes):
              y = np.random.randint(h)
              x = np.random.randint(w)
  
              y1 = np.clip(y - self.length // 2, 0, h)
              y2 = np.clip(y + self.length // 2, 0, h)
              x1 = np.clip(x - self.length // 2, 0, w)
              x2 = np.clip(x + self.length // 2, 0, w)
  
              mask[y1: y2, x1: x2] = 0.
  
          mask = torch.from_numpy(mask)
          mask = mask.expand_as(img)
          img = img * mask
  
          return img
  
  
  class BatchMeanCrossEntropyWithLogSoftmax(nn.Module):
      def forward(self, y_hat, y):
          return -(y_hat*y).sum(dim=1).mean(dim=0)
  
  class BatchMeanKLDivWithLogSoftmax(nn.Module):
      def forward(self, p, log_q,  log_p):
          return (p*log_p - p*log_q).sum(dim=1).mean(dim=0)
  
  
  class CrossEntropyWithLogSoftmax(nn.Module):
      def forward(self, y_hat, y):
          return -(y_hat*y).mean()
  
  
  def get_mean_and_std(dataset):
      '''Compute the mean and std value of dataset.'''
      dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=2)
      mean = torch.zeros(3)
      std = torch.zeros(3)
      print('==> Computing mean and std..')
      for inputs, targets in dataloader:
          for i in range(3):
              mean[i] += inputs[:,i,:,:].mean()
              std[i] += inputs[:,i,:,:].std()
      mean.div_(len(dataset))
      std.div_(len(dataset))
      return mean, std
  
  def init_params(net):
      '''Init layer parameters.'''
      for m in net.modules():
          if isinstance(m, nn.Conv2d):
              init.kaiming_normal(m.weight, mode='fan_out')
              if m.bias:
                  init.constant(m.bias, 0)
          elif isinstance(m, nn.BatchNorm2d):
              init.constant(m.weight, 1)
              init.constant(m.bias, 0)
          elif isinstance(m, nn.Linear):
              init.normal(m.weight, std=1e-3)
              if m.bias:
                  init.constant(m.bias, 0)
  
  
  try:
      _, term_width = os.popen('stty size', 'r').read().split()
      term_width = int(term_width)
  except:
      term_width = 40
  
  TOTAL_BAR_LENGTH = 65.
  last_time = time.time()
  begin_time = last_time
  def progress_bar(current, total, msg=None):
      global last_time, begin_time
      if current == 0:
          begin_time = time.time()  # Reset for new bar.
  
      cur_len = int(TOTAL_BAR_LENGTH*current/total)
      rest_len = int(TOTAL_BAR_LENGTH - cur_len) - 1
  
      sys.stdout.write(' [')
      for i in range(cur_len):
          sys.stdout.write('=')
      sys.stdout.write('>')
      for i in range(rest_len):
          sys.stdout.write('.')
      sys.stdout.write(']')
  
      cur_time = time.time()
      step_time = cur_time - last_time
      last_time = cur_time
      tot_time = cur_time - begin_time
  
      L = []
      L.append('  Step: %s' % format_time(step_time))
      L.append(' | Tot: %s' % format_time(tot_time))
      if msg:
          L.append(' | ' + msg)
  
      msg = ''.join(L)
      sys.stdout.write(msg)
      for i in range(term_width-int(TOTAL_BAR_LENGTH)-len(msg)-3):
          sys.stdout.write(' ')
  
      # Go back to the center of the bar.
      for i in range(term_width-int(TOTAL_BAR_LENGTH/2)+2):
          sys.stdout.write('\b')
      sys.stdout.write(' %d/%d ' % (current+1, total))
  
      if current < total-1:
          sys.stdout.write('\r')
      else:
          sys.stdout.write('\n')
      sys.stdout.flush()
  
  def format_time(seconds):
      days = int(seconds / 3600/24)
      seconds = seconds - days*3600*24
      hours = int(seconds / 3600)
      seconds = seconds - hours*3600
      minutes = int(seconds / 60)
      seconds = seconds - minutes*60
      secondsf = int(seconds)
      seconds = seconds - secondsf
      millis = int(seconds*1000)
  
      f = ''
      i = 1
      if days > 0:
          f += str(days) + 'D'
          i += 1
      if hours > 0 and i <= 2:
          f += str(hours) + 'h'
          i += 1
      if minutes > 0 and i <= 2:
          f += str(minutes) + 'm'
          i += 1
      if secondsf > 0 and i <= 2:
          f += str(secondsf) + 's'
          i += 1
      if millis > 0 and i <= 2:
          f += str(millis) + 'ms'
          i += 1
      if f == '':
          f = '0ms'
      return f
  
  def load_data(batch_size=128,is_cifar10=True,batch_test=50):
      transform_train = transforms.Compose([
          transforms.RandomCrop(32, padding=4),
          transforms.RandomHorizontalFlip(),
          transforms.ToTensor(),
          transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
      ])
  
      transform_test = transforms.Compose([
          transforms.ToTensor(),
          transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
      ])
      if is_cifar10:
          dataset = torchvision.datasets.CIFAR10
      else:
          dataset = torchvision.datasets.CIFAR100
      trainset = dataset(root='~/data', train=True, download=True, transform=transform_train)
      trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=8, pin_memory=True)
  
      testset = dataset(root='~/data', train=False, download=True, transform=transform_test)
      testloader = torch.utils.data.DataLoader(testset, batch_size=batch_test, shuffle=False, num_workers=8, pin_memory=True)
  
      return trainloader, testloader
  
  def to_one_hot(inp,num_classes):
      y_onehot = torch.cuda.FloatTensor(inp.size(0), num_classes)
      y_onehot.zero_()
  
      y_onehot.scatter_(1, inp.unsqueeze(1), 1)
      
      return y_onehot
  
  def get_distances(representations):
      rview = representations.view(representations.size(0),-1)
      distances = torch.cdist(rview,rview,p=2)
      return distances
  
  
  
  def representations_to_adj(representations, k=128, A_final=None,mult=None):
      rview = representations.view(representations.size(0),-1)
      rview =  torch.nn.functional.normalize(rview, p=2, dim=1)
      adj = torch.mm(rview,torch.t(rview))
      ind = np.diag_indices(adj.shape[0])
      adj[ind[0], ind[1]] = torch.zeros(adj.shape[0]).cuda()
      degree = torch.pow(adj.sum(dim=1),-0.5)
      degree_matrix = torch.diag(degree)
      adj = torch.matmul(degree_matrix,torch.matmul(adj,degree_matrix))
      if type(mult) == torch.Tensor:
          adj = adj*mult
      if k != 128:
        if type(A_final) == torch.Tensor:
          adj = adj*A_final
        else:
          y, ind = torch.sort(adj, 1)
          A = torch.zeros(*y.size()).cuda()
          k_biggest = ind[:,-k:].data
          for index1,value in enumerate(k_biggest):
              A_line = A[index1]
              A_line[value] = 1
          A_final = torch.min(torch.ones(*y.size()).cuda(),A+torch.t(A))
          adj = adj*A_final
  
      return adj,A_final
  
  def train(net,trainloader,scheduler,device,optimizer,teacher=None,lambda_hkd=0,lambda_gkd=0,lambda_rkd=0,pool3_only=False,temp=4,classes=10,power=1,k=128,intra_only=False,inter_only=False):
      net.train()
      train_loss = 0
      correct = 0
      total = 0
      criterion = BatchMeanCrossEntropyWithLogSoftmax()
      for batch_idx, (inputs, targets) in enumerate(trainloader):
          inputs, targets = inputs.to(device), targets.to(device)
          targets2 = to_one_hot(targets, classes)
          intra_class = torch.matmul(targets2,targets2.T)
          inter_class = 1 - intra_class
          targets = targets2.argmax(dim=1)
          optimizer.zero_grad()
          outputs, layers = net(inputs)
          loss = criterion(F.log_softmax(outputs,dim=-1),targets2)        
          if teacher:
              with torch.no_grad():
                  teacher_output, teacher_layers = teacher(inputs)
              if lambda_hkd > 0:
                  p = F.softmax(teacher_output/temp,dim=-1)
                  log_q = F.log_softmax(outputs/temp,dim=-1)
                  log_p = F.log_softmax(teacher_output/temp,dim=-1)
                  hkd_loss = BatchMeanKLDivWithLogSoftmax()(p=p,log_q=log_q,log_p=log_p)
                  loss += lambda_hkd*hkd_loss
              if lambda_rkd > 0:
                  loss_rkd = 0
                  zips = zip(layers,teacher_layers) if not pool3_only else zip([layers[-1]],[teacher_layers[-1]])
                  for student_layer,teacher_layer in zips:
                  
                      distances_teacher = get_distances(teacher_layer)
                      distances_teacher = distances_teacher[distances_teacher>0]
                      mean_teacher = distances_teacher.mean()
                      distances_teacher = distances_teacher/mean_teacher
  
                      distances_student = get_distances(student_layer)
                      distances_student = distances_student[distances_student>0]
                      mean_student = distances_student.mean()
                      distances_student = distances_student/mean_student
                      loss_rkd += lambda_rkd*F.smooth_l1_loss(distances_student, distances_teacher, reduction='none').mean()
                  loss += loss_rkd if pool3_only else loss_rkd/3
              elif lambda_gkd > 0:
                  loss_gkd = do_gkd(pool3_only, layers, teacher_layers, k, power, intra_only, lambda_gkd, intra_class, inter_only, inter_class)
                  loss += loss_gkd if pool3_only else loss_gkd/3
          loss.backward()
          optimizer.step()
          train_loss += loss.item()
          _, predicted = outputs.max(1)
          total += targets.size(0)
          correct += predicted.eq(targets).sum().item()
  
          progress_bar(batch_idx, len(trainloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
              % (train_loss/(batch_idx+1), 100.*correct/total, correct, total))
      scheduler.step()
  
  def do_gkd(pool3_only, layers, teacher_layers, k, power, intra_only, lambda_gkd, intra_class, inter_only, inter_class):
      loss_gkd = 0 
      zips = zip(layers,teacher_layers) if not pool3_only else zip([layers[-1]],[teacher_layers[-1]])
      mult = None
      if intra_only:
          mult=intra_class
      elif inter_only:
          mult=inter_class
      for student_layer,teacher_layer in zips:
          adj_teacher,A_final = representations_to_adj(teacher_layer,k,mult=mult)
          adj_student,A_final = representations_to_adj(student_layer,k,A_final,mult=mult)
          adj_teacher_p = adj_teacher
          adj_student_p = adj_student
          for _ in range(power-1):
              adj_teacher_p = torch.matmul(adj_teacher_p,adj_teacher)
              adj_student_p = torch.matmul(adj_student_p,adj_student)
          loss_gkd += lambda_gkd*F.mse_loss(adj_teacher_p, adj_student_p, reduction='none').sum()
      return loss_gkd
  
  def test(net,testloader, device,save_name="teacher",show="accuracy"):
      net.eval()
      test_loss = 0
      correct = 0
      total = 0
      criterion = nn.CrossEntropyLoss()
  
      with torch.no_grad():
          for batch_idx, (inputs, targets) in enumerate(testloader):
              inputs, targets = inputs.to(device), targets.to(device)
              outputs, layers = net(inputs)
              loss = criterion(outputs, targets)
  
              test_loss += loss.item()
              _, predicted = outputs.max(1)
              total += targets.size(0)
              correct += predicted.eq(targets).sum().item()
  
              if show=="accuracy":
                  progress_bar(batch_idx, len(testloader), 'Loss: %.3f | Acc: %.3f%% (%d/%d)'
                      % (test_loss/(batch_idx+1), 100.*correct/total, correct, total))                
      state = {
          'net': net,
      }
      if not os.path.isdir('checkpoint'):
          os.mkdir('checkpoint')
      if save_name != "no":
          torch.save(state, './checkpoint/{}.pth'.format(save_name))
      if show=="error":
          print("Test error: {:.2f}".format(100 - 100.*correct/total))
  