def main(args): use_cuda = not args['no_cuda'] and torch.cuda.is_available() torch.manual_seed(args['seed']) device = torch.device("cuda" if use_cuda else "cpu") kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {} #data_dir = os.path.join(args['data_dir'], nni.get_trial_id()) data_dir = os.path.join(args['data_dir'], 'data') train_loader = torch.utils.data.DataLoader(datasets.MNIST( data_dir, train=True, download=True, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307, ), (0.3081, )) ])), batch_size=args['batch_size'], shuffle=True, **kwargs) test_loader = torch.utils.data.DataLoader(datasets.MNIST( data_dir, train=False, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307, ), (0.3081, )) ])), batch_size=1000, shuffle=True, **kwargs) hidden_size = args['hidden_size'] model = Net(hidden_size=hidden_size).to(device) get_and_apply_next_architecture(model) optimizer = optim.SGD(model.parameters(), lr=args['lr'], momentum=args['momentum']) for epoch in range(1, args['epochs'] + 1): train(args, model, device, train_loader, optimizer, epoch) test_acc = test(args, model, device, test_loader) if epoch < args['epochs']: # report intermediate result nni.report_intermediate_result(test_acc) logger.debug('test accuracy %g', test_acc) logger.debug('Pipe send intermediate result done.') else: # report final result nni.report_final_result(test_acc) logger.debug('Final result is %g', test_acc) logger.debug('Send final result done.')
def test_classic_nas(self): for model_cls in self.default_cls: model = model_cls(self) get_and_apply_next_architecture(model) self.iterative_sample_and_forward(model)
def test_proxylessnas(self): model = self.model_module.LayerChoiceOnlySearchSpace(self) get_and_apply_next_architecture(model) self.iterative_sample_and_forward(model)
parser.add_argument("--train-iters", type=int, default=200) parser.add_argument("--test-batch-size", type=int, default=512) parser.add_argument("--log-frequency", type=int, default=10) args = parser.parse_args() # use a fixed set of image will improve the performance torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) np.random.seed(args.seed) random.seed(args.seed) torch.backends.cudnn.deterministic = True assert torch.cuda.is_available() model = ShuffleNetV2OneShot() criterion = CrossEntropyLabelSmooth(1000, 0.1) get_and_apply_next_architecture(model) model.load_state_dict(load_and_parse_state_dict(filepath=args.checkpoint)) model.cuda() train_loader = get_imagenet_iter_dali("train", args.imagenet_dir, args.train_batch_size, args.workers, spos_preprocessing=args.spos_preprocessing, seed=args.seed, device_id=0) val_loader = get_imagenet_iter_dali("val", args.imagenet_dir, args.test_batch_size, args.workers, spos_preprocessing=args.spos_preprocessing, shuffle=True, seed=args.seed, device_id=0) train_loader = cycle(train_loader) evaluate_acc(model, criterion, args, train_loader, val_loader)