import os import sys sys.path.insert(0, '../../') import time import glob import random import numpy as np import torch import shutil import nasbench201.utils as ig_utils import logging import argparse import torch.nn as nn import torch.utils import torchvision.datasets as dset import torch.backends.cudnn as cudnn import torchvision.transforms as transforms import json import copy from sota.cnn.model_search import Network as DartsNetwork from sota.cnn.model_search_darts_proj import DartsNetworkProj from sota.cnn.model_search_imagenet_proj import ImageNetNetworkProj # from optimizers.darts.architect import Architect as DartsArchitect from nasbench201.architect_ig import Architect from sota.cnn.spaces import spaces_dict from foresight.pruners import * from torch.utils.tensorboard import SummaryWriter from sota.cnn.init_projection import pt_project from hdf5 import H5Dataset torch.set_printoptions(precision=4, sci_mode=False) parser = argparse.ArgumentParser("sota") parser.add_argument('--data', type=str, default='../../data',help='location of the data corpus') parser.add_argument('--dataset', type=str, default='cifar10', help='choose dataset') parser.add_argument('--batch_size', type=int, default=64, help='batch size') parser.add_argument('--train_portion', type=float, default=0.5, help='portion of training data') parser.add_argument('--cutout', action='store_true', default=False, help='use cutout') parser.add_argument('--cutout_length', type=int, default=16, help='cutout length') parser.add_argument('--cutout_prob', type=float, default=1.0, help='cutout probability') parser.add_argument('--seed', type=int, default=666, help='random seed') #model opt related config parser.add_argument('--learning_rate', type=float, default=0.025, help='init learning rate') parser.add_argument('--learning_rate_min', type=float, default=0.001, help='min learning rate') parser.add_argument('--momentum', type=float, default=0.9, help='momentum') parser.add_argument('--nesterov', action='store_true', default=True, help='using nestrov momentum for SGD') parser.add_argument('--weight_decay', type=float, default=3e-4, help='weight decay') parser.add_argument('--grad_clip', type=float, default=5, help='gradient clipping') #system config parser.add_argument('--gpu', type=str, default='0', help='gpu device id') parser.add_argument('--save', type=str, default='exp', help='experiment name') parser.add_argument('--save_path', type=str, default='../../experiments/sota', help='experiment name') #search sapce config parser.add_argument('--init_channels', type=int, default=16, help='num of init channels') parser.add_argument('--layers', type=int, default=8, help='total number of layers') parser.add_argument('--search_space', type=str, default='s5', help='searching space to choose from') parser.add_argument('--pool_size', type=int, default=10, help='number of model to proposed') ## projection parser.add_argument('--edge_decision', type=str, default='random', choices=['random','reverse', 'order', 'global_op_greedy', 'global_op_once', 'global_edge_greedy', 'global_edge_once', 'sample'], help='used for both proj_op and proj_edge') parser.add_argument('--proj_crit_normal', type=str, default='meco', choices=['loss', 'acc', 'jacob', 'comb', 'synflow', 'snip', 'fisher', 'var', 'cor', 'norm', 'grad_norm', 'grasp', 'jacob_cov', 'meco', 'zico']) parser.add_argument('--proj_crit_reduce', type=str, default='meco', choices=['loss', 'acc', 'jacob', 'comb', 'synflow', 'snip', 'fisher', 'var', 'cor', 'norm', 'grad_norm', 'grasp', 'jacob_cov', 'meco', 'zico']) parser.add_argument('--proj_crit_edge', type=str, default='meco', choices=['loss', 'acc', 'jacob', 'comb', 'synflow', 'snip', 'fisher', 'var', 'cor', 'norm', 'grad_norm', 'grasp', 'jacob_cov', 'meco', 'zico']) parser.add_argument('--proj_mode_edge', type=str, default='reg', choices=['reg'], help='edge projection evaluation mode, reg: one edge at a time') args = parser.parse_args() #### args augment expid = args.save args.save = '{}/{}-search-{}-{}-{}-{}-{}'.format(args.save_path, args.dataset, args.save, args.search_space, args.seed, args.pool_size, args.proj_crit_normal) if not args.edge_decision == 'random': args.save += '-' + args.edge_decision scripts_to_save = glob.glob('*.py') + glob.glob('../../nasbench201/architect*.py') + glob.glob('../../optimizers/darts/architect.py') if os.path.exists(args.save): if input("WARNING: {} exists, override?[y/n]".format(args.save)) == 'y': print('proceed to override saving directory') shutil.rmtree(args.save) else: exit(0) ig_utils.create_exp_dir(args.save, scripts_to_save=scripts_to_save) #### logging log_format = '%(asctime)s %(message)s' logging.basicConfig(stream=sys.stdout, level=logging.INFO, format=log_format, datefmt='%m/%d %I:%M:%S %p') log_file = 'log.txt' log_path = os.path.join(args.save, log_file) logging.info('======> log filename: %s', log_file) if os.path.exists(log_path): if input("WARNING: {} exists, override?[y/n]".format(log_file)) == 'y': print('proceed to override log file directory') else: exit(0) fh = logging.FileHandler(log_path, mode='w') fh.setFormatter(logging.Formatter(log_format)) logging.getLogger().addHandler(fh) writer = SummaryWriter(args.save + '/runs') if args.dataset == 'cifar100': n_classes = 100 elif args.dataset == 'imagenet': n_classes = 1000 else: n_classes = 10 def main(): torch.set_num_threads(3) if not torch.cuda.is_available(): logging.info('no gpu device available') sys.exit(1) np.random.seed(args.seed) gpu = ig_utils.pick_gpu_lowest_memory() if args.gpu == 'auto' else int(args.gpu) torch.cuda.set_device(gpu) cudnn.benchmark = True torch.manual_seed(args.seed) cudnn.enabled = True torch.cuda.manual_seed(args.seed) logging.info('gpu device = %d' % gpu) logging.info("args = %s", args) #### model criterion = nn.CrossEntropyLoss() criterion = criterion.cuda() ## darts if args.dataset == 'imagenet': model = ImageNetNetworkProj(args.init_channels, n_classes, args.layers, criterion, spaces_dict[args.search_space], args) else: model = DartsNetworkProj(args.init_channels, n_classes, args.layers, criterion, spaces_dict[args.search_space], args) model = model.cuda() logging.info("param size = %fMB", ig_utils.count_parameters_in_MB(model)) #### data if args.dataset == 'imagenet': normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_transform = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ColorJitter( brightness=0.4, contrast=0.4, saturation=0.4, hue=0.2), transforms.ToTensor(), normalize, ]) #for test #from nasbench201.DownsampledImageNet import ImageNet16 # train_data = dset.CIFAR10(root=args.data, train=True, download=True, transform=train_transform) # n_classes = 10 train_data = H5Dataset(os.path.join(args.data, 'imagenet-train-256.h5'), transform=train_transform) #valid_data = H5Dataset(os.path.join(args.data, 'imagenet-val-256.h5'), transform=test_transform) train_queue = torch.utils.data.DataLoader( train_data, batch_size=args.batch_size, shuffle=True, pin_memory=True, num_workers=4) else: if args.dataset == 'cifar10': train_transform, valid_transform = ig_utils._data_transforms_cifar10(args) train_data = dset.CIFAR10(root=args.data, train=True, download=True, transform=train_transform) valid_data = dset.CIFAR10(root=args.data, train=False, download=True, transform=valid_transform) elif args.dataset == 'cifar100': train_transform, valid_transform = ig_utils._data_transforms_cifar100(args) train_data = dset.CIFAR100(root=args.data, train=True, download=True, transform=train_transform) valid_data = dset.CIFAR100(root=args.data, train=False, download=True, transform=valid_transform) elif args.dataset == 'svhn': train_transform, valid_transform = ig_utils._data_transforms_svhn(args) train_data = dset.SVHN(root=args.data, split='train', download=True, transform=train_transform) valid_data = dset.SVHN(root=args.data, split='test', download=True, transform=valid_transform) num_train = len(train_data) indices = list(range(num_train)) split = int(np.floor(args.train_portion * num_train)) train_queue = torch.utils.data.DataLoader( train_data, batch_size=args.batch_size, sampler=torch.utils.data.sampler.SubsetRandomSampler(indices[:split]), pin_memory=True) valid_queue = torch.utils.data.DataLoader( train_data, batch_size=args.batch_size, sampler=torch.utils.data.sampler.SubsetRandomSampler(indices[split:num_train]), pin_memory=True) # for x, y in train_queue: # from torchvision import transforms # unloader = transforms.ToPILImage() # image = x.cpu().clone() # clone the tensor # image = image.squeeze(0) # remove the fake batch dimension # image = unloader(image) # image.save('example.jpg') # print(x.size()) # exit() #### projection networks_pool={} networks_pool['search_space'] = args.search_space networks_pool['dataset'] = args.dataset networks_pool['networks'] = [] for i in range(args.pool_size): network_info={} logging.info('{} MODEL HAS SEARCHED'.format(i+1)) pt_project(train_queue, model, args) ## logging num_params = ig_utils.count_parameters_in_Compact(model) genotype = model.genotype() json_data = {} json_data['normal'] = genotype.normal json_data['normal_concat'] = [x for x in genotype.normal_concat] json_data['reduce'] = genotype.reduce json_data['reduce_concat'] = [x for x in genotype.reduce_concat] json_string = json.dumps(json_data) logging.info(json_string) network_info['id'] = str(i) network_info['genotype'] = json_string networks_pool['networks'].append(network_info) model.reset_arch_parameters() with open(os.path.join(args.save,'networks_pool.json'), 'w') as save_file: json.dump(networks_pool, save_file) if __name__ == '__main__': main()