def _test_searchspace_on_dataset(searchspace, dataset='cifar10', arch=None): _reset() # dataset supports cifar10 and imagenet model, mutators = extract_mutation_from_pt_module(searchspace) if arch is None: model = try_mutation_until_success(model, mutators, 10) arch = { mut.mutator.label: _unpack_if_only_one(mut.samples) for mut in model.history } print('Selected model:', arch) with fixed_arch(arch): model = model.python_class(**model.python_init_params) if dataset == 'cifar10': train_data = FakeData(size=200, image_size=(3, 32, 32), num_classes=10, transform=transforms.ToTensor()) valid_data = FakeData(size=200, image_size=(3, 32, 32), num_classes=10, transform=transforms.ToTensor()) elif dataset == 'imagenet': train_data = FakeData(size=200, image_size=(3, 224, 224), num_classes=1000, transform=transforms.ToTensor()) valid_data = FakeData(size=200, image_size=(3, 224, 224), num_classes=1000, transform=transforms.ToTensor()) train_dataloader = pl.DataLoader(train_data, batch_size=4, shuffle=True) valid_dataloader = pl.DataLoader(valid_data, batch_size=6) evaluator = pl.Classification( train_dataloader=train_dataloader, val_dataloaders=valid_dataloader, export_onnx=False, max_epochs=1, limit_train_batches=2, limit_val_batches=3, ) evaluator.fit(model) # cleanup to avoid affecting later test cases _reset()
def _get_model_with_mutators(self, pytorch_model): return extract_mutation_from_pt_module(pytorch_model)