text: """
  
  Pre-train a CNN on the whole dataset for evaluation purpose
  
  """
  import os
  import argparse
  import shutil
  import timeit
  
  import torch
  import torchvision
  import torchvision.transforms as transforms
  import numpy as np
  import torch.nn as nn
  import torch.backends.cudnn as cudnn
  import random
  import matplotlib.pyplot as plt
  import matplotlib as mpl
  from torch import autograd
  from torchvision.utils import save_image
  import csv
  from tqdm import tqdm
  import gc
  import h5py
  
  from models import *
  from utils import IMGs_dataset
  
  
  #############################
  # Settings
  #############################
  
  parser = argparse.ArgumentParser(description='Pre-train CNNs')
  parser.add_argument('--root_path', type=str, default='')
  parser.add_argument('--data_path', type=str, default='')
  parser.add_argument('--num_workers', type=int, default=0)
  parser.add_argument('--CNN', type=str, default='ResNet34_class',
                      help='CNN for training; ResNetXX')
  parser.add_argument('--epochs', type=int, default=200, metavar='N',
                      help='number of epochs to train CNNs (default: 200)')
  parser.add_argument('--batch_size_train', type=int, default=128, metavar='N',
                      help='input batch size for training')
  parser.add_argument('--batch_size_valid', type=int, default=10, metavar='N',
                      help='input batch size for testing')
  parser.add_argument('--base_lr', type=float, default=0.01,
                      help='learning rate, default=0.1')
  parser.add_argument('--weight_dacay', type=float, default=1e-4,
                      help='Weigth decay, default=1e-4')
  parser.add_argument('--seed', type=int, default=2020, metavar='S',
                      help='random seed (default: 1)')
  parser.add_argument('--CVMode', action='store_true', default=False,
                      help='CV mode?')
  parser.add_argument('--valid_proport', type=float, default=0.1,
                      help='Proportion of validation samples')
  parser.add_argument('--img_size', type=int, default=128, metavar='N')
  parser.add_argument('--min_label', type=float, default=0.0)
  parser.add_argument('--max_label', type=float, default=90.0)
  args = parser.parse_args()
  
  
  wd = args.root_path
  os.chdir(wd)
  from models import *
  from utils import IMGs_dataset
  
  # cuda
  device = torch.device("cuda")
  ngpu = torch.cuda.device_count()  # number of gpus
  
  # random seed
  random.seed(args.seed)
  torch.manual_seed(args.seed)
  torch.backends.cudnn.deterministic = True
  np.random.seed(args.seed)
  
  # directories for checkpoint, images and log files
  save_models_folder = wd + '/output/eval_models/'
  os.makedirs(save_models_folder, exist_ok=True)
  
  
  # data loader
  data_filename = args.data_path + '/RC-49_' + str(args.img_size) + 'x' + str(args.img_size) + '.h5'
  hf = h5py.File(data_filename, 'r')
  angles = hf['labels'][:]
  angles = angles.astype(float)
  labels = hf['types'][:]
  images = hf['images'][:]
  hf.close()
  num_classes = len(set(labels))
  assert max(labels)==num_classes-1
  
  q1 = args.min_label
  q2 = args.max_label
  indx = np.where((angles>q1)*(angles<q2)==True)[0]
  angles = angles[indx]
  labels = labels[indx]
  images = images[indx]
  assert len(labels)==len(images)
  assert len(angles)==len(images)
  
  
  # define training (and validaiton) set
  if args.CVMode:
      for i in range(num_classes):
          indx_i = np.where(labels==i)[0] #i-th class
          np.random.shuffle(indx_i)
          num_imgs_all_i = len(indx_i)
          num_imgs_valid_i = int(num_imgs_all_i*args.valid_proport)
          if i == 0:
              indx_valid = indx_i[0:num_imgs_valid_i]
              indx_train = indx_i[num_imgs_valid_i:]
          else:
              indx_valid = np.concatenate((indx_valid, indx_i[0:num_imgs_valid_i]))
              indx_train = np.concatenate((indx_train, indx_i[num_imgs_valid_i:]))
      #end for i
      trainset = IMGs_dataset(images[indx_train], labels[indx_train], normalize=True)
      trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size_train, shuffle=True, num_workers=args.num_workers)
      validset = IMGs_dataset(images[indx_valid], labels[indx_valid], normalize=True)
      validloader = torch.utils.data.DataLoader(validset, batch_size=args.batch_size_valid, shuffle=False, num_workers=args.num_workers)
  else:
      trainset = IMGs_dataset(images, labels, normalize=True)
      trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size_train, shuffle=True, num_workers=args.num_workers)
  
  
  
  ###########################################################################################################
  # Necessary functions
  ###########################################################################################################
  
  #initialize CNNs
  def net_initialization(Pretrained_CNN_Name, ngpu = 1, num_classes=num_classes):
      if Pretrained_CNN_Name == "ResNet18_class":
          net = ResNet18_class_eval(num_classes=num_classes, ngpu = ngpu)
      elif Pretrained_CNN_Name == "ResNet34_class":
          net = ResNet34_class_eval(num_classes=num_classes, ngpu = ngpu)
      elif Pretrained_CNN_Name == "ResNet50_class":
          net = ResNet50_class_eval(num_classes=num_classes, ngpu = ngpu)
      elif Pretrained_CNN_Name == "ResNet101_class":
          net = ResNet101_class_eval(num_classes=num_classes, ngpu = ngpu)
  
      net_name = 'PreCNNForEvalGANs_' + Pretrained_CNN_Name #get the net's name
      net = net.to(device)
  
      return net, net_name
  
  #adjust CNN learning rate
  def adjust_learning_rate(optimizer, epoch, BASE_LR_CNN):
      lr = BASE_LR_CNN
      # if epoch >= 35:
      #     lr /= 10
      # if epoch >= 70:
      #     lr /= 10
      if epoch >= 50:
          lr /= 10
      if epoch >= 120:
          lr /= 10
      for param_group in optimizer.param_groups:
          param_group['lr'] = lr
  
  
  def train_CNN():
  
      start_tmp = timeit.default_timer()
      for epoch in range(args.epochs):
          net.train()
          train_loss = 0
          adjust_learning_rate(optimizer, epoch, args.base_lr)
          for batch_idx, (batch_train_images, batch_train_labels) in enumerate(trainloader):
  
              # batch_train_images = nn.functional.interpolate(batch_train_images, size = (299,299), scale_factor=None, mode='bilinear', align_corners=False)
  
              batch_train_images = batch_train_images.type(torch.float).cuda()
              batch_train_labels = batch_train_labels.type(torch.long).cuda()
  
              #Forward pass
              outputs,_ = net(batch_train_images)
              loss = criterion(outputs, batch_train_labels)
  
              #backward pass
              optimizer.zero_grad()
              loss.backward()
              optimizer.step()
  
              train_loss += loss.cpu().item()
          #end for batch_idx
          train_loss = train_loss / len(trainloader)
  
          if args.CVMode:
              valid_acc = valid_CNN(verbose=False)
              print('CNN: [epoch %d/%d] train_loss:%f valid_acc:%f Time: %.4f' % (epoch+1, args.epochs, train_loss, valid_acc, timeit.default_timer()-start_tmp))
          else:
              print('CNN: [epoch %d/%d] train_loss:%f Time: %.4f' % (epoch+1, args.epochs, train_loss, timeit.default_timer()-start_tmp))
      #end for epoch
  
      return net, optimizer
  
  
  if args.CVMode:
      def valid_CNN(verbose=True):
          net.eval()  # eval mode (batchnorm uses moving mean/variance instead of mini-batch mean/variance)
          with torch.no_grad():
              correct = 0
              total = 0
              for batch_idx, (images, labels) in enumerate(validloader):
                  images = images.type(torch.float).cuda()
                  labels = labels.type(torch.long).cuda()
                  outputs,_ = net(images)
                  _, predicted = torch.max(outputs.data, 1)
                  total += labels.size(0)
                  correct += (predicted == labels).sum().item()
              if verbose:
                  print('Valid Accuracy of the model on the validation set: {} %'.format(100.0 * correct / total))
          return 100.0 * correct / total
  
  
  
  ###########################################################################################################
  # Training and Testing
  ###########################################################################################################
  # model initialization
  net, net_name = net_initialization(args.CNN, ngpu = ngpu, num_classes = num_classes)
  criterion = nn.CrossEntropyLoss()
  optimizer = torch.optim.SGD(net.parameters(), lr = args.base_lr, momentum= 0.9, weight_decay=args.weight_dacay)
  
  filename_ckpt = save_models_folder + '/ckpt_{}_epoch_{}_seed_{}_classify_{}_chair_types_CVMode_{}.pth'.format(net_name, args.epochs, args.seed, num_classes, args.CVMode)
  
  # training
  if not os.path.isfile(filename_ckpt):
      # TRAIN CNN
      print("\n Begin training CNN: ")
      start = timeit.default_timer()
      net, optimizer = train_CNN()
      stop = timeit.default_timer()
      print("Time elapses: {}s".format(stop - start))
      # save model
      torch.save({
      'net_state_dict': net.state_dict(),
      }, filename_ckpt)
  else:
      print("\n Ckpt already exists")
      print("\n Loading...")
      checkpoint = torch.load(filename_ckpt)
      net.load_state_dict(checkpoint['net_state_dict'])
  torch.cuda.empty_cache()#release GPU mem which is  not references
  
  if args.CVMode:
      #validation
      _ = valid_CNN(True)
      torch.cuda.empty_cache()
  