text: import math
  import torch.nn as nn
  import torch.nn.functional as nnF
  from typing import Tuple, List, Union
  import os
  import numpy as np
  import torch
  import torchvision
  import torch.utils.data as Data
  from torch.utils.data.dataset import Dataset
  from torchvision import datasets, transforms
  from PIL import Image
  from torch.nn import Parameter
  from spectral_clustering import spectral_clustering, pairwise_cosine_similarity, KMeans
  import time
  import torchvision.transforms.functional as F
  import sys
  from typing import Type, Any, Callable, Union, List, Optional
  from torch import Tensor
  import copy
  import cv2
  ###########################
  ###set initial prameters###
  ###########################
  initializing = False # if initializing, save a .pkl files, else read saved .pkl files
  AMP= True 
  CLD = True # use cld loss
  maxLambda = 2.0
  increaseLambda = False
  
  
  BYOL = False
  if BYOL:
      FrameWork = 'BYOL'
  else:
      FrameWork = 'SimSiam'
  
  RGB_3F = True # use RGB channels for past future current
  
  if RGB_3F:
      temporal = '_RGB3F'
      p_RGB_3F = 1.0
  else:
      temporal = ''
      p_RGB_3F = 0.0
  
  innerShuffle = True #only use this for data with dramatica batch effect, like mice data. Don't use this with fly!
  if innerShuffle:
      Ishuffle = '_innerShuffle_'
  else:
      Ishuffle = '_randomShuffle_'
  
  
  num_workers = 16
  batch_size = 128
  frame_interval = 2 #2 for mice and 1 for flies
  if AMP:
      batch_size=batch_size*2
  proj_size = [2048,2048,2048,2048]
  pred_size = [2048,512,2048]
  
  steps = 20000  
  input_size = [256,192] #for flies, use 224,224 
  k_eigen=10
  clusters=10
  num_iters=10
  
  base_lr=0.05 #per batchsize256
  videoSets = '' #fill with the dir name of your dataset
  modelName = FrameWork+'_ResNet50_maxLambda'+str(maxLambda)+temporal+str(p_RGB_3F)+Ishuffle+videoSets+"_"
  
  # change '/run/media/harold/Data_Storage/' to where your data are stored
  ValidDir = '/run/media/harold/Data_Storage/'+videoSets+'/Valid_Set'
  TrainDir = '/run/media/harold/Data_Storage/'+videoSets+'/Train_Set'
  TestDir = '/run/media/harold/Data_Storage/'+videoSets+'/Test_Set'
  
  SavedDir=os.path.dirname(__file__)+"/Saved_Models/"+FrameWork+"_CLD_ResNet50_initial_FlyCourtship.pkl"
  CheckpointDir=os.path.dirname(__file__)+"/Saved_Models/"+modelName
  log=os.path.dirname(__file__)+"/Saved_Models/"+modelName
  embeddingName="/embedded_by_"+modelName+"steps"
  ValidEmbedDir = os.path.dirname(__file__)+"/ValidEmbed"
  
  f = open(log+"train.log",'a')
  localtime = time.asctime( time.localtime(time.time()) )
  f.write(str(localtime)+'\n')
  f.write('frame_interval = '+str(frame_interval)+'\n')
  f.write("Learning Rate per 256 images = "+str(base_lr)+'\n')
  f.write("Initialte weights with "+SavedDir.split('/')[-1]+'\n')
  f.write("max Lambda = "+str(maxLambda)+'\n')
  f.write("Lambda increase in cosine style = "+str(increaseLambda)+'\n')
  f.close()
  f = open(log+"test.log",'a')
  localtime = time.asctime( time.localtime(time.time()) )
  f.write(str(localtime)+'\n')
  f.close()
  
  class NormedLinear(nn.Module):
      def __init__(self, in_features, out_features):
          super(NormedLinear, self).__init__()
          self.weight = Parameter(torch.Tensor(in_features, out_features))
          self.weight.data.uniform_(-1, 1).renorm_(2, 1, 1e-5).mul_(1e5)
  
      def forward(self, x):
          out = nnF.normalize(x, dim=1).mm(nnF.normalize(self.weight, dim=0))
          return out
  
  class Normalize(nn.Module):
  
      def __init__(self, power=2):
          super(Normalize, self).__init__()
          self.power = power
      
      def forward(self, x):
          norm = x.pow(self.power).sum(1, keepdim=True).pow(1./self.power)
          out = x.div(norm)
          return out
  
  class AddProjector(nn.Module):
  
      def __init__(self,backbone,proj_size=[512,512,512]):
          super(AddProjector, self).__init__()  
          self.backbone = backbone
          self.backbone.fc = nn.Identity()   
          proj_layer=[]
          for i in range(len(proj_size) - 2):
              proj_layer.append(nn.Linear(proj_size[i], proj_size[i + 1]))
              proj_layer.append(nn.BatchNorm1d(proj_size[i + 1]))
              proj_layer.append(nn.ReLU(inplace=True))
          proj_layer.append(nn.Linear(proj_size[-2], proj_size[-1]))
          proj_layer.append(nn.BatchNorm1d(proj_size[-1]))
          self.projector = nn.Sequential(*proj_layer)
  
      def forward(self, x):
          out = self.backbone(x)
          out = self.projector(out)
          return out
  
  class PredAndCLD(torch.nn.Module):
      def __init__(self,pred_size=[512,128,512]):
          super(PredAndCLD, self).__init__()  
          #branches for BYOL
          self.predictor = nn.Sequential(
              nn.Linear(pred_size[0], pred_size[1]),
              nn.BatchNorm1d(pred_size[1]),
              nn.ReLU(inplace=True),
              nn.Linear(pred_size[1], pred_size[2]))
  
          #branches for CLD
          self.groupDis = nn.Sequential(
              NormedLinear(pred_size[0], pred_size[1]*2),
              Normalize(2))
          
      def forward(self, x, online = True):
            
          CLD = self.groupDis(x)
  
          if online:  #for online network only
              pred = self.predictor(x)
              return pred, CLD
  
          return CLD
  
  # exponential moving average
  class EMA():
      def __init__(self, beta):
          super().__init__()
          self.beta = beta
  
      def update_average(self, old, new):
          if old is None:
              return new
          return old * self.beta + (1 - self.beta) * new
  
  def update_moving_average(ema_updater, ma_model, current_model):
      for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()):
          old_weight, up_weight = ma_params.data, current_params.data
          ma_params.data = ema_updater.update_average(old_weight, up_weight)
  
  def set_requires_grad(model, val):
      for p in model.parameters():
          p.requires_grad = val
          
  def update_moving_average(ema_updater, ma_model, current_model):
      for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()):
          old_weight, up_weight = ma_params.data, current_params.data
          ma_params.data = ema_updater.update_average(old_weight, up_weight)
  
  class CLD_Byol(nn.Module):
  
      def __init__(self, backbone,predictor,CLD,moving_average_decay = 0.99):
          super().__init__()
          self.online_encoder = backbone
          self.target_encoder = self._get_target_encoder()
          self.CLD = CLD
          self.predictor = predictor
          self.target_ema_updater =  EMA(moving_average_decay)
  
      def _get_target_encoder(self):
          target_encoder = copy.deepcopy(self.online_encoder)
          set_requires_grad(target_encoder, False)
          return target_encoder
  
      def update_moving_average(self):
          update_moving_average(self.target_ema_updater, self.target_encoder, self.online_encoder)
  
      def loss_fn_SS(self,z1,z2): #z1 = (emb1,pred1) z2 = (emb2,pred2)
          t1, p1 = z1
          t2, p2 = z2
          return 2 - nnF.cosine_similarity(p1, t2.detach(), dim=-1).mean() - nnF.cosine_similarity(p2, t1.detach(), dim=-1).mean()
  
      def loss_fn_cld(self,z1,z2,cld_t=0.07): # z1 = (img1_online_cld,img1_target_cld) z2 = (img2_online_cld,img2_target_cld)
          cluster_label1_1, centroids1_1 = KMeans(z1[0], K=clusters, Niters=num_iters)
          cluster_label2_1, centroids2_1 = KMeans(z2[1], K=clusters, Niters=num_iters)       
          affnity1_1 = torch.mm(z1[0], centroids2_1.t()) 
          affnity2_1 = torch.mm(z2[1], centroids1_1.t())
          loss_CLD_1 = 0.5*(nnF.cross_entropy(affnity1_1.div_(cld_t), cluster_label2_1)+nnF.cross_entropy(affnity2_1.div_(cld_t), cluster_label1_1))
          cluster_label1_2, centroids1_2 = KMeans(z1[1], K=clusters, Niters=num_iters)
          cluster_label2_2, centroids2_2 = KMeans(z2[0], K=clusters, Niters=num_iters)       
          affnity1_2 = torch.mm(z1[1], centroids2_2.t()) 
          affnity2_2 = torch.mm(z2[0], centroids1_2.t())
          loss_CLD_2 = 0.5*(nnF.cross_entropy(affnity1_2.div_(cld_t), cluster_label2_2)+nnF.cross_entropy(affnity2_2.div_(cld_t), cluster_label1_2))
          return 0.5*(loss_CLD_1+loss_CLD_2)
  
      def forward(self, img1, img2,Lambda):
          y1 = self.online_encoder(img1)
          y2 = self.online_encoder(img2)
          pred1, online_cld1 = self.predictor(y1)
          pred2, online_cld2 = self.predictor(y2)
  
          with torch.no_grad():
              emb1 = self.target_encoder(img1)
              emb2 = self.target_encoder(img2)
              target_cld1 = self.predictor(emb1,online = False)
              target_cld2 = self.predictor(emb2,online = False)
              #target_cld1.detach_()
              #target_cld2.detach_()
  
          loss_SS = self.loss_fn_SS((emb1,pred1),(emb2,pred2))
          if self.CLD:       
              loss_CLD = self.loss_fn_cld((online_cld1,target_cld1),(online_cld2,target_cld2))
          else:
              loss_CLD=torch.tensor(0)
  
          loss= loss_SS+Lambda*loss_CLD
          return loss, loss_SS, loss_CLD, y1, online_cld1
          
      @torch.no_grad()
      def infer(self,x):
          return self.online_encoder(x)
  
  class CLD_SimSiam(nn.Module):
      def __init__(self, backbone,predictor,CLD):
          super().__init__()
          self.backbone = backbone
          self.loss_fn_cld = nn.CrossEntropyLoss()
          self.CLD = CLD
          self.predictor = predictor
  
      def loss_fn_SS(self,z1,z2):
          t1, p1 = z1
          t2, p2 = z2
          return 2 - nnF.cosine_similarity(p1, t2.detach(), dim=-1).mean() - nnF.cosine_similarity(p2, t1.detach(), dim=-1).mean()
  
      def forward(self, y1, y2,Lambda,cld_t=0.07):
          y1 = self.backbone(y1)
          y2 = self.backbone(y2)
          z1_H,z1_L = self.predictor(y1)
          z2_H,z2_L = self.predictor(y2)
          loss_SS = self.loss_fn_SS((y1,z1_H),(y2,z2_H))
          if self.CLD:
              cluster_label1, centroids1 = KMeans(z1_L, K=clusters, Niters=num_iters)
              cluster_label2, centroids2 = KMeans(z2_L, K=clusters, Niters=num_iters)       
              affnity1 = torch.mm(z1_L, centroids2.t()) 
              affnity2 = torch.mm(z2_L, centroids1.t())
              loss_CLD = 0.5*(self.loss_fn_cld(affnity1.div_(cld_t), cluster_label2)+self.loss_fn_cld(affnity2.div_(cld_t), cluster_label1))
          else:
              loss_CLD=torch.tensor(0)
  
          loss= loss_SS+Lambda*loss_CLD
          return loss, loss_SS, loss_CLD, z1_H, z1_L
          
      @torch.no_grad()
      def infer(self,x):
          return self.backbone(x)
  
  ################
  ###Dataloader###
  ################
  
  TurboColor =  np.array([[ 48,  18,  59],
                          [ 50,  21,  67],
                          [ 51,  24,  74],
                          [ 52,  27,  81],
                          [ 53,  30,  88],
                          [ 54,  33,  95],
                          [ 55,  36, 102],
                          [ 56,  39, 109],
                          [ 57,  42, 115],
                          [ 58,  45, 121],
                          [ 59,  47, 128],
                          [ 60,  50, 134],
                          [ 61,  53, 139],
                          [ 62,  56, 145],
                          [ 63,  59, 151],
                          [ 63,  62, 156],
                          [ 64,  64, 162],
                          [ 65,  67, 167],
                          [ 65,  70, 172],
                          [ 66,  73, 177],
                          [ 66,  75, 181],
                          [ 67,  78, 186],
                          [ 68,  81, 191],
                          [ 68,  84, 195],
                          [ 68,  86, 199],
                          [ 69,  89, 203],
                          [ 69,  92, 207],
                          [ 69,  94, 211],
                          [ 70,  97, 214],
                          [ 70, 100, 218],
                          [ 70, 102, 221],
                          [ 70, 105, 224],
                          [ 70, 107, 227],
                          [ 71, 110, 230],
                          [ 71, 113, 233],
                          [ 71, 115, 235],
                          [ 71, 118, 238],
                          [ 71, 120, 240],
                          [ 71, 123, 242],
                          [ 70, 125, 244],
                          [ 70, 128, 246],
                          [ 70, 130, 248],
                          [ 70, 133, 250],
                          [ 70, 135, 251],
                          [ 69, 138, 252],
                          [ 69, 140, 253],
                          [ 68, 143, 254],
                          [ 67, 145, 254],
                          [ 66, 148, 255],
                          [ 65, 150, 255],
                          [ 64, 153, 255],
                          [ 62, 155, 254],
                          [ 61, 158, 254],
                          [ 59, 160, 253],
                          [ 58, 163, 252],
                          [ 56, 165, 251],
                          [ 55, 168, 250],
                          [ 53, 171, 248],
                          [ 51, 173, 247],
                          [ 49, 175, 245],
                          [ 47, 178, 244],
                          [ 46, 180, 242],
                          [ 44, 183, 240],
                          [ 42, 185, 238],
                          [ 40, 188, 235],
                          [ 39, 190, 233],
                          [ 37, 192, 231],
                          [ 35, 195, 228],
                          [ 34, 197, 226],
                          [ 32, 199, 223],
                          [ 31, 201, 221],
                          [ 30, 203, 218],
                          [ 28, 205, 216],
                          [ 27, 208, 213],
                          [ 26, 210, 210],
                          [ 26, 212, 208],
                          [ 25, 213, 205],
                          [ 24, 215, 202],
                          [ 24, 217, 200],
                          [ 24, 219, 197],
                          [ 24, 221, 194],
                          [ 24, 222, 192],
                          [ 24, 224, 189],
                          [ 25, 226, 187],
                          [ 25, 227, 185],
                          [ 26, 228, 182],
                          [ 28, 230, 180],
                          [ 29, 231, 178],
                          [ 31, 233, 175],
                          [ 32, 234, 172],
                          [ 34, 235, 170],
                          [ 37, 236, 167],
                          [ 39, 238, 164],
                          [ 42, 239, 161],
                          [ 44, 240, 158],
                          [ 47, 241, 155],
                          [ 50, 242, 152],
                          [ 53, 243, 148],
                          [ 56, 244, 145],
                          [ 60, 245, 142],
                          [ 63, 246, 138],
                          [ 67, 247, 135],
                          [ 70, 248, 132],
                          [ 74, 248, 128],
                          [ 78, 249, 125],
                          [ 82, 250, 122],
                          [ 85, 250, 118],
                          [ 89, 251, 115],
                          [ 93, 252, 111],
                          [ 97, 252, 108],
                          [101, 253, 105],
                          [105, 253, 102],
                          [109, 254,  98],
                          [113, 254,  95],
                          [117, 254,  92],
                          [121, 254,  89],
                          [125, 255,  86],
                          [128, 255,  83],
                          [132, 255,  81],
                          [136, 255,  78],
                          [139, 255,  75],
                          [143, 255,  73],
                          [146, 255,  71],
                          [150, 254,  68],
                          [153, 254,  66],
                          [156, 254,  64],
                          [159, 253,  63],
                          [161, 253,  61],
                          [164, 252,  60],
                          [167, 252,  58],
                          [169, 251,  57],
                          [172, 251,  56],
                          [175, 250,  55],
                          [177, 249,  54],
                          [180, 248,  54],
                          [183, 247,  53],
                          [185, 246,  53],
                          [188, 245,  52],
                          [190, 244,  52],
                          [193, 243,  52],
                          [195, 241,  52],
                          [198, 240,  52],
                          [200, 239,  52],
                          [203, 237,  52],
                          [205, 236,  52],
                          [208, 234,  52],
                          [210, 233,  53],
                          [212, 231,  53],
                          [215, 229,  53],
                          [217, 228,  54],
                          [219, 226,  54],
                          [221, 224,  55],
                          [223, 223,  55],
                          [225, 221,  55],
                          [227, 219,  56],
                          [229, 217,  56],
                          [231, 215,  57],
                          [233, 213,  57],
                          [235, 211,  57],
                          [236, 209,  58],
                          [238, 207,  58],
                          [239, 205,  58],
                          [241, 203,  58],
                          [242, 201,  58],
                          [244, 199,  58],
                          [245, 197,  58],
                          [246, 195,  58],
                          [247, 193,  58],
                          [248, 190,  57],
                          [249, 188,  57],
                          [250, 186,  57],
                          [251, 184,  56],
                          [251, 182,  55],
                          [252, 179,  54],
                          [252, 177,  54],
                          [253, 174,  53],
                          [253, 172,  52],
                          [254, 169,  51],
                          [254, 167,  50],
                          [254, 164,  49],
                          [254, 161,  48],
                          [254, 158,  47],
                          [254, 155,  45],
                          [254, 153,  44],
                          [254, 150,  43],
                          [254, 147,  42],
                          [254, 144,  41],
                          [253, 141,  39],
                          [253, 138,  38],
                          [252, 135,  37],
                          [252, 132,  35],
                          [251, 129,  34],
                          [251, 126,  33],
                          [250, 123,  31],
                          [249, 120,  30],
                          [249, 117,  29],
                          [248, 114,  28],
                          [247, 111,  26],
                          [246, 108,  25],
                          [245, 105,  24],
                          [244, 102,  23],
                          [243,  99,  21],
                          [242,  96,  20],
                          [241,  93,  19],
                          [240,  91,  18],
                          [239,  88,  17],
                          [237,  85,  16],
                          [236,  83,  15],
                          [235,  80,  14],
                          [234,  78,  13],
                          [232,  75,  12],
                          [231,  73,  12],
                          [229,  71,  11],
                          [228,  69,  10],
                          [226,  67,  10],
                          [225,  65,   9],
                          [223,  63,   8],
                          [221,  61,   8],
                          [220,  59,   7],
                          [218,  57,   7],
                          [216,  55,   6],
                          [214,  53,   6],
                          [212,  51,   5],
                          [210,  49,   5],
                          [208,  47,   5],
                          [206,  45,   4],
                          [204,  43,   4],
                          [202,  42,   4],
                          [200,  40,   3],
                          [197,  38,   3],
                          [195,  37,   3],
                          [193,  35,   2],
                          [190,  33,   2],
                          [188,  32,   2],
                          [185,  30,   2],
                          [183,  29,   2],
                          [180,  27,   1],
                          [178,  26,   1],
                          [175,  24,   1],
                          [172,  23,   1],
                          [169,  22,   1],
                          [167,  20,   1],
                          [164,  19,   1],
                          [161,  18,   1],
                          [158,  16,   1],
                          [155,  15,   1],
                          [152,  14,   1],
                          [149,  13,   1],
                          [146,  11,   1],
                          [142,  10,   1],
                          [139,   9,   2],
                          [136,   8,   2],
                          [133,   7,   2],
                          [129,   6,   2],
                          [126,   5,   2],
                          [122,   4,   3]], dtype=np.uint8)
  
  def Turbo(img):
      np_img=np.array(img)
      b, g ,r =cv2.split(np_img) 
      b_color=np.array([TurboColor[xi,0] for xi in b])
      g_color=np.array([TurboColor[xi,1] for xi in g])
      r_color=np.array([TurboColor[xi,2] for xi in r])
      np_color = cv2.merge([b_color,g_color,r_color])
      return Image.fromarray(np_color)
  
  def RandomTransformed(img,resize=input_size,min_scale=0.7,degrees=37,kernel=2,p_hflip=0.5,p_vflip=0.5,p_turbo=0.5,brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1, is_train=True):
      #37 degree for XXH 45 for others
      if is_train:
          width, height = F._get_image_size(img)
          w = torch.randint(int(width*min_scale), width, size=(1,)).item()
          h = torch.randint(int(height*min_scale), height, size=(1,)).item()
          i = torch.randint(0, height - h + 1, size=(1,)).item()
          j = torch.randint(0, width - w + 1, size=(1,)).item()
          angle = float(torch.empty(1).uniform_(float(-degrees), float(degrees)).item())
          r1 = torch.rand(1)
          r2 = torch.rand(1)
          r3 = torch.rand(1)
          #blur_factor = torch.randint(0, kernel, size=(1,)).item() *2 +1
          brightness_factor = float(torch.empty(1).uniform_(1-brightness, 1+brightness))
          contrast_factor = float(torch.empty(1).uniform_(1-contrast, 1+contrast))
          saturation_factor = float(torch.empty(1).uniform_(1-saturation, 1+saturation))
          hue_factor = float(torch.empty(1).uniform_(-hue, hue))
          fn_idx = torch.randperm(4)
          img = F.resized_crop(img, i, j, h, w, resize)
          img = F.rotate(img, angle)
          if r1 < p_hflip:
              img = F.hflip(img)
          if r2 < p_vflip:
              img = F.vflip(img)
          if r3 < p_turbo:
              img = Turbo(img)
          #img = F.gaussian_blur(img,(blur_factor,blur_factor))
          for fn_id in fn_idx:
              if fn_id == 0:
                  img = F.adjust_brightness(img, brightness_factor)
              elif fn_id == 1:
                  img = F.adjust_contrast(img, contrast_factor)
              elif fn_id == 2:
                  img = F.adjust_saturation(img, saturation_factor)
              elif fn_id == 3:
                  img = F.adjust_hue(img, hue_factor)
      else:
          img = F.resize(img,resize)
      img = F.to_tensor(img)
      return img
  
  def pil_loader(path: str) -> Image.Image:
      # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
      with open(path, 'rb') as f:
          img = Image.open(f)
          return img.convert('RGB')
  
  class preShuffleImageDataset(Dataset): 
      '''
      your folder should have following structrues:
      
      root
          -folder1
              --1.jpg
              --2.jpg
              --...
          -folder2
              --1.jpg
              --2.jpg
              --...
          -...
      This dataset should be sampled sequentially. It return batches within same videos respectively. 
      That means, for each batch, we only sampled within one video to avoid batch effect.
      '''
      def __init__(self,
                   root: str,
                   transform = None,
                   p_RGB_3F = p_RGB_3F,
                   frame_interval = frame_interval, # interval = 2 equals to 30fps -> 6fps
                   batchsize = batch_size,
                   innerShuffle = innerShuffle,
                   shuffle = True):
          self.p_RGB_3F = p_RGB_3F
          self.frame_interval = frame_interval
          self.batchsize = batchsize
          self.transform = transform
          self.videoList, self.videoLen = self.videoList(root)
          self.innerShuffle = innerShuffle
          if shuffle:
              self.list = self.preShuffleSampleList()
          else:
              self.list = self.SampleList()
  
      def videoList(self,root):
          dirlist=os.listdir(root)
          videoList = []
          videoLen = []
          for dirs in dirlist:
              if os.path.isdir(os.path.join(root,dirs)):
                  filenum=len(os.listdir(os.path.join(root,dirs)))
                  videoList.append(os.path.join(root,dirs))
                  videoLen.append(filenum)
          return videoList, videoLen
  
      def preShuffleSampleList(self):  
          if self.innerShuffle: #force to use innerShuffle for mice data due to batch effects
              for i,lens in enumerate(self.videoLen):
                  droplastLen = lens // self.batchsize * self.batchsize
                  videoIndex = i * torch.ones(droplastLen).unsqueeze(1)
                  imgIndex = torch.randperm(droplastLen).unsqueeze(1)
                  try:
                      indexMat = torch.cat((indexMat,torch.cat((videoIndex,imgIndex),1).reshape(-1,self.batchsize,2)),0)
                  except:
                      indexMat = torch.cat((videoIndex,imgIndex),1).reshape(-1,self.batchsize,2)
              indexMat=indexMat[torch.randperm(len(indexMat))]
          else:
              for i,lens in enumerate(self.videoLen):
                  droplastLen = lens // self.batchsize * self.batchsize
                  videoIndex = i * torch.ones(droplastLen).unsqueeze(1)
                  imgIndex = torch.arange(droplastLen).unsqueeze(1)
                  try:
                      indexMat = torch.cat((indexMat,torch.cat((videoIndex,imgIndex),1)),0)
                  except:
                      indexMat = torch.cat((videoIndex,imgIndex),1)
              indexMat=indexMat[torch.randperm(len(indexMat))]
              indexMat=indexMat.reshape(-1,self.batchsize,2)
  
          for _ in range(len(indexMat)):
              p_rgb =  torch.rand(1)*torch.ones(self.batchsize).unsqueeze(1).unsqueeze(0)
              try:
                  probMat = torch.cat((probMat,p_rgb),0)
              except:
                  probMat = p_rgb
          indexMat = torch.cat((indexMat,probMat),2)
          indexMat=indexMat.reshape(-1,3)
          return indexMat.tolist()
      
      def SampleList(self):
          for i,lens in enumerate(self.videoLen):
              droplastLen = lens
              videoIndex = i * torch.ones(droplastLen).unsqueeze(1)
              imgIndex = torch.arange(droplastLen).unsqueeze(1)
              try:
                  indexMat = torch.cat((indexMat,torch.cat((videoIndex,imgIndex),1)),0)
              except:
                  indexMat = torch.cat((videoIndex,imgIndex),1)
          probMat = torch.zeros(len(indexMat)).unsqueeze(1)
          indexMat = torch.cat((indexMat,probMat),1)
          return indexMat.tolist()
  
      def __getitem__(self, index):
          path ,num ,p_rgb = self.list[index]
          path = int(path)
          num = int (num)
          #print(path ,num ,p_rgb)
          if p_rgb < self.p_RGB_3F:
              current = cv2.imread(self.videoList[path]+'/'+str(num+1)+'.jpg', cv2.IMREAD_GRAYSCALE)
              past = cv2.imread(self.videoList[path]+'/'+str(num-self.frame_interval+1)+'.jpg', cv2.IMREAD_GRAYSCALE)
              if past is None:
                  past = current
              future = cv2.imread(self.videoList[path]+'/'+str(num+self.frame_interval+1)+'.jpg', cv2.IMREAD_GRAYSCALE)
              if future is None:
                  future = current
              img = cv2.merge([past,current,future])
              img = Image.fromarray(img)
          else:
              img = pil_loader(self.videoList[path]+'/'+str(num+1)+'.jpg')
          if self.transform != None:
              img = self.transform(img)
          return img,self.videoList[path]+'/'+str(num+1)+'.jpg'
  
      def __len__(self):
          return len(self.list)
  
  class CollateFunction(nn.Module):
  
      def __init__(self, is_train: bool = True):
  
          super(CollateFunction, self).__init__()
          self.is_train = is_train
  
      def forward(self, batch: List[tuple]):
   
          batch_size = len(batch)
          # list of labels
          labels = [item[1] for item in batch]
          if self.is_train:
  
              # list of transformed images
              transformed = [RandomTransformed(batch[i % batch_size][0]).unsqueeze_(0)
                          for i in range(2 * batch_size)]
              # tuple of transforms
              transformed = (
                  torch.cat(transformed[:batch_size], 0),
                  torch.cat(transformed[batch_size:], 0)
              )
  
          else:
              transformed = torch.cat([RandomTransformed(batch[i % batch_size][0],is_train=False).unsqueeze_(0)
                          for i in range(batch_size)])
              # tuple of transforms            
          return transformed, labels
  
  train_set=preShuffleImageDataset(TrainDir,innerShuffle=innerShuffle)
  print("TrainSet:"+str(train_set.__len__()))
  test_set=preShuffleImageDataset(TestDir,innerShuffle=innerShuffle)
  print("TestSet:"+str(test_set.__len__()))
  valid_set=preShuffleImageDataset(ValidDir,shuffle=False)
  print("ValidSet:"+str(valid_set.__len__()))
  
  TrainLoader = torch.utils.data.DataLoader(
      train_set,
      batch_size=batch_size,
      shuffle= False,
      drop_last=True,
      num_workers=16,
      collate_fn= CollateFunction(is_train=True)
  )
    
  TestLoader = torch.utils.data.DataLoader(
      test_set,
      batch_size=batch_size,
      shuffle= False,
      drop_last=True,
      num_workers=16,
      collate_fn= CollateFunction(is_train=True)
  )
  
  ValidLoader = torch.utils.data.DataLoader(
      valid_set,
      batch_size=batch_size,
      shuffle=False,
      drop_last=False,
      num_workers=16,
      collate_fn= CollateFunction(is_train=False)
  )
  
  #################
  ###Train loops###
  #################
  
  ResNet = torchvision.models.resnet50(pretrained = True)
  ResNetAndProj = AddProjector(ResNet,proj_size)
  predictor = PredAndCLD(pred_size)
  
  if BYOL:
      model = CLD_Byol(ResNetAndProj,predictor,CLD)
  else:
      model = CLD_SimSiam(ResNetAndProj,predictor,CLD)
  
  model.eval() ##always save and load under eval()
  if initializing:
      torch.save(model.state_dict(),SavedDir)
  
  model.load_state_dict(torch.load(SavedDir),strict=False)
  model.train()    #This behaviour can cause problems when storing the state_dict() of a model while in a mode and lately loading it in a model with a different mode, as the attributes of this class change. To avoid this issue, we recommend converting the model to eval mode before storing or loading the state dictionary.
  device =  torch.device('cuda')
  model.to(device)
  
  #SDG optimizer and adjust lr
  optimizer = torch.optim.SGD(model.parameters(),
                              lr=base_lr*batch_size/256,
                              momentum=0.9,
                              weight_decay=1e-4)
  
  scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=batch_size*base_lr/256, total_steps=steps,pct_start=0.025)
  
  avg_loss = np.array([0.,0.,0.])
  avg_collapse_level = np.array([0.,0.])
  
  # main training loop
  TrainGenerator = iter(TrainLoader)
  TestGenerator = iter(TestLoader)
  
  scaler=torch.cuda.amp.GradScaler(enabled=AMP)
  time_start=time.time()
  
  for i in range(steps):
  
      if increaseLambda:
          Lambda = maxLambda*0.5*(1-np.cos(i*np.pi/steps))
      else:
          Lambda = maxLambda
  
      optimizer.zero_grad()
  
      try:
          # Samples the batch
          (x0, x1), _= next(TrainGenerator)
  
      except StopIteration:
          # restart the generator if the previous generator is exhausted.
          train_set=preShuffleImageDataset(TrainDir,innerShuffle=innerShuffle)
          TrainLoader = torch.utils.data.DataLoader(
              train_set,
              batch_size=batch_size,
              shuffle= False,
              drop_last=True,
              num_workers=16,
              collate_fn= CollateFunction(is_train=True)
          )
          TrainGenerator = iter(TrainLoader)
          (x0, x1), _ = next(TrainGenerator)
  
      # move images to the gpu
      x0 = x0.to(device)
      x1 = x1.to(device)
  
      with torch.cuda.amp.autocast(enabled=AMP):
          loss, loss_SS, loss_CLD,out_H,output_L=model(x0,x1,Lambda)
  
      #loss.backward()
      scaler.scale(loss).backward()
      #optimizer.step()
      scaler.step(optimizer)
      scaler.update()
      if BYOL:
          model.update_moving_average()
      #use one cycle scheduler
      scheduler.step()
      
      # calculate the per-dimension standard deviation of the outputs
      # we can use this later to check whether the embeddings are collapsing
      output_H = out_H.detach()
      output_H = torch.nn.functional.normalize(output_H, dim=1)
      output_H_std = torch.std(output_H, 0)
      output_H_std = output_H_std.mean()
  
      output_L = output_L.detach()
      output_L = torch.nn.functional.normalize(output_L, dim=1)
      output_L_std = torch.std(output_L, 0)
      output_L_std = output_L_std.mean()
      loss_array = np.array([loss.item(),loss_SS.item(),loss_CLD.item()])
      collapse_level = np.array([1-math.sqrt(proj_size[-1]) * output_H_std.item(),1-math.sqrt(pred_size[1]*2) * output_L_std.item()])
      # use moving averages to track the loss and standard deviation
      w = 0.9
      avg_loss = avg_loss * w + loss_array * (1 - w)
      avg_collapse_level = avg_collapse_level * w + collapse_level * (1 - w)
      #Lambda = k * loss_array[1] / loss_array[2]
  
      if (i+1)  % 5000 == 0 or i == 0:
          ###validation
          model.eval()
          torch.save(model.state_dict(),CheckpointDir+str(i+1)+".pkl")
          output=[]
          for img, fname in ValidLoader:
              img = img.to(device)
              with torch.cuda.amp.autocast(enabled=AMP):
                  embeddings=model.infer(img)
              for j in range(len(embeddings)):
                  output.append((embeddings[j].cpu().numpy(),fname[j]))     
          output=np.array(output,dtype=object)    
          np.save(ValidEmbedDir+embeddingName+str(i+1)+".npy",output)   
          model.train()
          del img,embeddings
          
      if (i+1)  % 100 == 0:
          
          ###evaluation
          model.eval()
          try:
              (t0, t1), _= next(TestGenerator)
          except StopIteration:
              test_set=preShuffleImageDataset(TestDir,innerShuffle=innerShuffle)
              TestLoader = torch.utils.data.DataLoader(
                  test_set,
                  batch_size=batch_size,
                  shuffle= False,
                  drop_last=True,
                  num_workers=16,
                  collate_fn= CollateFunction(is_train=True)
              )
              TestGenerator = iter(TestLoader)
              (t0, t1), _= next(TestGenerator)
          t0 = t0.to(device)
          t1 = t1.to(device)
          with torch.no_grad():  
              with torch.cuda.amp.autocast(enabled=AMP):
                  T_loss, T_loss_SS, T_loss_CLD,NR_1,NR_2 =model(t0,t1,Lambda)            
          print(f'[Test {i+1:8d}] '
              f'Loss = {T_loss.item():.5f} | '
              f'Loss_SS = {T_loss_SS.item():.5f} | '
              f'Loss_CLD = {T_loss_CLD.item():.5f} ')
          model.train()
          # the level of collapse is large if the standard deviation of the l2
          # normalized output is much smaller than 1 / sqrt(dim)
          # print intermediate results
          print(f'[Steps {i+1:8d}] '
              f'Loss = {avg_loss[0]:.5f} | '
              f'Loss_SS = {avg_loss[1]:.5f} | '
              f'Loss_CLD = {avg_loss[2]:.5f} | '
              f'Collapse Level Low Dim: {collapse_level[1]:.5f} | '
              f'Collapse Level High Dim: {collapse_level[0]:.5f}')
          f = open(log+"train.log",'a')
          f.write(f'[Steps {i+1:8d}] '
              f'Loss = {avg_loss[0]:.5f} | '
              f'Loss_SS = {avg_loss[1]:.5f} | '
              f'Loss_CLD = {avg_loss[2]:.5f} | '
              f'Collapse Level Low Dim: {collapse_level[1]:.5f} | '
              f'Collapse Level High Dim: {collapse_level[0]:.5f}\n')
          f.close()
          f = open(log+"test.log",'a')
          f.write(f'[Test {i+1:8d}] '
              f'Loss = {T_loss.item():.5f} | '
              f'Loss_SS = {T_loss_SS.item():.5f} | '
              f'Loss_CLD = {T_loss_CLD.item():.5f} \n')            
          f.close()
          time_end=time.time()
          print('time cost',time_end-time_start,'s')
          time_start=time.time()
          del t0,t1,T_loss, T_loss_SS, T_loss_CLD,NR_1,NR_2   
  
       
  