def main(): # Select the search space to search in if args.search_space == '1': search_space = SearchSpace1() elif args.search_space == '2': search_space = SearchSpace2() elif args.search_space == '3': search_space = SearchSpace3() else: raise ValueError('Unknown search space') if not torch.cuda.is_available(): logging.info('no gpu device available') sys.exit(1) np.random.seed(args.seed) torch.cuda.set_device(args.gpu) cudnn.benchmark = True torch.manual_seed(args.seed) cudnn.enabled = True torch.cuda.manual_seed(args.seed) logging.info('gpu device = %d' % args.gpu) logging.info("args = %s", args) criterion = nn.CrossEntropyLoss() criterion = criterion.cuda() model = Network(args.init_channels, CIFAR_CLASSES, args.layers, criterion, output_weights=args.output_weights, steps=search_space.num_intermediate_nodes, search_space=search_space) model = model.cuda() logging.info("param size = %fMB", utils.count_parameters_in_MB(model)) optimizer = torch.optim.SGD(model.parameters(), args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) train_transform, valid_transform = utils._data_transforms_cifar10(args) train_data = dset.CIFAR10(root=args.data, train=True, download=True, transform=train_transform) train_data_non_augm = dset.CIFAR10(root=args.data, train=True, 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) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, float(args.epochs), eta_min=args.learning_rate_min) # Validation data with data augmentations augm_valid_subset = torch.utils.data.Subset(train_data, indices[split:num_train]) # Validation data with no data augmentations non_augm_valid_subset = torch.utils.data.Subset(train_data_non_augm, indices[split:num_train]) architect = Architect(model, args) for epoch in range(args.epochs): scheduler.step() lr = scheduler.get_lr()[0] logging.info('epoch %d lr %e', epoch, lr) # Save the one shot model architecture weights for later analysis filehandler = open( os.path.join(args.save, 'one_shot_architecture_{}.obj'.format(epoch)), 'wb') numpy_tensor_list = [] for tensor in model.arch_parameters(): numpy_tensor_list.append(tensor.detach().cpu().numpy()) pickle.dump(numpy_tensor_list, filehandler) # Save the entire one-shot-model filepath = os.path.join(args.save, 'one_shot_model_{}.obj'.format(epoch)) torch.save(model.state_dict(), filepath) logging.info('architecture', numpy_tensor_list) # training train_acc, train_obj = train( train_queue, [augm_valid_subset, non_augm_valid_subset], model, architect, criterion, optimizer, lr, epoch) logging.info('train_acc %f', train_acc) # validation valid_acc, valid_obj = infer(valid_queue, model, criterion) logging.info('valid_acc %f', valid_acc) utils.save(model, os.path.join(args.save, 'weights.pt'))
def main(): # Select the search space to search in if args.search_space == '1': search_space = SearchSpace1() elif args.search_space == '2': search_space = SearchSpace2() elif args.search_space == '3': search_space = SearchSpace3() else: raise ValueError('Unknown search space') if not torch.cuda.is_available(): logging.info('no gpu device available') sys.exit(1) np.random.seed(args.seed) torch.cuda.set_device(args.gpu) cudnn.benchmark = True torch.manual_seed(args.seed) cudnn.enabled = True torch.cuda.manual_seed(args.seed) logging.info('gpu device = %d' % args.gpu) logging.info("args = %s", args) criterion = nn.CrossEntropyLoss() criterion = criterion.cuda() model = Network(args.init_channels, CIFAR_CLASSES, args.layers, criterion, output_weights=args.output_weights, steps=search_space.num_intermediate_nodes, search_space=search_space) model = model.cuda() logging.info("param size = %fMB", utils.count_parameters_in_MB(model)) optimizer = torch.optim.SGD(model.parameters(), args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) train_transform, valid_transform = utils._data_transforms_cifar10(args) train_data = dset.CIFAR10(root=args.data, train=True, download=True, transform=train_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) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, float(args.epochs), eta_min=args.learning_rate_min) architect = Architect(model, args) for epoch in range(args.epochs): scheduler.step() lr = scheduler.get_lr()[0] # increase the cutout probability linearly throughout search train_transform.transforms[ -1].cutout_prob = args.cutout_prob * epoch / (args.epochs - 1) logging.info('epoch %d lr %e cutout_prob %e', epoch, lr, train_transform.transforms[-1].cutout_prob) # Save the one shot model architecture weights for later analysis arch_filename = os.path.join( args.save, 'one_shot_architecture_{}.obj'.format(epoch)) with open(arch_filename, 'wb') as filehandler: numpy_tensor_list = [] for tensor in model.arch_parameters(): numpy_tensor_list.append(tensor.detach().cpu().numpy()) pickle.dump(numpy_tensor_list, filehandler) # Save the entire one-shot-model filepath = os.path.join(args.save, 'one_shot_model_{}.obj'.format(epoch)) torch.save(model.state_dict(), filepath) logging.info('architecture', numpy_tensor_list) # training train_acc, train_obj = train(train_queue, valid_queue, model, architect, criterion, optimizer, lr, epoch) logging.info('train_acc %f', train_acc) # validation valid_acc, valid_obj = infer(valid_queue, model, criterion) logging.info('valid_acc %f', valid_acc) utils.save(model, os.path.join(args.save, 'weights.pt')) logging.info('STARTING EVALUATION') test, valid, runtime, params = naseval.eval_one_shot_model( config=args.__dict__, model=arch_filename) index = np.random.choice(list(range(3))) logging.info( 'TEST ERROR: %.3f | VALID ERROR: %.3f | RUNTIME: %f | PARAMS: %d' % (test[index], valid[index], runtime[index], params[index]))
def main(): if not 'debug' in args.save: from nasbench_analysis import eval_darts_one_shot_model_in_nasbench as naseval # Select the search space to search in if args.search_space == '1': search_space = SearchSpace1() elif args.search_space == '2': search_space = SearchSpace2() elif args.search_space == '3': search_space = SearchSpace3() else: raise ValueError('Unknown search space') 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) torch.cuda.set_device(args.gpu) cudnn.benchmark = True torch.manual_seed(args.seed) cudnn.enabled = True torch.cuda.manual_seed(args.seed) logging.info('gpu device = %d' % args.gpu) logging.info("args = %s", args) if args.perturb_alpha == 'none': perturb_alpha = None elif args.perturb_alpha == 'pgd_linf': perturb_alpha = Linf_PGD_alpha elif args.perturb_alpha == 'random': perturb_alpha = Random_alpha criterion = nn.CrossEntropyLoss() criterion = criterion.cuda() model = Network(args.init_channels, CIFAR_CLASSES, args.layers, criterion, output_weights=args.output_weights, steps=search_space.num_intermediate_nodes, search_space=search_space) model = model.cuda() logging.info("param size = %fMB", utils.count_parameters_in_MB(model)) optimizer = torch.optim.SGD(model.parameters(), args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay) train_transform, valid_transform = utils._data_transforms_cifar10(args) train_data = dset.CIFAR10(root=args.data, train=True, download=True, transform=train_transform) num_train = len(train_data) indices = list(range(num_train)) split = int(np.floor(args.train_portion * num_train)) if 'debug' in args.save: split = args.batch_size num_train = 2 * args.batch_size 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) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, float(args.epochs), eta_min=args.learning_rate_min) analyzer = Analyzer(model, args) architect = Architect(model, args) for epoch in range(args.epochs): scheduler.step() lr = scheduler.get_lr()[0] if args.cutout: # increase the cutout probability linearly throughout search train_transform.transforms[ -1].cutout_prob = args.cutout_prob * epoch / (args.epochs - 1) logging.info('epoch %d lr %e cutout_prob %e', epoch, lr, train_transform.transforms[-1].cutout_prob) else: logging.info('epoch %d lr %e', epoch, lr) if args.perturb_alpha: epsilon_alpha = 0.03 + (args.epsilon_alpha - 0.03) * epoch / args.epochs logging.info('epoch %d epsilon_alpha %e', epoch, epsilon_alpha) # Save the one shot model architecture weights for later analysis arch_filename = os.path.join( args.save, 'one_shot_architecture_{}.obj'.format(epoch)) with open(arch_filename, 'wb') as filehandler: numpy_tensor_list = [] for tensor in model.arch_parameters(): numpy_tensor_list.append(tensor.detach().cpu().numpy()) pickle.dump(numpy_tensor_list, filehandler) # # Save the entire one-shot-model # filepath = os.path.join(args.save, 'one_shot_model_{}.obj'.format(epoch)) # torch.save(model.state_dict(), filepath) if not 'debug' in args.save: for i in numpy_tensor_list: logging.info(str(i)) # training train_acc, train_obj, ev = train(train_queue, valid_queue, model, architect, criterion, optimizer, lr, epoch, analyzer, perturb_alpha, epsilon_alpha) logging.info('train_acc %f', train_acc) logging.info('eigenvalue %f', ev) writer.add_scalar('Acc/train', train_acc, epoch) writer.add_scalar('Obj/train', train_obj, epoch) writer.add_scalar('Analysis/eigenvalue', ev, epoch) # validation valid_acc, valid_obj = infer(valid_queue, model, criterion) logging.info('valid_acc %f', valid_acc) writer.add_scalar('Acc/valid', valid_acc, epoch) writer.add_scalar('Obj/valid', valid_obj, epoch) utils.save(model, os.path.join(args.save, 'weights.pt')) if not 'debug' in args.save: # benchmark logging.info('STARTING EVALUATION') test, valid, runtime, params = naseval.eval_one_shot_model( config=args.__dict__, model=arch_filename) index = np.random.choice(list(range(3))) test, valid, runtime, params = np.mean(test), np.mean( valid), np.mean(runtime), np.mean(params) logging.info( 'TEST ERROR: %.3f | VALID ERROR: %.3f | RUNTIME: %f | PARAMS: %d' % (test, valid, runtime, params)) writer.add_scalar('Analysis/test', test, epoch) writer.add_scalar('Analysis/valid', valid, epoch) writer.add_scalar('Analysis/runtime', runtime, epoch) writer.add_scalar('Analysis/params', params, epoch) writer.close()