code: import os
  import sys
  import time
  import re
  import argparse
  import numpy as np
  from tqdm import tqdm
  from tensorboardX import SummaryWriter
  import matplotlib.pyplot as plt
  
  plt.switch_backend('agg')
  
  sys.path.append('../Utils')
  from CPCTrans.dataset_3d import *
  from CPCTrans.model_3d import *
  from Backbone.resnet import neq_load_customized
  from Utils.augmentation import *
  from Utils.utils import AverageMeter, save_checkpoint, denorm, calc_topk_accuracy
  
  import torch
  import torch.optim as optim
  from torch.utils import data
  from torchvision import datasets, models, transforms
  import torchvision.utils as vutils
  
  torch.backends.cudnn.benchmark = True
  
  parser = argparse.ArgumentParser()
  parser.add_argument('--net', default='resnet18', type=str)
  parser.add_argument('--model', default='cpc-trans', type=str)
  parser.add_argument('--dataset', default='ucf101', type=str)
  parser.add_argument('--num_seq', default=8, type=int, help='number of video blocks')
  parser.add_argument('--pred_step', default=3, type=int)
  parser.add_argument('--ds', default=3, type=int, help='frame downsampling rate')
  parser.add_argument('--batch_size', default=16, type=int)
  parser.add_argument('--lr', default=1e-3, type=float, help='learning rate')
  parser.add_argument('--wd', default=1e-5, type=float, help='weight decay')
  parser.add_argument('--resume', default='', type=str, help='path of model to resume')
  parser.add_argument('--pretrain', default='', type=str, help='path of pretrained model')
  parser.add_argument('--epochs', default=300, type=int, help='number of total epochs to run')
  parser.add_argument('--start_epoch', default=1, type=int, help='manual epoch number (useful on restarts)')
  parser.add_argument('--gpu', default='0', type=str)
  parser.add_argument('--print_freq', default=200, type=int, help='frequency of printing output during training')
  parser.add_argument('--reset_lr', action='store_true', help='Reset learning rate when resume training?')
  parser.add_argument('--prefix', default='tmp', type=str, help='prefix of checkpoint filename')
  parser.add_argument('--train_what', default='all', type=str)
  parser.add_argument('--img_dim', default=128, type=int)
  
  
  def main():
      torch.manual_seed(0)
      np.random.seed(0)
      global args;
  
      args = parser.parse_args()
      os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
      global cuda;
      cuda = torch.device('cuda')
  
      ### CPC with TransformerEncoder model ###
      if args.model == 'cpc-trans':
          model = CPC_Trans(sample_size=args.img_dim,
                            num_seq=args.num_seq,
                            network=args.net,
                            pred_step=args.pred_step)
      else:
          raise ValueError('wrong model!')
  
      model = model.to(cuda)
      global criterion;
      criterion = nn.CrossEntropyLoss()
  
      params = model.parameters()
      optimizer = optim.Adam(params, lr=args.lr, weight_decay=args.wd)
      scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=100, gamma=0.5)
      args.old_lr = None
  
      best_acc = 0
      global iteration;
      iteration = 0
  
      ### restart training ###
      if args.resume:
          if os.path.isfile(args.resume):
              args.old_lr = float(re.search('_lr(.+?)_', args.resume).group(1))
              print("=> loading resumed checkpoint '{}'".format(args.resume))
              checkpoint = torch.load(args.resume, map_location=torch.device('cpu'))
              args.start_epoch = checkpoint['epoch']
              iteration = checkpoint['iteration']
              best_acc = checkpoint['best_acc']
              model.load_state_dict(checkpoint['state_dict'])
              if not args.reset_lr:  # if didn't reset lr, load old optimizer
                  optimizer.load_state_dict(checkpoint['optimizer'])
              else:
                  print('==== Change lr from %f to %f ====' % (args.old_lr, args.lr))
              print("=> loaded resumed checkpoint '{}' (epoch {}) with best_acc {}".format(args.resume, checkpoint['epoch'], best_acc))
          else:
              print("[Warning] no checkpoint found at '{}'".format(args.resume))
  
      if args.pretrain:
          if os.path.isfile(args.pretrain):
              print("=> loading pretrained checkpoint '{}'".format(args.pretrain))
              checkpoint = torch.load(args.pretrain, map_location=torch.device('cpu'))
              model = neq_load_customized(model, checkpoint['state_dict'])
              print("=> loaded pretrained checkpoint '{}' (epoch {})"
                    .format(args.pretrain, checkpoint['epoch']))
          else:
              print("=> no checkpoint found at '{}'".format(args.pretrain))
  
      ### load data ###
      if args.dataset == 'ucf101':  # designed for ucf101, short size=256, rand crop to 224x224 then scale to 128x128
          transform = transforms.Compose([
              RandomHorizontalFlip(consistent=True),
              RandomCrop(size=224, consistent=True),
              Scale(size=(args.img_dim, args.img_dim)),
              RandomGray(consistent=False, p=0.5),
              ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.25, p=1.0),
              ToTensor(),
              Normalize()
          ])
      elif args.dataset == 'k400':  # designed for kinetics400, short size=150, rand crop to 128x128
          transform = transforms.Compose([
              RandomSizedCrop(size=args.img_dim, consistent=True, p=1.0),
              RandomHorizontalFlip(consistent=True),
              RandomGray(consistent=False, p=0.5),
              ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.25, p=1.0),
              ToTensor(),
              Normalize()
          ])
  
      train_loader = get_data(transform, 'train')
      val_loader = get_data(transform, 'val')
  
      # setup tools
      global de_normalize;
      de_normalize = denorm()
      global img_path;
      img_path, model_path = set_path(args)
      global writer_train
      try:  # old version
          writer_val = SummaryWriter(log_dir=os.path.join(img_path, 'val'))
          writer_train = SummaryWriter(log_dir=os.path.join(img_path, 'train'))
      except:  # v1.7
          writer_val = SummaryWriter(logdir=os.path.join(img_path, 'val'))
          writer_train = SummaryWriter(logdir=os.path.join(img_path, 'train'))
  
      print('-- start main loop --')
  
      ### main loop ###
      for epoch in range(args.start_epoch, args.epochs):
          train_loss, train_acc, train_accuracy_list = train(train_loader, model, optimizer, epoch)
          val_loss, val_acc, val_accuracy_list = validate(val_loader, model, epoch)
          scheduler.step()
          print('\t Epoch: ', epoch, 'with lr: ', scheduler.get_last_lr())
  
          # save curve
          writer_train.add_scalar('global/loss', train_loss, epoch)
          writer_train.add_scalar('global/accuracy', train_acc, epoch)
          writer_val.add_scalar('global/loss', val_loss, epoch)
          writer_val.add_scalar('global/accuracy', val_acc, epoch)
          writer_train.add_scalar('accuracy/top1', train_accuracy_list[0], epoch)
          writer_train.add_scalar('accuracy/top3', train_accuracy_list[1], epoch)
          writer_train.add_scalar('accuracy/top5', train_accuracy_list[2], epoch)
          writer_val.add_scalar('accuracy/top1', val_accuracy_list[0], epoch)
          writer_val.add_scalar('accuracy/top3', val_accuracy_list[1], epoch)
          writer_val.add_scalar('accuracy/top5', val_accuracy_list[2], epoch)
  
          # save check_point
          is_best = val_acc > best_acc;
          best_acc = max(val_acc, best_acc)
          save_checkpoint({'epoch': epoch + 1,
                           'net': args.net,
                           'state_dict': model.state_dict(),
                           'best_acc': best_acc,
                           'optimizer': optimizer.state_dict(),
                           'iteration': iteration},
                          is_best, filename=os.path.join(model_path, 'epoch%s.pth.tar' % str(epoch + 1)), keep_all=False)
  
      print('Training from ep %d to ep %d finished' % (args.start_epoch, args.epochs))
  
  
  
  def process_output(mask):
      '''task mask as input, compute the target for contrastive loss'''
      (B, NP, SQ, B2, NS, _) = mask.size()  # [B, P, SQ, B, N, SQ]
      target = mask == 1
      target = target * 1
      target.requires_grad = False
      return target, (B, B2, NS, NP, SQ)
  
  
  def train(data_loader, model, optimizer, epoch):
      losses = AverageMeter()
      accuracy = AverageMeter()
      accuracy_list = [AverageMeter(), AverageMeter(), AverageMeter()]
      model.train()
      global iteration
  
      for idx, input_seq in enumerate(data_loader):
          tic = time.time()
          input_seq = input_seq.to(cuda)
          B = input_seq.size(0)
          [score_, mask_] = model(input_seq)
          # visualize
          if (iteration == 0) or (iteration == args.print_freq):
              if B > 2: input_seq = input_seq[0:2, :]
              writer_train.add_image('input_seq',
                                     de_normalize(vutils.make_grid(
                                         input_seq.view(-1, 3, args.img_dim, args.img_dim),
                                         nrow=args.num_seq)),
                                     iteration)
          del input_seq
  
          if idx == 0: target_, (_, B2, NS, NP, SQ) = process_output(mask_)
  
          score_flattened = score_.view(B * NP * SQ, B2 * NS * SQ)
          target_flattened = target_.contiguous().view(B * NP * SQ, B2 * NS * SQ)
          target_flattened = target_flattened.argmax(dim=1)
  
          loss = criterion(score_flattened, target_flattened)
          top1, top3, top5 = calc_topk_accuracy(score_flattened, target_flattened, (1, 3, 5))
  
          accuracy_list[0].update(top1.item(), B)
          accuracy_list[1].update(top3.item(), B)
          accuracy_list[2].update(top5.item(), B)
  
          losses.update(loss.item(), B)
          accuracy.update(top1.item(), B)
  
          del score_
  
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
  
          del loss
  
          if idx % args.print_freq == 0:
              print('Epoch: [{0}][{1}/{2}]\t'
                    'Loss {loss.val:.6f} ({loss.local_avg:.4f})\t'
                    'Acc: top1 {3:.4f}; top3 {4:.4f}; top5 {5:.4f} T:{6:.2f}\t'.format(
                  epoch, idx, len(data_loader), top1, top3, top5, time.time() - tic, loss=losses))
  
              writer_train.add_scalar('local/loss', losses.val, iteration)
              writer_train.add_scalar('local/accuracy', accuracy.val, iteration)
  
              iteration += 1
  
      return losses.local_avg, accuracy.local_avg, [i.local_avg for i in accuracy_list]
  
  
  def validate(data_loader, model, epoch):
      losses = AverageMeter()
      accuracy = AverageMeter()
      accuracy_list = [AverageMeter(), AverageMeter(), AverageMeter()]
      model.eval()
  
      with torch.no_grad():
          for idx, input_seq in tqdm(enumerate(data_loader), total=len(data_loader)):
              input_seq = input_seq.to(cuda)
              B = input_seq.size(0)
              [score_, mask_] = model(input_seq)
              del input_seq
  
              if idx == 0: target_, (_, B2, NS, NP, SQ) = process_output(mask_)
  
              score_flattened = score_.view(B * NP * SQ, B2 * NS * SQ)
              target_flattened = target_.contiguous().view(B * NP * SQ, B2 * NS * SQ)
              target_flattened = target_flattened.argmax(dim=1)
  
              loss = criterion(score_flattened, target_flattened)
              top1, top3, top5 = calc_topk_accuracy(score_flattened, target_flattened, (1, 3, 5))
  
              losses.update(loss.item(), B)
              accuracy.update(top1.item(), B)
  
              accuracy_list[0].update(top1.item(), B)
              accuracy_list[1].update(top3.item(), B)
              accuracy_list[2].update(top5.item(), B)
  
      print('[{0}/{1}] Loss {loss.local_avg:.4f}\t'
            'Acc: top1 {2:.4f}; top3 {3:.4f}; top5 {4:.4f} \t'.format(
          epoch, args.epochs, *[i.avg for i in accuracy_list], loss=losses))
      return losses.local_avg, accuracy.local_avg, [i.local_avg for i in accuracy_list]
  
  
  def get_data(transform, mode='train'):
      print('Loading data for "%s" ...' % mode)
      if args.dataset == 'k400':
          pass
      elif args.dataset == 'ucf101':
          dataset = UCF101_3d(mode=mode,
                              transform=transform,
                              num_seq=args.num_seq,
                              downsample=args.ds,
                              which_split=3)
      else:
          raise ValueError('dataset not supported')
  
      sampler = data.RandomSampler(dataset)
  
      if mode == 'train':
          data_loader = data.DataLoader(dataset,
                                        batch_size=args.batch_size,
                                        sampler=sampler,
                                        shuffle=False,
                                        num_workers=2,
                                        pin_memory=True,
                                        drop_last=True)
      elif mode == 'val':
          data_loader = data.DataLoader(dataset,
                                        batch_size=args.batch_size,
                                        sampler=sampler,
                                        shuffle=False,
                                        num_workers=2,
                                        pin_memory=True,
                                        drop_last=True)
      print('"%s" dataset size: %d' % (mode, len(dataset)))
      return data_loader
  
  
  def set_path(args):
      if args.resume:
          exp_path = os.path.dirname(os.path.dirname(args.resume))
      else:
          exp_path = 'log_{args.prefix}/{args.dataset}-{args.img_dim}_{0}_{args.model}_\
  bs{args.batch_size}_lr{1}_seq{args.num_seq}_pred{args.pred_step}_ds{args.ds}_\
  train-{args.train_what}{2}'.format(
              'r%s' % args.net[6::], \
              args.old_lr if args.old_lr is not None else args.lr, \
              '_pt=%s' % args.pretrain.replace('/', '-') if args.pretrain else '', \
              args=args)
      img_path = os.path.join(exp_path, 'img')
      model_path = os.path.join(exp_path, 'model')
      if not os.path.exists(img_path): os.makedirs(img_path)
      if not os.path.exists(model_path): os.makedirs(model_path)
      return img_path, model_path
  
  
  if __name__ == '__main__':
      main()
  
apis: torch.utils.data.DataLoader matplotlib.pyplot.switch_backend torch.manual_seed numpy.random.seed torch.no_grad torch.optim.Adam torch.optim.lr_scheduler.StepLR torch.utils.data.RandomSampler torch.device