text: import os
  import sys
  import shutil
  import numpy as np
  import time, datetime
  import torch
  import random
  import logging
  import argparse
  from tqdm import tqdm
  import torch.nn as nn
  import torch.utils
  import torch.backends.cudnn as cudnn
  import torch.distributed as dist
  import torch.utils.data.distributed
  
  sys.path.append("../../")
  from utils.utils import *
  from utils import KD_loss
  from torchvision import datasets, transforms
  from torch.autograd import Variable
  from birealnet import birealnet18
  import torchvision.models as models
  
  parser = argparse.ArgumentParser("birealnet18")
  parser.add_argument('--batch_size', type=int, default=512, help='batch size')
  parser.add_argument('--epochs', type=int, default=256, help='num of training epochs')
  parser.add_argument('--learning_rate', type=float, default=0.001, help='init learning rate')
  parser.add_argument('--momentum', type=float, default=0.9, help='momentum')
  parser.add_argument('--weight_decay', type=float, default=0, help='weight decay')
  parser.add_argument('--save', type=str, default='./models', help='path for saving trained models')
  parser.add_argument('--data', metavar='DIR', help='path to dataset')
  parser.add_argument('--label_smooth', type=float, default=0.1, help='label smoothing')
  parser.add_argument('--teacher', type=str, default='resnet34', help='path of ImageNet')
  parser.add_argument('-j', '--workers', default=40, type=int, metavar='N',
                      help='number of data loading workers (default: 4)')
  args = parser.parse_args()
  
  CLASSES = 1000
  
  if not os.path.exists('log'):
      os.mkdir('log')
  
  log_format = '%(asctime)s %(message)s'
  logging.basicConfig(stream=sys.stdout, level=logging.INFO,
      format=log_format, datefmt='%m/%d %I:%M:%S %p')
  fh = logging.FileHandler(os.path.join('log/log.txt'))
  fh.setFormatter(logging.Formatter(log_format))
  logging.getLogger().addHandler(fh)
  
  def main():
      if not torch.cuda.is_available():
          sys.exit(1)
      start_t = time.time()
  
      cudnn.benchmark = True
      cudnn.enabled=True
      logging.info("args = %s", args)
  
      # load model
      model_teacher = models.__dict__[args.teacher](pretrained=True)
      model_teacher = nn.DataParallel(model_teacher).cuda()
      for p in model_teacher.parameters():
          p.requires_grad = False
      model_teacher.eval()
  
      model_student = birealnet18()
      logging.info('student:')
      logging.info(model_student)
      model_student = nn.DataParallel(model_student).cuda()
  
      criterion = nn.CrossEntropyLoss()
      criterion = criterion.cuda()
      criterion_smooth = CrossEntropyLabelSmooth(CLASSES, args.label_smooth)
      criterion_smooth = criterion_smooth.cuda()
      criterion_kd = KD_loss.DistributionLoss()
  
      all_parameters = model_student.parameters()
      weight_parameters = []
      for pname, p in model_student.named_parameters():
          if p.ndimension() == 4 or 'conv' in pname:
              weight_parameters.append(p)
      weight_parameters_id = list(map(id, weight_parameters))
      other_parameters = list(filter(lambda p: id(p) not in weight_parameters_id, all_parameters))
  
      optimizer = torch.optim.Adam(
              [{'params' : other_parameters},
              {'params' : weight_parameters, 'weight_decay' : args.weight_decay}],
              lr=args.learning_rate,)
  
      scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda step : (1.0-step/args.epochs), last_epoch=-1)
      start_epoch = 0
      best_top1_acc= 0
  
      checkpoint_tar = os.path.join(args.save, 'checkpoint.pth.tar')
      logging.info('loading checkpoint {} ..........'.format(checkpoint_tar))
      checkpoint = torch.load(checkpoint_tar)
      start_epoch = checkpoint['epoch']
      best_top1_acc = checkpoint['best_top1_acc']
      model_student.load_state_dict(checkpoint['state_dict'], strict=False)
      logging.info("loaded checkpoint {} epoch = {}" .format(checkpoint_tar, checkpoint['epoch']))
      start_epoch = 255
  
      # adjust the learning rate according to the checkpoint
      for epoch in range(start_epoch):
          scheduler.step()
  
      # load training data
      traindir = os.path.join(args.data, 'train')
      valdir = os.path.join(args.data, 'val')
      normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                       std=[0.229, 0.224, 0.225])
  
      # data augmentation
      crop_scale = 0.08
      lighting_param = 0.1
      train_transforms = transforms.Compose([
          transforms.RandomResizedCrop(224, scale=(crop_scale, 1.0)),
          Lighting(lighting_param),
          transforms.RandomHorizontalFlip(),
          transforms.ToTensor(),
          normalize])
  
      train_dataset = datasets.ImageFolder(
          traindir,
          transform=train_transforms)
  
      train_loader = torch.utils.data.DataLoader(
          train_dataset, batch_size=args.batch_size, shuffle=True,
          num_workers=args.workers, pin_memory=True)
  
      # load validation data
      val_loader = torch.utils.data.DataLoader(
          datasets.ImageFolder(valdir, transforms.Compose([
              transforms.Resize(256),
              transforms.CenterCrop(224),
              transforms.ToTensor(),
              normalize,
          ])),
          batch_size=args.batch_size, shuffle=False,
          num_workers=args.workers, pin_memory=True)
  
      # train the model
      epoch = start_epoch
      while epoch < args.epochs:
          train_obj, train_top1_acc,  train_top5_acc = train(epoch,  train_loader, model_student, model_teacher, criterion_kd, optimizer, scheduler)
          valid_obj, valid_top1_acc, valid_top5_acc = validate(epoch, val_loader, model_student, criterion, args)
  
          is_best = False
          if valid_top1_acc > best_top1_acc:
              best_top1_acc = valid_top1_acc
              is_best = True
  
          save_checkpoint({
              'epoch': epoch,
              'state_dict': model_student.state_dict(),
              'best_top1_acc': best_top1_acc,
              'optimizer' : optimizer.state_dict(),
              }, is_best, args.save)
  
          epoch += 1
  
      training_time = (time.time() - start_t) / 3600
      print('total training time = {} hours'.format(training_time))
  
  
  def train(epoch, train_loader, model_student, model_teacher, criterion, optimizer, scheduler):
      batch_time = AverageMeter('Time', ':6.3f')
      data_time = AverageMeter('Data', ':6.3f')
      losses = AverageMeter('Loss', ':.4e')
      top1 = AverageMeter('Acc@1', ':6.2f')
      top5 = AverageMeter('Acc@5', ':6.2f')
  
      progress = ProgressMeter(
          len(train_loader),
          [batch_time, data_time, losses, top1, top5],
          prefix="Epoch: [{}]".format(epoch))
  
      model_student.train()
      model_teacher.eval()
      end = time.time()
      scheduler.step()
  
      for param_group in optimizer.param_groups:
          cur_lr = param_group['lr']
      print('learning_rate:', cur_lr)
  
      for i, (images, target) in enumerate(train_loader):
          data_time.update(time.time() - end)
          images = images.cuda()
          target = target.cuda()
  
          # compute outputy
          logits_student = model_student(images)
          logits_teacher = model_teacher(images)
          loss = criterion(logits_student, logits_teacher)
  
          # measure accuracy and record loss
          prec1, prec5 = accuracy(logits_student, target, topk=(1, 5))
          n = images.size(0)
          losses.update(loss.item(), n)   #accumulated loss
          top1.update(prec1.item(), n)
          top5.update(prec5.item(), n)
  
          # compute gradient and do SGD step
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
  
          # measure elapsed time
          batch_time.update(time.time() - end)
          end = time.time()
  
          progress.display(i)
  
      return losses.avg, top1.avg, top5.avg
  
  def validate(epoch, val_loader, model, criterion, args):
      batch_time = AverageMeter('Time', ':6.3f')
      losses = AverageMeter('Loss', ':.4e')
      top1 = AverageMeter('Acc@1', ':6.2f')
      top5 = AverageMeter('Acc@5', ':6.2f')
      # progress = ProgressMeter(
      #     len(val_loader),
      #     [batch_time, losses, top1, top5],
      #     prefix='Test: ')
  
      # switch to evaluation mode
      model.eval()
      with torch.no_grad():
          end = time.time()
          for i, (images, target) in tqdm(enumerate(val_loader), total=len(val_loader)):
              images = images.cuda()
              target = target.cuda()
  
              # compute output
              logits = model(images)
              loss = criterion(logits, target)
  
              # measure accuracy and record loss
              pred1, pred5 = accuracy(logits, target, topk=(1, 5))
              n = images.size(0)
              losses.update(loss.item(), n)
              top1.update(pred1[0], n)
              top5.update(pred5[0], n)
  
              # measure elapsed time
              batch_time.update(time.time() - end)
              end = time.time()
  
          #     progress.display(i)
  
          print(' * acc@1 {top1.avg:.3f} acc@5 {top5.avg:.3f}'
                .format(top1=top1, top5=top5))
  
      return losses.avg, top1.avg, top5.avg
  
  
  if __name__ == '__main__':
      main()
  