def _main(port):
base_model = ShuffleNetV2(32)
base_predictor = 'cortexA76cpu_tflite21'
transf = [
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip()
]
normalize = [
transforms.ToTensor(),
transforms.Normalize([0.49139968, 0.48215827, 0.44653124], [0.24703233, 0.24348505, 0.26158768])
]
train_dataset = serialize(CIFAR10, 'data', train=True, download=True, transform=transforms.Compose(transf + normalize))
test_dataset = serialize(CIFAR10, 'data', train=False, transform=transforms.Compose(normalize))
trainer = pl.Classification(train_dataloader=pl.DataLoader(train_dataset, batch_size=64),
val_dataloaders=pl.DataLoader(test_dataset, batch_size=64),
max_epochs=2, gpus=1)
simple_strategy = strategy.Random(model_filter=LatencyFilter(threshold=100, predictor=base_predictor))
exp = RetiariiExperiment(base_model, trainer, strategy=simple_strategy)
exp_config = RetiariiExeConfig('local')
exp_config.trial_concurrency = 2
exp_config.max_trial_number = 2
exp_config.trial_gpu_number = 1
exp_config.training_service.use_active_gpu = False
exp_config.execution_engine = 'base'
exp_config.dummy_input = [1, 3, 32, 32]
exp.run(exp_config, port)
print('Exported models:')
for model in exp.export_top_models(formatter='dict'):
print(model)
def test_oneshot_experiment():
base_model = Net()
evaluator = get_mnist_evaluator()
search_strategy = strategy.RandomOneShot()
exp = RetiariiExperiment(base_model, evaluator, strategy=search_strategy)
exp_config = RetiariiExeConfig()
exp_config.execution_engine = 'oneshot'
exp.run(exp_config)
assert isinstance(exp.export_top_models()[0], dict)
def test_hub_oneshot(space_type, strategy_type):
NDS_SPACES = ['amoeba', 'darts', 'pnas', 'enas', 'nasnet']
if strategy_type == 'proxyless':
if 'width' in space_type or 'depth' in space_type or \
any(space_type.startswith(prefix) for prefix in NDS_SPACES + ['proxylessnas', 'mobilenetv3']):
pytest.skip('The space has used unsupported APIs.')
if strategy_type in ['darts', 'gumbel'] and space_type == 'mobilenetv3':
pytest.skip('Skip as it consumes too much memory.')
model_space = _hub_factory(space_type)
dataset_type = 'cifar10'
if 'imagenet' in space_type or space_type in [
'mobilenetv3', 'proxylessnas', 'shufflenet', 'autoformer'
]:
dataset_type = 'imagenet'
subset_size = 4
if strategy_type in ['darts', 'gumbel'] and any(
space_type.startswith(prefix)
for prefix in NDS_SPACES) and '_' in space_type:
subset_size = 2
train_dataset, valid_dataset = _dataset_factory(dataset_type,
subset=subset_size)
train_loader = pl.DataLoader(train_dataset,
batch_size=2,
num_workers=2,
shuffle=True)
valid_loader = pl.DataLoader(valid_dataset,
batch_size=2,
num_workers=2,
shuffle=False)
evaluator = pl.Classification(
train_dataloaders=train_loader,
val_dataloaders=valid_loader,
max_epochs=1,
export_onnx=False,
gpus=1 if torch.cuda.is_available() else 0, # 0 for my debug
logger=False, # disable logging and checkpoint to avoid too much log
enable_checkpointing=False,
enable_model_summary=False
# profiler='advanced'
)
# To test on final model:
# model = type(model_space).load_searched_model('darts-v2')
# evaluator.fit(model)
strategy = _strategy_factory(strategy_type, space_type)
config = RetiariiExeConfig()
config.execution_engine = 'oneshot'
experiment = RetiariiExperiment(model_space, evaluator, strategy=strategy)
experiment.run(config)
def _main(port):
base_model = ShuffleNetV2OneShot(32)
base_predictor = 'cortexA76cpu_tflite21'
transf = [
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip()
]
normalize = [
transforms.ToTensor(),
transforms.Normalize([0.49139968, 0.48215827, 0.44653124],
[0.24703233, 0.24348505, 0.26158768])
]
# FIXME
# CIFAR10 is used here temporarily.
# Actually we should load weight from supernet and evaluate on imagenet.
train_dataset = serialize(CIFAR10,
'data',
train=True,
download=True,
transform=transforms.Compose(transf + normalize))
test_dataset = serialize(CIFAR10,
'data',
train=False,
transform=transforms.Compose(normalize))
trainer = pl.Classification(train_dataloader=pl.DataLoader(train_dataset,
batch_size=64),
val_dataloaders=pl.DataLoader(test_dataset,
batch_size=64),
max_epochs=2,
gpus=1)
simple_strategy = strategy.RegularizedEvolution(model_filter=LatencyFilter(
threshold=100, predictor=base_predictor),
sample_size=1,
population_size=2,
cycles=2)
exp = RetiariiExperiment(base_model, trainer, strategy=simple_strategy)
exp_config = RetiariiExeConfig('local')
exp_config.trial_concurrency = 2
# exp_config.max_trial_number = 2
exp_config.trial_gpu_number = 1
exp_config.training_service.use_active_gpu = False
exp_config.execution_engine = 'base'
exp_config.dummy_input = [1, 3, 32, 32]
exp.run(exp_config, port)
print('Exported models:')
for i, model in enumerate(exp.export_top_models(formatter='dict')):
print(model)
with open(f'architecture_final_{i}.json', 'w') as f:
json.dump(get_archchoice_by_model(model), f, indent=4)
def _multi_trial_test(epochs, batch_size, port, benchmark):
# initalize dataset. Note that 50k+10k is used. It's a little different from paper
transf = [
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip()
]
normalize = [
transforms.ToTensor(),
transforms.Normalize([0.49139968, 0.48215827, 0.44653124],
[0.24703233, 0.24348505, 0.26158768])
]
train_dataset = serialize(CIFAR10,
'data',
train=True,
download=True,
transform=transforms.Compose(transf + normalize))
test_dataset = serialize(CIFAR10,
'data',
train=False,
transform=transforms.Compose(normalize))
# specify training hyper-parameters
training_module = NasBench101TrainingModule(max_epochs=epochs)
# FIXME: need to fix a bug in serializer for this to work
# lr_monitor = serialize(LearningRateMonitor, logging_interval='step')
trainer = pl.Trainer(max_epochs=epochs, gpus=1)
lightning = pl.Lightning(
lightning_module=training_module,
trainer=trainer,
train_dataloader=pl.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True),
val_dataloaders=pl.DataLoader(test_dataset, batch_size=batch_size),
)
strategy = Random()
model = NasBench101()
exp = RetiariiExperiment(model, lightning, [], strategy)
exp_config = RetiariiExeConfig('local')
exp_config.trial_concurrency = 2
exp_config.max_trial_number = 20
exp_config.trial_gpu_number = 1
exp_config.training_service.use_active_gpu = False
if benchmark:
exp_config.benchmark = 'nasbench101'
exp_config.execution_engine = 'benchmark'
exp.run(exp_config, port)
def _multi_trial_test(epochs, batch_size, port):
# initalize dataset. Note that 50k+10k is used. It's a little different from paper
transf = [
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip()
]
normalize = [
transforms.ToTensor(),
transforms.Normalize([x / 255 for x in [129.3, 124.1, 112.4]],
[x / 255 for x in [68.2, 65.4, 70.4]])
]
train_dataset = serialize(CIFAR100,
'data',
train=True,
download=True,
transform=transforms.Compose(transf + normalize))
test_dataset = serialize(CIFAR100,
'data',
train=False,
transform=transforms.Compose(normalize))
# specify training hyper-parameters
training_module = NasBench201TrainingModule(max_epochs=epochs)
# FIXME: need to fix a bug in serializer for this to work
# lr_monitor = serialize(LearningRateMonitor, logging_interval='step')
trainer = pl.Trainer(max_epochs=epochs, gpus=1)
lightning = pl.Lightning(
lightning_module=training_module,
trainer=trainer,
train_dataloader=pl.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True),
val_dataloaders=pl.DataLoader(test_dataset, batch_size=batch_size),
)
strategy = Random()
model = NasBench201()
exp = RetiariiExperiment(model, lightning, [], strategy)
exp_config = RetiariiExeConfig('local')
exp_config.trial_concurrency = 2
exp_config.max_trial_number = 20
exp_config.trial_gpu_number = 1
exp_config.training_service.use_active_gpu = False
exp.run(exp_config, port)
def _test_experiment_in_separate_process(rootpath):
try:
base_model, evaluator = _mnist_net('simple', {'max_epochs': 1})
search_strategy = strategy.Random()
exp = RetiariiExperiment(base_model, evaluator, strategy=search_strategy)
exp_config = RetiariiExeConfig('local')
exp_config.experiment_name = 'mnist_unittest'
exp_config.trial_concurrency = 1
exp_config.max_trial_number = 1
exp_config._trial_command_params = nas_experiment_trial_params(rootpath)
exp.run(exp_config)
ensure_success(exp)
assert isinstance(exp.export_top_models()[0], dict)
finally:
# https://stackoverflow.com/questions/34506638/how-to-register-atexit-function-in-pythons-multiprocessing-subprocess
import atexit
atexit._run_exitfuncs()
def _test_strategy(strategy_, support_value_choice=True):
to_test = [
# (model, evaluator), support_or_net
(_mnist_net('simple'), True),
(_mnist_net('simple_value_choice'), support_value_choice),
(_mnist_net('value_choice'), support_value_choice),
(_mnist_net('repeat'), False), # no strategy supports repeat currently
(_mnist_net('custom_op'), False), # this is definitely a NO
(_multihead_attention_net(), support_value_choice),
]
for (base_model, evaluator), support_or_not in to_test:
if isinstance(strategy_, BaseStrategy):
strategy = strategy_
else:
strategy = strategy_(base_model, evaluator)
print('Testing:',
type(strategy).__name__,
type(base_model).__name__,
type(evaluator).__name__, support_or_not)
experiment = RetiariiExperiment(base_model,
evaluator,
strategy=strategy)
config = RetiariiExeConfig()
config.execution_engine = 'oneshot'
if support_or_not:
experiment.run(config)
assert isinstance(experiment.export_top_models()[0], dict)
else:
with pytest.raises(TypeError, match='not supported'):
experiment.run(config)
def test_multitrial_experiment(pytestconfig):
base_model = Net()
evaluator = get_mnist_evaluator()
search_strategy = strategy.Random()
exp = RetiariiExperiment(base_model, evaluator, strategy=search_strategy)
exp_config = RetiariiExeConfig('local')
exp_config.trial_concurrency = 1
exp_config.max_trial_number = 1
exp_config._trial_command_params = nas_experiment_trial_params(pytestconfig.rootpath)
exp.run(exp_config)
ensure_success(exp)
assert isinstance(exp.export_top_models()[0], dict)
exp.stop()
def test_multi_trial(model, pytestconfig):
evaluator_kwargs = {
'max_epochs': 1
}
base_model, evaluator = _mnist_net(model, evaluator_kwargs)
search_strategy = strategy.Random()
exp = RetiariiExperiment(base_model, evaluator, strategy=search_strategy)
exp_config = RetiariiExeConfig('local')
exp_config.experiment_name = 'mnist_unittest'
exp_config.trial_concurrency = 1
exp_config.max_trial_number = 1
exp_config._trial_command_params = nas_experiment_trial_params(pytestconfig.rootpath)
exp.run(exp_config)
ensure_success(exp)
assert isinstance(exp.export_top_models()[0], dict)
exp.stop()
def ensure_success(exp: RetiariiExperiment):
# check experiment directory exists
exp_dir = os.path.join(
exp.config.canonical_copy().experiment_working_directory,
exp.id
)
assert os.path.exists(exp_dir) and os.path.exists(os.path.join(exp_dir, 'trials'))
# check job status
job_stats = exp.get_job_statistics()
if not (len(job_stats) == 1 and job_stats[0]['trialJobStatus'] == 'SUCCEEDED'):
print('Experiment jobs did not all succeed. Status is:', job_stats, file=sys.stderr)
print('Trying to fetch trial logs.', file=sys.stderr)
for root, _, files in os.walk(os.path.join(exp_dir, 'trials')):
for file in files:
fpath = os.path.join(root, file)
print('=' * 10 + ' ' + fpath + ' ' + '=' * 10, file=sys.stderr)
print(open(fpath).read(), file=sys.stderr)
raise RuntimeError('Experiment jobs did not all succeed.')
def _test_strategy(strategy_, support_value_choice=True, multi_gpu=False):
evaluator_kwargs = {'max_epochs': 1}
if multi_gpu:
evaluator_kwargs.update(strategy='ddp',
accelerator='gpu',
devices=torch.cuda.device_count())
to_test = [
# (model, evaluator), support_or_net
(_mnist_net('simple', evaluator_kwargs), True),
(_mnist_net('simple_value_choice',
evaluator_kwargs), support_value_choice),
(_mnist_net('value_choice', evaluator_kwargs), support_value_choice),
(_mnist_net('repeat', evaluator_kwargs),
support_value_choice), # no strategy supports repeat currently
(_mnist_net('custom_op',
evaluator_kwargs), False), # this is definitely a NO
(_multihead_attention_net(evaluator_kwargs), support_value_choice),
]
for (base_model, evaluator), support_or_not in to_test:
if isinstance(strategy_, BaseStrategy):
strategy = strategy_
else:
strategy = strategy_(base_model, evaluator)
print('Testing:',
type(strategy).__name__,
type(base_model).__name__,
type(evaluator).__name__, support_or_not)
experiment = RetiariiExperiment(base_model,
evaluator,
strategy=strategy)
config = RetiariiExeConfig()
config.execution_engine = 'oneshot'
if support_or_not:
experiment.run(config)
assert isinstance(experiment.export_top_models()[0], dict)
else:
with pytest.raises(TypeError, match='not supported'):
experiment.run(config)
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
valid_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
train_dataset = serialize(CIFAR10, root='data/cifar10', train=True, download=True, transform=train_transform)
test_dataset = serialize(CIFAR10, root='data/cifar10', train=False, download=True, transform=valid_transform)
trainer = pl.Classification(train_dataloader=pl.DataLoader(train_dataset, batch_size=100),
val_dataloaders=pl.DataLoader(test_dataset, batch_size=100),
max_epochs=1, limit_train_batches=0.2)
simple_strategy = strategy.Random()
exp = RetiariiExperiment(base_model, trainer, [], simple_strategy)
exp_config = RetiariiExeConfig('remote')
exp_config.experiment_name = 'darts_search'
exp_config.trial_concurrency = 2
exp_config.max_trial_number = 10
exp_config.trial_gpu_number = 1
exp_config.training_service.use_active_gpu = True
exp_config.training_service.reuse_mode = True
exp_config.training_service.gpu_indices = [0, 1, 2]
rm_conf = RemoteMachineConfig()
rm_conf.host = '127.0.0.1'
rm_conf.user = '******'
rm_conf.password = '******'
rm_conf.port = 22
root='data/mnist',
train=False,
download=True,
transform=transform)
trainer = pl.Classification(train_dataloader=pl.DataLoader(train_dataset,
batch_size=100),
val_dataloaders=pl.DataLoader(test_dataset,
batch_size=100),
max_epochs=2,
gpus=1,
limit_train_batches=0.1,
limit_val_batches=0.1)
simple_strategy = strategy.Random()
exp = RetiariiExperiment(base_model, trainer, [], simple_strategy)
exp_config = RetiariiExeConfig('local')
exp_config.experiment_name = 'mnist_search'
exp_config.trial_concurrency = 2
exp_config.max_trial_number = args.budget
exp_config.trial_gpu_number = 1
exp_config.training_service.use_active_gpu = True # Integration test GPU has a Xorg running
export_formatter = 'dict'
if args.exec == 'graph':
exp_config.execution_engine = 'base'
export_formatter = 'code'
exp.run(exp_config, args.port)
print('Final model:')
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))])
train_dataset = serialize(MNIST,
root='data/mnist',
train=True,
download=True,
transform=transform)
test_dataset = serialize(MNIST,
root='data/mnist',
train=False,
download=True,
transform=transform)
trainer = pl.Classification(train_dataloader=pl.DataLoader(train_dataset,
batch_size=100),
val_dataloaders=pl.DataLoader(test_dataset,
batch_size=100),
max_epochs=2)
simple_strategy = strategy.Random()
exp = RetiariiExperiment(base_model, trainer, [], simple_strategy)
exp_config = RetiariiExeConfig('local')
exp_config.experiment_name = 'mnist_search'
exp_config.trial_concurrency = 2
exp_config.max_trial_number = 10
exp_config.training_service.use_active_gpu = False
exp.run(exp_config, 8081 + random.randint(0, 100))
correct_k = correct[:k].view(-1).float().sum(0)
res["acc{}".format(k)] = correct_k.mul_(1.0 / batch_size).item()
return res
if __name__ == '__main__':
base_model = CNN(32, 3, 16, 10, 8)
dataset_train, dataset_valid = get_dataset("cifar10")
criterion = nn.CrossEntropyLoss()
optim = torch.optim.SGD(base_model.parameters(),
0.025,
momentum=0.9,
weight_decay=3.0E-4)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optim,
50,
eta_min=0.001)
trainer = DartsTrainer(
model=base_model,
loss=criterion,
metrics=lambda output, target: accuracy(output, target, topk=(1, )),
optimizer=optim,
num_epochs=50,
dataset=dataset_train,
batch_size=32,
log_frequency=10,
unrolled=False)
exp = RetiariiExperiment(base_model, trainer)
exp.run()
best_val_ppl
) # reports best validation ppl to nni as final result of one trial
if __name__ == "__main__":
train_iter = WikiText2(split='train')
tokenizer = get_tokenizer('basic_english')
vocab = build_vocab_from_iterator(map(tokenizer, train_iter),
specials=['<unk>'])
vocab.set_default_index(vocab['<unk>'])
n_token = len(vocab)
base_model = Transformer(n_token)
evaluator = FunctionalEvaluator(fit)
exp = RetiariiExperiment(base_model, evaluator, [], strategy.Random())
exp_config = RetiariiExeConfig('local')
exp_config.experiment_name = 'transformer tuning'
exp_config.trial_concurrency = 3 # please change configurations accordingly
exp_config.max_trial_number = 25
exp_config.trial_gpu_number = 1
exp_config.training_service.use_active_gpu = False
export_formatter = 'dict'
exp.run(exp_config, 8081)
print('Final model:')
for model_code in exp.export_top_models(optimize_mode='minimize',
formatter=export_formatter):
print(model_code)
transform=transform)
test_dataset = serialize(MNIST,
root='data/mnist',
train=False,
download=True,
transform=transform)
trainer = pl.Classification(train_dataloader=pl.DataLoader(train_dataset,
batch_size=100),
val_dataloaders=pl.DataLoader(test_dataset,
batch_size=100),
max_epochs=2)
# uncomment the following two lines to debug a generated model
#debug_mutated_model(base_model, trainer, [])
#exit(0)
simple_strategy = strategy.Random()
exp = RetiariiExperiment(base_model, trainer, [], simple_strategy)
exp_config = RetiariiExeConfig('local')
exp_config.experiment_name = 'mnist_search'
exp_config.trial_concurrency = 2
exp_config.max_trial_number = 2
exp_config.training_service.use_active_gpu = False
exp.run(exp_config, 8081 + random.randint(0, 100))
print('Final model:')
for model_code in exp.export_top_models():
print(model_code)
# test the model for one epoch
accuracy = test_epoch(model, device, test_loader)
# call report intermediate result. Result can be float or dict
nni.report_intermediate_result(accuracy)
# report final test result
nni.report_final_result(accuracy)
if __name__ == '__main__':
base_model = Net()
search_strategy = strategy.Random()
model_evaluator = FunctionalEvaluator(evaluate_model)
exp = RetiariiExperiment(base_model, model_evaluator, [], search_strategy)
exp_config = RetiariiExeConfig('local')
exp_config.experiment_name = 'mnist_search'
exp_config.trial_concurrency = 2
exp_config.max_trial_number = 20
exp_config.training_service.use_active_gpu = False
export_formatter = 'dict'
# uncomment this for graph-based execution engine
# exp_config.execution_engine = 'base'
# export_formatter = 'code'
exp.run(exp_config, 8080)
print('Final model:')
for model_code in exp.export_top_models(formatter=export_formatter):
def _main():
parser = argparse.ArgumentParser("SPOS Evolutional Search")
parser.add_argument("--port", type=int, default=8084)
parser.add_argument("--imagenet-dir", type=str, default="./data/imagenet")
parser.add_argument("--checkpoint",
type=str,
default="./data/checkpoint-150000.pth.tar")
parser.add_argument(
"--spos-preprocessing",
action="store_true",
default=False,
help="When true, image values will range from 0 to 255 and use BGR "
"(as in original repo).")
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--workers", type=int, default=6)
parser.add_argument("--train-batch-size", type=int, default=128)
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)
parser.add_argument("--label-smoothing", type=float, default=0.1)
parser.add_argument("--evolution-sample-size", type=int, default=10)
parser.add_argument("--evolution-population-size", type=int, default=50)
parser.add_argument("--evolution-cycles", type=int, default=10)
parser.add_argument(
"--latency-filter",
type=str,
default=None,
help="Apply latency filter by calling the name of the applied hardware."
)
parser.add_argument("--latency-threshold", type=float, default=100)
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()
base_model = ShuffleNetV2OneShot()
criterion = CrossEntropyLabelSmooth(1000, args.label_smoothing)
if args.latency_filter:
latency_filter = LatencyFilter(threshold=args.latency_threshold,
predictor=args.latency_filter)
else:
latency_filter = None
evaluator = FunctionalEvaluator(evaluate_acc,
criterion=criterion,
args=args)
evolution_strategy = strategy.RegularizedEvolution(
model_filter=latency_filter,
sample_size=args.evolution_sample_size,
population_size=args.evolution_population_size,
cycles=args.evolution_cycles)
exp = RetiariiExperiment(base_model,
evaluator,
strategy=evolution_strategy)
exp_config = RetiariiExeConfig('local')
exp_config.trial_concurrency = 2
exp_config.trial_gpu_number = 1
exp_config.max_trial_number = args.evolution_cycles
exp_config.training_service.use_active_gpu = False
exp_config.execution_engine = 'base'
exp_config.dummy_input = [1, 3, 224, 224]
exp.run(exp_config, args.port)
print('Exported models:')
for i, model in enumerate(exp.export_top_models(formatter='dict')):
print(model)
with open(f'architecture_final_{i}.json', 'w') as f:
json.dump(get_archchoice_by_model(model), f, indent=4)
transform=valid_transform)
# trainer = pl.Classification(train_dataloader=pl.DataLoader(train_dataset, batch_size=100),
# val_dataloaders=pl.DataLoader(test_dataset, batch_size=100),
# max_epochs=1, limit_train_batches=0.2)
trainer = cgo.Classification(train_dataloader=pl.DataLoader(
train_dataset, batch_size=100),
val_dataloaders=pl.DataLoader(test_dataset,
batch_size=100),
max_epochs=1,
limit_train_batches=0.2)
applied_mutators = [BlockMutator('mutable_0'), BlockMutator('mutable_1')]
simple_strategy = TPEStrategy()
exp = RetiariiExperiment(base_model, trainer, applied_mutators,
simple_strategy)
exp_config = RetiariiExeConfig('remote')
exp_config.experiment_name = 'darts_search'
exp_config.trial_concurrency = 3
exp_config.max_trial_number = 10
exp_config.trial_gpu_number = 1
exp_config.training_service.reuse_mode = True
rm_conf = RemoteMachineConfig()
rm_conf.host = '127.0.0.1'
rm_conf.user = '******'
rm_conf.password = '******'
rm_conf.port = 22
rm_conf.python_path = '/home/xxx/py38/bin'
rm_conf.gpu_indices = [0, 1, 2]
# %%
#
# The ``train_epoch`` and ``test_epoch`` here can be any customized function,
# where users can write their own training recipe.
#
# It is recommended that the ``evaluate_model`` here accepts no additional arguments other than ``model_cls``.
# However, in the :doc:`advanced tutorial </nas/evaluator>`, we will show how to use additional arguments in case you actually need those.
# In future, we will support mutation on the arguments of evaluators, which is commonly called "Hyper-parmeter tuning".
#
# Launch an Experiment
# --------------------
#
# After all the above are prepared, it is time to start an experiment to do the model search. An example is shown below.
from nni.retiarii.experiment.pytorch import RetiariiExperiment, RetiariiExeConfig
exp = RetiariiExperiment(model_space, evaluator, [], search_strategy)
exp_config = RetiariiExeConfig('local')
exp_config.experiment_name = 'mnist_search'
# %%
# The following configurations are useful to control how many trials to run at most / at the same time.
exp_config.max_trial_number = 4 # spawn 4 trials at most
exp_config.trial_concurrency = 2 # will run two trials concurrently
# %%
# Remember to set the following config if you want to GPU.
# ``use_active_gpu`` should be set true if you wish to use an occupied GPU (possibly running a GUI).
exp_config.trial_gpu_number = 1
exp_config.training_service.use_active_gpu = True