def test_pickle_trainer():
import nni.retiarii.evaluator.pytorch.lightning as pl
from pytorch_lightning import Trainer
trainer = pl.Trainer(max_epochs=1)
data = pickle.dumps(trainer)
trainer = pickle.loads(data)
assert isinstance(trainer, Trainer)
def test_lightning_earlystop():
import nni.retiarii.evaluator.pytorch.lightning as pl
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
trainer = pl.Trainer(
callbacks=[nni.trace(EarlyStopping)(monitor="val_loss")])
trainer = nni.load(nni.dump(trainer))
assert any(
isinstance(callback, EarlyStopping) for callback in trainer.callbacks)
def test_lightning_earlystop():
import nni.retiarii.evaluator.pytorch.lightning as pl
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
trainer = pl.Trainer(
callbacks=[nni.trace(EarlyStopping)(monitor="val_loss")])
pickle_size_limit = 4096 if sys.platform == 'linux' else 32768
trainer = nni.load(nni.dump(trainer, pickle_size_limit=pickle_size_limit))
assert any(
isinstance(callback, EarlyStopping) for callback in trainer.callbacks)
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)