mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-09-10 12:21:57 +08:00
decouple returns from each step (#307)
* decoupled training metrics from logging metrics * decoupled validation metrics from log metrics * updated docs * updated docs * updated docs * Fixed test * merged master * merged master * merged master * merged master * merged master * merged master * merged master * merged master * merged master * merged master * merged master * merged master * merged master * merged master * merged master * merged master * merged master * merged master * merged master * merged master * merged master * merged master * merged master * merged master * merged master
This commit is contained in:
William Falcon
+33
-156
@@ -14,6 +14,7 @@ from torch.utils.data import DataLoader
|
||||
from torchvision.datasets import MNIST
|
||||
import numpy as np
|
||||
import pdb
|
||||
from . import test_models
|
||||
|
||||
|
||||
class CoolModel(pl.LightningModule):
|
||||
@@ -59,156 +60,6 @@ class CoolModel(pl.LightningModule):
|
||||
return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
|
||||
|
||||
|
||||
def get_model(use_test_model=False):
|
||||
# set up model with these hyperparams
|
||||
hparams = get_hparams()
|
||||
|
||||
if use_test_model:
|
||||
model = LightningTestModel(hparams)
|
||||
else:
|
||||
model = LightningTemplateModel(hparams)
|
||||
|
||||
return model, hparams
|
||||
|
||||
|
||||
def get_exp(debug=True, version=None):
|
||||
# set up exp object without actually saving logs
|
||||
root_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
save_dir = os.path.join(root_dir, 'save_dir')
|
||||
exp = Experiment(debug=debug, save_dir=save_dir, name='tests_tt_dir', version=version)
|
||||
return exp
|
||||
|
||||
|
||||
def init_save_dir():
|
||||
root_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
save_dir = os.path.join(root_dir, 'save_dir')
|
||||
|
||||
if os.path.exists(save_dir):
|
||||
shutil.rmtree(save_dir)
|
||||
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
return save_dir
|
||||
|
||||
|
||||
def clear_save_dir():
|
||||
root_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
save_dir = os.path.join(root_dir, 'save_dir')
|
||||
if os.path.exists(save_dir):
|
||||
shutil.rmtree(save_dir)
|
||||
|
||||
|
||||
def load_model(exp, save_dir, on_gpu, map_location=None, module_class=LightningTemplateModel):
|
||||
|
||||
# load trained model
|
||||
tags_path = exp.get_data_path(exp.name, exp.version)
|
||||
tags_path = os.path.join(tags_path, 'meta_tags.csv')
|
||||
|
||||
checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
|
||||
weights_dir = os.path.join(save_dir, checkpoints[0])
|
||||
|
||||
trained_model = module_class.load_from_metrics(weights_path=weights_dir,
|
||||
tags_csv=tags_path,
|
||||
on_gpu=on_gpu,
|
||||
)
|
||||
|
||||
assert trained_model is not None, 'loading model failed'
|
||||
|
||||
return trained_model
|
||||
|
||||
|
||||
def run_prediction(dataloader, trained_model):
|
||||
# run prediction on 1 batch
|
||||
for batch in dataloader:
|
||||
break
|
||||
|
||||
x, y = batch
|
||||
x = x.view(x.size(0), -1)
|
||||
|
||||
y_hat = trained_model(x)
|
||||
|
||||
# acc
|
||||
labels_hat = torch.argmax(y_hat, dim=1)
|
||||
val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
|
||||
val_acc = torch.tensor(val_acc)
|
||||
val_acc = val_acc.item()
|
||||
assert val_acc > 0.70, 'this model is expected to get > 0.7 in test set (it got %f)' % val_acc
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------
|
||||
def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
|
||||
save_dir = init_save_dir()
|
||||
|
||||
# exp file to get meta
|
||||
exp = get_exp(False)
|
||||
exp.argparse(hparams)
|
||||
exp.save()
|
||||
|
||||
# exp file to get weights
|
||||
checkpoint = ModelCheckpoint(save_dir)
|
||||
|
||||
# add these to the trainer options
|
||||
trainer_options['checkpoint_callback'] = checkpoint
|
||||
trainer_options['experiment'] = exp
|
||||
|
||||
# fit model
|
||||
trainer = Trainer(**trainer_options)
|
||||
result = trainer.fit(model)
|
||||
|
||||
# correct result and ok accuracy
|
||||
assert result == 1, 'amp + ddp model failed sto complete'
|
||||
|
||||
# test model loading
|
||||
pretrained_model = load_model(exp, save_dir, on_gpu)
|
||||
|
||||
# test new model accuracy
|
||||
run_prediction(model.test_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, exp)
|
||||
trainer.hpc_load(save_dir, on_gpu=on_gpu)
|
||||
|
||||
clear_save_dir()
|
||||
|
||||
|
||||
def assert_ok_val_acc(trainer):
|
||||
# this model should get 0.80+ acc
|
||||
acc = trainer.training_tqdm_dict['val_acc']
|
||||
assert acc > 0.50, f'model failed to get expected 0.50 validation accuracy. Got: {acc}'
|
||||
|
||||
|
||||
def assert_ok_test_acc(trainer):
|
||||
# this model should get 0.80+ acc
|
||||
acc = trainer.training_tqdm_dict['test_acc']
|
||||
assert acc > 0.50, f'model failed to get expected 0.50 validation accuracy. Got: {acc}'
|
||||
|
||||
|
||||
def get_hparams(continue_training=False, hpc_exp_number=0):
|
||||
root_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
|
||||
args = {
|
||||
'drop_prob': 0.2,
|
||||
'batch_size': 32,
|
||||
'in_features': 28 * 28,
|
||||
'learning_rate': 0.001 * 8,
|
||||
'optimizer_name': 'adam',
|
||||
'data_root': os.path.join(root_dir, 'mnist'),
|
||||
'out_features': 10,
|
||||
'hidden_dim': 1000}
|
||||
|
||||
if continue_training:
|
||||
args['test_tube_do_checkpoint_load'] = True
|
||||
args['hpc_exp_number'] = hpc_exp_number
|
||||
|
||||
hparams = Namespace(**args)
|
||||
return hparams
|
||||
|
||||
|
||||
def main():
|
||||
"""
|
||||
Make sure DDP + AMP continue training correctly
|
||||
@@ -218,19 +69,45 @@ def main():
|
||||
Make sure DDP2 works
|
||||
:return:
|
||||
"""
|
||||
os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
|
||||
model, hparams = get_model()
|
||||
hparams = test_models.get_hparams()
|
||||
model = LightningTestModel(hparams)
|
||||
|
||||
save_dir = test_models.init_save_dir()
|
||||
|
||||
# logger file to get meta
|
||||
logger = test_models.get_test_tube_logger(False)
|
||||
logger.log_hyperparams(hparams)
|
||||
logger.save()
|
||||
|
||||
# logger file to get weights
|
||||
checkpoint = ModelCheckpoint(save_dir)
|
||||
|
||||
trainer_options = dict(
|
||||
show_progress_bar=True,
|
||||
max_nb_epochs=1,
|
||||
train_percent_check=0.4,
|
||||
val_percent_check=0.2,
|
||||
gpus=2,
|
||||
print_weights_summary=True,
|
||||
distributed_backend='ddp2'
|
||||
checkpoint_callback=checkpoint,
|
||||
logger=logger,
|
||||
gpus=[0, 1],
|
||||
distributed_backend='dp'
|
||||
)
|
||||
|
||||
run_gpu_model_test(trainer_options, model, hparams)
|
||||
# fit model
|
||||
trainer = Trainer(**trainer_options)
|
||||
result = trainer.fit(model)
|
||||
|
||||
# correct result and ok accuracy
|
||||
assert result == 1, 'training failed to complete'
|
||||
pretrained_model = test_models.load_model(logger.experiment, save_dir,
|
||||
module_class=LightningTestModel)
|
||||
|
||||
new_trainer = Trainer(**trainer_options)
|
||||
new_trainer.test(pretrained_model)
|
||||
|
||||
# test we have good test accuracy
|
||||
test_models.assert_ok_test_acc(new_trainer)
|
||||
test_models.clear_save_dir()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
Reference in New Issue
Block a user