def test_testtube_pickle():
"""
Verify that pickling a trainer containing a test tube logger works
"""
tutils.reset_seed()
hparams = tutils.get_hparams()
model = LightningTestModel(hparams)
save_dir = tutils.init_save_dir()
logger = tutils.get_test_tube_logger(False)
logger.log_hyperparams(hparams)
logger.save()
trainer_options = dict(max_nb_epochs=1,
train_percent_check=0.01,
logger=logger)
trainer = Trainer(**trainer_options)
pkl_bytes = pickle.dumps(trainer)
trainer2 = pickle.loads(pkl_bytes)
trainer2.logger.log_metrics({"acc": 1.0})
tutils.clear_save_dir()
def test_simple_cpu():
"""
Verify continue training session on CPU
:return:
"""
tutils.reset_seed()
hparams = tutils.get_hparams()
model = LightningTestModel(hparams)
save_dir = tutils.init_save_dir()
# logger file to get meta
trainer_options = dict(
max_nb_epochs=1,
val_percent_check=0.1,
train_percent_check=0.1,
)
# fit model
trainer = Trainer(**trainer_options)
result = trainer.fit(model)
# traning complete
assert result == 1, 'amp + ddp model failed to complete'
tutils.clear_save_dir()
def test_running_test_after_fitting():
"""Verify test() on fitted model"""
tutils.reset_seed()
hparams = tutils.get_hparams()
model = LightningTestModel(hparams)
save_dir = tutils.init_save_dir()
# logger file to get meta
logger = tutils.get_test_tube_logger(False)
# logger file to get weights
checkpoint = tutils.init_checkpoint_callback(logger)
trainer_options = dict(show_progress_bar=False,
max_nb_epochs=1,
train_percent_check=0.4,
val_percent_check=0.2,
test_percent_check=0.2,
checkpoint_callback=checkpoint,
logger=logger)
# fit model
trainer = Trainer(**trainer_options)
result = trainer.fit(model)
assert result == 1, 'training failed to complete'
trainer.test()
# test we have good test accuracy
tutils.assert_ok_test_acc(trainer)
tutils.clear_save_dir()
def test_model_saving_loading():
"""
Tests use case where trainer saves the model, and user loads it from tags independently
:return:
"""
tutils.reset_seed()
hparams = tutils.get_hparams()
model = LightningTestModel(hparams)
save_dir = tutils.init_save_dir()
# logger file to get meta
logger = tutils.get_test_tube_logger(False)
trainer_options = dict(
max_nb_epochs=1,
logger=logger,
checkpoint_callback=ModelCheckpoint(save_dir)
)
# fit model
trainer = Trainer(**trainer_options)
result = trainer.fit(model)
# traning complete
assert result == 1, 'amp + ddp model failed to complete'
# make a prediction
for dataloader in model.test_dataloader():
for batch in dataloader:
break
x, y = batch
x = x.view(x.size(0), -1)
# generate preds before saving model
model.eval()
pred_before_saving = model(x)
# save model
new_weights_path = os.path.join(save_dir, 'save_test.ckpt')
trainer.save_checkpoint(new_weights_path)
# load new model
tags_path = logger.experiment.get_data_path(logger.experiment.name, logger.experiment.version)
tags_path = os.path.join(tags_path, 'meta_tags.csv')
model_2 = LightningTestModel.load_from_metrics(weights_path=new_weights_path,
tags_csv=tags_path)
model_2.eval()
# make prediction
# assert that both predictions are the same
new_pred = model_2(x)
assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1
tutils.clear_save_dir()
def test_running_test_pretrained_model_ddp():
"""Verify test() on pretrained model"""
if not tutils.can_run_gpu_test():
return
tutils.reset_seed()
tutils.set_random_master_port()
hparams = tutils.get_hparams()
model = LightningTestModel(hparams)
save_dir = tutils.init_save_dir()
# exp file to get meta
logger = tutils.get_test_tube_logger(False)
# exp file to get weights
checkpoint = tutils.init_checkpoint_callback(logger)
trainer_options = dict(
show_progress_bar=False,
max_nb_epochs=1,
train_percent_check=0.4,
val_percent_check=0.2,
checkpoint_callback=checkpoint,
logger=logger,
gpus=[0, 1],
distributed_backend='ddp'
)
# fit model
trainer = Trainer(**trainer_options)
result = trainer.fit(model)
exp = logger.experiment
logging.info(os.listdir(exp.get_data_path(exp.name, exp.version)))
# correct result and ok accuracy
assert result == 1, 'training failed to complete'
pretrained_model = tutils.load_model(logger.experiment,
trainer.checkpoint_callback.filepath,
module_class=LightningTestModel)
# run test set
new_trainer = Trainer(**trainer_options)
new_trainer.test(pretrained_model)
for dataloader in model.test_dataloader():
tutils.run_prediction(dataloader, pretrained_model)
tutils.clear_save_dir()
def test_running_test_pretrained_model_dp():
tutils.reset_seed()
"""Verify test() on pretrained model"""
if not tutils.can_run_gpu_test():
return
hparams = tutils.get_hparams()
model = LightningTestModel(hparams)
save_dir = tutils.init_save_dir()
# logger file to get meta
logger = tutils.get_test_tube_logger(False)
# logger file to get weights
checkpoint = tutils.init_checkpoint_callback(logger)
trainer_options = dict(
show_progress_bar=True,
max_nb_epochs=1,
train_percent_check=0.4,
val_percent_check=0.2,
checkpoint_callback=checkpoint,
logger=logger,
gpus=[0, 1],
distributed_backend='dp'
)
# fit model
trainer = Trainer(**trainer_options)
result = trainer.fit(model)
# correct result and ok accuracy
assert result == 1, 'training failed to complete'
pretrained_model = tutils.load_model(logger.experiment,
trainer.checkpoint_callback.filepath,
module_class=LightningTestModel)
new_trainer = Trainer(**trainer_options)
new_trainer.test(pretrained_model)
# test we have good test accuracy
tutils.assert_ok_test_acc(new_trainer)
tutils.clear_save_dir()
def test_no_val_end_module():
"""
Tests use case where trainer saves the model, and user loads it from tags independently
:return:
"""
tutils.reset_seed()
class CurrentTestModel(LightningValidationStepMixin,
LightningTestModelBase):
pass
hparams = tutils.get_hparams()
model = CurrentTestModel(hparams)
save_dir = tutils.init_save_dir()
# logger file to get meta
logger = tutils.get_test_tube_logger(False)
trainer_options = dict(max_nb_epochs=1,
logger=logger,
checkpoint_callback=ModelCheckpoint(save_dir))
# fit model
trainer = Trainer(**trainer_options)
result = trainer.fit(model)
# traning complete
assert result == 1, 'amp + ddp model failed to complete'
# save model
new_weights_path = os.path.join(save_dir, 'save_test.ckpt')
trainer.save_checkpoint(new_weights_path)
# load new model
tags_path = logger.experiment.get_data_path(logger.experiment.name,
logger.experiment.version)
tags_path = os.path.join(tags_path, 'meta_tags.csv')
model_2 = LightningTestModel.load_from_metrics(
weights_path=new_weights_path, tags_csv=tags_path)
model_2.eval()
# make prediction
tutils.clear_save_dir()
def test_load_model_from_checkpoint():
tutils.reset_seed()
"""Verify test() on pretrained model"""
hparams = tutils.get_hparams()
model = LightningTestModel(hparams)
save_dir = tutils.init_save_dir()
trainer_options = dict(
show_progress_bar=False,
max_nb_epochs=1,
train_percent_check=0.4,
val_percent_check=0.2,
checkpoint_callback=True,
logger=False,
default_save_path=save_dir
)
# fit model
trainer = Trainer(**trainer_options)
result = trainer.fit(model)
# correct result and ok accuracy
assert result == 1, 'training failed to complete'
pretrained_model = LightningTestModel.load_from_checkpoint(
os.path.join(trainer.checkpoint_callback.filepath, "_ckpt_epoch_0.ckpt")
)
# test that hparams loaded correctly
for k, v in vars(hparams).items():
assert getattr(pretrained_model.hparams, k) == v
new_trainer = Trainer(**trainer_options)
new_trainer.test(pretrained_model)
# test we have good test accuracy
tutils.assert_ok_test_acc(new_trainer)
tutils.clear_save_dir()
def test_testtube_logger():
"""
verify that basic functionality of test tube logger works
"""
tutils.reset_seed()
hparams = tutils.get_hparams()
model = LightningTestModel(hparams)
save_dir = tutils.init_save_dir()
logger = tutils.get_test_tube_logger(False)
trainer_options = dict(max_nb_epochs=1,
train_percent_check=0.01,
logger=logger)
trainer = Trainer(**trainer_options)
result = trainer.fit(model)
assert result == 1, "Training failed"
tutils.clear_save_dir()
def test_running_test_without_val():
tutils.reset_seed()
"""Verify test() works on a model with no val_loader"""
class CurrentTestModel(LightningTestMixin, LightningTestModelBase):
pass
hparams = tutils.get_hparams()
model = CurrentTestModel(hparams)
save_dir = tutils.init_save_dir()
# logger file to get meta
logger = tutils.get_test_tube_logger(False)
# logger file to get weights
checkpoint = tutils.init_checkpoint_callback(logger)
trainer_options = dict(show_progress_bar=False,
max_nb_epochs=1,
train_percent_check=0.4,
val_percent_check=0.2,
test_percent_check=0.2,
checkpoint_callback=checkpoint,
logger=logger)
# fit model
trainer = Trainer(**trainer_options)
result = trainer.fit(model)
assert result == 1, 'training failed to complete'
trainer.test()
# test we have good test accuracy
tutils.assert_ok_test_acc(trainer)
tutils.clear_save_dir()
def test_cpu_restore_training():
"""
Verify continue training session on CPU
:return:
"""
tutils.reset_seed()
hparams = tutils.get_hparams()
model = LightningTestModel(hparams)
save_dir = tutils.init_save_dir()
# logger file to get meta
test_logger_version = 10
logger = tutils.get_test_tube_logger(False, version=test_logger_version)
trainer_options = dict(
max_nb_epochs=2,
val_check_interval=0.50,
val_percent_check=0.2,
train_percent_check=0.2,
logger=logger,
checkpoint_callback=ModelCheckpoint(save_dir)
)
# fit model
trainer = Trainer(**trainer_options)
result = trainer.fit(model)
real_global_epoch = trainer.current_epoch
# traning complete
assert result == 1, 'amp + ddp model failed to complete'
# wipe-out trainer and model
# retrain with not much data... this simulates picking training back up after slurm
# we want to see if the weights come back correctly
new_logger = tutils.get_test_tube_logger(False, version=test_logger_version)
trainer_options = dict(
max_nb_epochs=2,
val_check_interval=0.50,
val_percent_check=0.2,
train_percent_check=0.2,
logger=new_logger,
checkpoint_callback=ModelCheckpoint(save_dir),
)
trainer = Trainer(**trainer_options)
model = LightningTestModel(hparams)
# set the epoch start hook so we can predict before the model does the full training
def assert_good_acc():
assert trainer.current_epoch == real_global_epoch
assert trainer.current_epoch >= 0
# if model and state loaded correctly, predictions will be good even though we
# haven't trained with the new loaded model
trainer.model.eval()
for dataloader in trainer.get_val_dataloaders():
tutils.run_prediction(dataloader, trainer.model)
model.on_sanity_check_start = assert_good_acc
# by calling fit again, we trigger training, loading weights from the cluster
# and our hook to predict using current model before any more weight updates
trainer.fit(model)
tutils.clear_save_dir()
def test_dp_resume():
"""
Make sure DP continues training correctly
:return:
"""
if not tutils.can_run_gpu_test():
return
tutils.reset_seed()
hparams = tutils.get_hparams()
model = LightningTestModel(hparams)
trainer_options = dict(
show_progress_bar=True,
max_nb_epochs=2,
gpus=2,
distributed_backend='dp',
)
save_dir = tutils.init_save_dir()
# get logger
logger = tutils.get_test_tube_logger(debug=False)
# exp file to get weights
# logger file to get weights
checkpoint = tutils.init_checkpoint_callback(logger)
# add these to the trainer options
trainer_options['logger'] = logger
trainer_options['checkpoint_callback'] = checkpoint
# fit model
trainer = Trainer(**trainer_options)
trainer.is_slurm_managing_tasks = True
result = trainer.fit(model)
# track epoch before saving
real_global_epoch = trainer.current_epoch
# correct result and ok accuracy
assert result == 1, 'amp + dp model failed to complete'
# ---------------------------
# HPC LOAD/SAVE
# ---------------------------
# save
trainer.hpc_save(save_dir, logger)
# init new trainer
new_logger = tutils.get_test_tube_logger(version=logger.version)
trainer_options['logger'] = new_logger
trainer_options['checkpoint_callback'] = ModelCheckpoint(save_dir)
trainer_options['train_percent_check'] = 0.2
trainer_options['val_percent_check'] = 0.2
trainer_options['max_nb_epochs'] = 1
new_trainer = Trainer(**trainer_options)
# set the epoch start hook so we can predict before the model does the full training
def assert_good_acc():
assert new_trainer.current_epoch == real_global_epoch and new_trainer.current_epoch > 0
# if model and state loaded correctly, predictions will be good even though we
# haven't trained with the new loaded model
dp_model = new_trainer.model
dp_model.eval()
dataloader = trainer.get_train_dataloader()
tutils.run_prediction(dataloader, dp_model, dp=True)
# new model
model = LightningTestModel(hparams)
model.on_sanity_check_start = assert_good_acc
# fit new model which should load hpc weights
new_trainer.fit(model)
# test freeze on gpu
model.freeze()
model.unfreeze()
tutils.clear_save_dir()
def test_amp_gpu_ddp_slurm_managed():
"""
Make sure DDP + AMP work
:return:
"""
if not tutils.can_run_gpu_test():
return
tutils.reset_seed()
# simulate setting slurm flags
tutils.set_random_master_port()
os.environ['SLURM_LOCALID'] = str(0)
hparams = tutils.get_hparams()
model = LightningTestModel(hparams)
trainer_options = dict(show_progress_bar=True,
max_nb_epochs=1,
gpus=[0],
distributed_backend='ddp',
use_amp=True)
save_dir = tutils.init_save_dir()
# exp file to get meta
logger = tutils.get_test_tube_logger(False)
# exp file to get weights
checkpoint = tutils.init_checkpoint_callback(logger)
# add these to the trainer options
trainer_options['checkpoint_callback'] = checkpoint
trainer_options['logger'] = logger
# fit model
trainer = Trainer(**trainer_options)
trainer.is_slurm_managing_tasks = True
result = trainer.fit(model)
# correct result and ok accuracy
assert result == 1, 'amp + ddp model failed to complete'
# test root model address
assert trainer.resolve_root_node_address('abc') == 'abc'
assert trainer.resolve_root_node_address('abc[23]') == 'abc23'
assert trainer.resolve_root_node_address('abc[23-24]') == 'abc23'
assert trainer.resolve_root_node_address(
'abc[23-24, 45-40, 40]') == 'abc23'
# test model loading with a map_location
pretrained_model = tutils.load_model(logger.experiment,
trainer.checkpoint_callback.filepath)
# test model preds
for dataloader in trainer.get_test_dataloaders():
tutils.run_prediction(dataloader, pretrained_model)
if trainer.use_ddp:
# on hpc this would work fine... but need to hack it for the purpose of the test
trainer.model = pretrained_model
trainer.optimizers, trainer.lr_schedulers = pretrained_model.configure_optimizers(
)
# test HPC loading / saving
trainer.hpc_save(save_dir, logger)
trainer.hpc_load(save_dir, on_gpu=True)
# test freeze on gpu
model.freeze()
model.unfreeze()
tutils.clear_save_dir()
def test_tbptt_cpu_model():
"""
Test truncated back propagation through time works.
:return:
"""
tutils.reset_seed()
save_dir = tutils.init_save_dir()
truncated_bptt_steps = 2
sequence_size = 30
batch_size = 30
x_seq = torch.rand(batch_size, sequence_size, 1)
y_seq_list = torch.rand(batch_size, sequence_size, 1).tolist()
class MockSeq2SeqDataset(torch.utils.data.Dataset):
def __getitem__(self, i):
return x_seq, y_seq_list
def __len__(self):
return 1
class BpttTestModel(LightningTestModelBase):
def __init__(self, hparams):
super().__init__(hparams)
self.test_hidden = None
def training_step(self, batch, batch_idx, hiddens):
assert hiddens == self.test_hidden, "Hidden state not persistent between tbptt steps"
self.test_hidden = torch.rand(1)
x_tensor, y_list = batch
assert x_tensor.shape[
1] == truncated_bptt_steps, "tbptt split Tensor failed"
y_tensor = torch.tensor(y_list, dtype=x_tensor.dtype)
assert y_tensor.shape[
1] == truncated_bptt_steps, "tbptt split list failed"
pred = self.forward(x_tensor.view(batch_size,
truncated_bptt_steps))
loss_val = torch.nn.functional.mse_loss(
pred, y_tensor.view(batch_size, truncated_bptt_steps))
return {
'loss': loss_val,
'hiddens': self.test_hidden,
}
@data_loader
def train_dataloader(self):
return torch.utils.data.DataLoader(
dataset=MockSeq2SeqDataset(),
batch_size=batch_size,
shuffle=False,
sampler=None,
)
trainer_options = dict(
max_nb_epochs=1,
truncated_bptt_steps=truncated_bptt_steps,
val_percent_check=0,
weights_summary=None,
)
hparams = tutils.get_hparams()
hparams.batch_size = batch_size
hparams.in_features = truncated_bptt_steps
hparams.hidden_dim = truncated_bptt_steps
hparams.out_features = truncated_bptt_steps
model = BpttTestModel(hparams)
# fit model
trainer = Trainer(**trainer_options)
result = trainer.fit(model)
assert result == 1, 'training failed to complete'
tutils.clear_save_dir()
def test_cpu_slurm_save_load():
"""
Verify model save/load/checkpoint on CPU
:return:
"""
tutils.reset_seed()
hparams = tutils.get_hparams()
model = LightningTestModel(hparams)
save_dir = tutils.init_save_dir()
# logger file to get meta
logger = tutils.get_test_tube_logger(False)
version = logger.version
trainer_options = dict(max_nb_epochs=1,
logger=logger,
checkpoint_callback=ModelCheckpoint(save_dir))
# fit model
trainer = Trainer(**trainer_options)
result = trainer.fit(model)
real_global_step = trainer.global_step
# traning complete
assert result == 1, 'amp + ddp model failed to complete'
# predict with trained model before saving
# make a prediction
for dataloader in model.test_dataloader():
for batch in dataloader:
break
x, y = batch
x = x.view(x.size(0), -1)
model.eval()
pred_before_saving = model(x)
# test HPC saving
# simulate snapshot on slurm
saved_filepath = trainer.hpc_save(save_dir, logger)
assert os.path.exists(saved_filepath)
# new logger file to get meta
logger = tutils.get_test_tube_logger(False, version=version)
trainer_options = dict(
max_nb_epochs=1,
logger=logger,
checkpoint_callback=ModelCheckpoint(save_dir),
)
trainer = Trainer(**trainer_options)
model = LightningTestModel(hparams)
# set the epoch start hook so we can predict before the model does the full training
def assert_pred_same():
assert trainer.global_step == real_global_step and trainer.global_step > 0
# predict with loaded model to make sure answers are the same
trainer.model.eval()
new_pred = trainer.model(x)
assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1
model.on_epoch_start = assert_pred_same
# by calling fit again, we trigger training, loading weights from the cluster
# and our hook to predict using current model before any more weight updates
trainer.fit(model)
tutils.clear_save_dir()
def test_model_checkpoint_options():
"""
Test ModelCheckpoint options
:return:
"""
def mock_save_function(filepath):
open(filepath, 'a').close()
hparams = tutils.get_hparams()
model = LightningTestModel(hparams)
# simulated losses
save_dir = tutils.init_save_dir()
losses = [10, 9, 2.8, 5, 2.5]
# -----------------
# CASE K=-1 (all)
w = ModelCheckpoint(save_dir, save_top_k=-1, verbose=1)
w.save_function = mock_save_function
for i, loss in enumerate(losses):
w.on_epoch_end(i, logs={'val_loss': loss})
file_lists = set(os.listdir(save_dir))
assert len(file_lists) == len(
losses), "Should save all models when save_top_k=-1"
# verify correct naming
for i in range(0, len(losses)):
assert f'_ckpt_epoch_{i}.ckpt' in file_lists
tutils.clear_save_dir()
# -----------------
# CASE K=0 (none)
w = ModelCheckpoint(save_dir, save_top_k=0, verbose=1)
w.save_function = mock_save_function
for i, loss in enumerate(losses):
w.on_epoch_end(i, logs={'val_loss': loss})
file_lists = os.listdir(save_dir)
assert len(file_lists) == 0, "Should save 0 models when save_top_k=0"
tutils.clear_save_dir()
# -----------------
# CASE K=1 (2.5, epoch 4)
w = ModelCheckpoint(save_dir,
save_top_k=1,
verbose=1,
prefix='test_prefix')
w.save_function = mock_save_function
for i, loss in enumerate(losses):
w.on_epoch_end(i, logs={'val_loss': loss})
file_lists = set(os.listdir(save_dir))
assert len(file_lists) == 1, "Should save 1 model when save_top_k=1"
assert 'test_prefix_ckpt_epoch_4.ckpt' in file_lists
tutils.clear_save_dir()
# -----------------
# CASE K=2 (2.5 epoch 4, 2.8 epoch 2)
# make sure other files don't get deleted
w = ModelCheckpoint(save_dir, save_top_k=2, verbose=1)
open(f'{save_dir}/other_file.ckpt', 'a').close()
w.save_function = mock_save_function
for i, loss in enumerate(losses):
w.on_epoch_end(i, logs={'val_loss': loss})
file_lists = set(os.listdir(save_dir))
assert len(file_lists) == 3, 'Should save 2 model when save_top_k=2'
assert '_ckpt_epoch_4.ckpt' in file_lists
assert '_ckpt_epoch_2.ckpt' in file_lists
assert 'other_file.ckpt' in file_lists
tutils.clear_save_dir()
# -----------------
# CASE K=4 (save all 4 models)
# multiple checkpoints within same epoch
w = ModelCheckpoint(save_dir, save_top_k=4, verbose=1)
w.save_function = mock_save_function
for loss in losses:
w.on_epoch_end(0, logs={'val_loss': loss})
file_lists = set(os.listdir(save_dir))
assert len(
file_lists
) == 4, 'Should save all 4 models when save_top_k=4 within same epoch'
tutils.clear_save_dir()
# -----------------
# CASE K=3 (save the 2nd, 3rd, 4th model)
# multiple checkpoints within same epoch
w = ModelCheckpoint(save_dir, save_top_k=3, verbose=1)
w.save_function = mock_save_function
for loss in losses:
w.on_epoch_end(0, logs={'val_loss': loss})
file_lists = set(os.listdir(save_dir))
assert len(file_lists) == 3, 'Should save 3 models when save_top_k=3'
assert '_ckpt_epoch_0_v2.ckpt' in file_lists
assert '_ckpt_epoch_0_v1.ckpt' in file_lists
assert '_ckpt_epoch_0.ckpt' in file_lists
tutils.clear_save_dir()
