code: def main():
      global args
      args = parser.parse_args()
      print(args)
  
      os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_ids
      cudnn.benchmark = True
      cudnn.enabled = True
  
      random.seed(args.seed)
      np.random.seed(args.seed)
      torch.manual_seed(args.seed)
      torch.cuda.manual_seed(args.seed)
  
      dataset = (
          "lamp" if args.dataset == "LAMP-HQ" else args.dataset.lower()
      )
  
      # train loader of real data
      real_dataset_paired = Real_Dataset_112_paired(args)
      vis_pos, nir_pos = GenIdx(
          real_dataset_paired.vis_labels, real_dataset_paired.nir_labels
      )
      sampler = IdentitySampler(
          real_dataset_paired.vis_labels,
          real_dataset_paired.nir_labels,
          vis_pos,
          nir_pos,
          args.batch_size,
          4,
      )
  
      real_dataset_paired.visIndex = sampler.visIndex
      real_dataset_paired.nirIndex = sampler.nirIndex
  
      train_loader_real_paired = torch.utils.data.DataLoader(
          real_dataset_paired,
          batch_size=args.batch_size,
          sampler=sampler,
          num_workers=args.workers,
          pin_memory=True,
      )
  
      num_classes = real_dataset_paired.num_classes
  
      model = LightCNN_29Layers_cosface(num_classes=num_classes)
  
      model = torch.nn.DataParallel(model).cuda()
  
      # load pre trained model
      if args.pre_epoch:
          print("load pretrained model of epoch %d" % args.pre_epoch)
          load_model(
              model,
              "./model/lightCNN_epoch_%d.pth.tar" % args.pre_epoch,
          )
      else:
          print(
              "=> loading pretrained lightcnn '{}'".format(
                  args.weights_lightcnn
              )
          )
          load_model_train_lightcnn(model, args.weights_lightcnn)
  
      # criterion
      criterion = nn.CrossEntropyLoss().cuda()
      criterion_idmmd = IDMMD().cuda()
      margin_softmax = CosFace(s=64.0, m=0.4).cuda()
  
      """
      Stage I: model pretrained for last fc2 parameters
      """
      params_pretrain = []
      for name, value in model.named_parameters():
          if name == "module.weight":
              params_pretrain += [{"params": value, "lr": 1 * args.lr}]
  
      print("Stage I: trainable params ", len(params_pretrain))
      assert len(params_pretrain) > 0
  
      # optimizer
      optimizer_pretrain = torch.optim.SGD(
          params_pretrain,
          args.lr,
          momentum=args.momentum,
          weight_decay=args.weight_decay,
      )
  
      for epoch in range(1, 5):
          pre_train_pair(
              train_loader_real_paired,
              model,
              criterion,
              margin_softmax,
              optimizer_pretrain,
              epoch,
          )
          # save_checkpoint(model, epoch, args.save_name+"_pretrain", dataset)
  
      """
      Stage II: model finetune for full network
      """
      # optimizer
      optimizer = torch.optim.SGD(
          model.parameters(),
          args.lr,
          momentum=args.momentum,
          weight_decay=args.weight_decay,
      )
  
      start_epoch = args.pre_epoch + 1
      for epoch in range(start_epoch, args.epochs + 1):
          adjust_learning_rate(
              args.lr, args.step_size, optimizer, epoch
          )
          train(
              train_loader_real_paired,
              model,
              criterion,
              criterion_idmmd,
              margin_softmax,
              optimizer,
              epoch,
          )
          if epoch == args.epochs or epoch % 10 == 0:
              save_checkpoint(model, epoch, args.save_name, dataset)
  
file_path: insightface/recognition/idmmd/train.py
imports: import os
  import argparse
  import random
  import numpy
  import torch
  import torch.nn
  import torch.backends.cudnn
  from utils import *
  from network.lightcnn112 import LightCNN_29Layers_cosface
  from losses import IDMMD, CosFace
  from dataset_mix import Real_Dataset_112_paired, IdentitySampler, GenIdx
label_desc: Use image node_id 1 for calling a global function with example usage: main() without return types