import os, sys, time, glob, random, argparse import numpy as np from copy import deepcopy import torch import torch.nn as nn import torch.nn.functional as F import torchvision.datasets as dset import torch.backends.cudnn as cudnn import torchvision.transforms as transforms from pathlib import Path lib_dir = (Path(__file__).parent / '..' / 'lib').resolve() if str(lib_dir) not in sys.path: sys.path.insert(0, str(lib_dir)) from utils import AverageMeter, time_string, convert_secs2time from utils import print_log, obtain_accuracy from utils import Cutout, count_parameters_in_MB from nas import model_types as models from train_utils import main_procedure from train_utils_imagenet import main_procedure_imagenet from scheduler import load_config parser = argparse.ArgumentParser("Train-CNN") parser.add_argument('--data_path', type=str, help='Path to dataset') parser.add_argument('--dataset', type=str, choices=['imagenet', 'cifar10', 'cifar100'], help='Choose between Cifar10/100 and ImageNet.') parser.add_argument('--arch', type=str, choices=models.keys(), help='the searched model.') # parser.add_argument('--grad_clip', type=float, help='gradient clipping') parser.add_argument('--model_config', type=str , help='the model configuration') parser.add_argument('--init_channels', type=int , help='the initial number of channels') parser.add_argument('--layers', type=int , help='the number of layers.') # log parser.add_argument('--workers', type=int, default=2, help='number of data loading workers (default: 2)') parser.add_argument('--save_path', type=str, help='Folder to save checkpoints and log.') parser.add_argument('--print_freq', type=int, help='print frequency (default: 200)') parser.add_argument('--manualSeed', type=int, help='manual seed') args = parser.parse_args() if 'CUDA_VISIBLE_DEVICES' not in os.environ: print('Can not find CUDA_VISIBLE_DEVICES in os.environ') else : print('Find CUDA_VISIBLE_DEVICES={:}'.format(os.environ['CUDA_VISIBLE_DEVICES'])) assert torch.cuda.is_available(), 'torch.cuda is not available' if args.manualSeed is None or args.manualSeed < 0: args.manualSeed = random.randint(1, 10000) random.seed(args.manualSeed) cudnn.benchmark = True cudnn.enabled = True torch.manual_seed(args.manualSeed) torch.cuda.manual_seed_all(args.manualSeed) def main(): # Init logger #args.save_path = os.path.join(args.save_path, 'seed-{:}'.format(args.manualSeed)) if not os.path.isdir(args.save_path): os.makedirs(args.save_path) log = open(os.path.join(args.save_path, 'seed-{:}-log.txt'.format(args.manualSeed)), 'w') print_log('Save Path : {:}'.format(args.save_path), log) state = {k: v for k, v in args._get_kwargs()} print_log(state, log) print_log("Random Seed : {:}".format(args.manualSeed), log) print_log("Python version : {:}".format(sys.version.replace('\n', ' ')), log) print_log("Torch version : {:}".format(torch.__version__), log) print_log("CUDA version : {:}".format(torch.version.cuda), log) print_log("cuDNN version : {:}".format(cudnn.version()), log) print_log("Num of GPUs : {:}".format(torch.cuda.device_count()), log) args.dataset = args.dataset.lower() config = load_config(args.model_config) genotype = models[args.arch] print_log('configuration : {:}'.format(config), log) print_log('genotype : {:}'.format(genotype), log) # clear GPU cache torch.cuda.empty_cache() if args.dataset == 'imagenet': main_procedure_imagenet(config, args.data_path, args, genotype, args.init_channels, args.layers, None, log) else: main_procedure(config, args.dataset, args.data_path, args, genotype, args.init_channels, args.layers, None, log) log.close() if __name__ == '__main__': main()