220 lines
7.0 KiB
Python
220 lines
7.0 KiB
Python
import os
|
|
import warnings
|
|
from argparse import Namespace
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from pl_examples import LightningTemplateModel
|
|
from pytorch_lightning import Trainer
|
|
from pytorch_lightning.callbacks import ModelCheckpoint
|
|
from pytorch_lightning.loggers import TestTubeLogger, TensorBoardLogger
|
|
from pytorch_lightning.testing import LightningTestModel
|
|
|
|
# generate a list of random seeds for each test
|
|
RANDOM_PORTS = list(np.random.randint(12000, 19000, 1000))
|
|
ROOT_SEED = 1234
|
|
torch.manual_seed(ROOT_SEED)
|
|
np.random.seed(ROOT_SEED)
|
|
RANDOM_SEEDS = list(np.random.randint(0, 10000, 1000))
|
|
ROOT_PATH = os.path.abspath(os.path.dirname(__file__))
|
|
|
|
|
|
def run_model_test_no_loggers(trainer_options, model, min_acc=0.50):
|
|
save_dir = trainer_options['default_save_path']
|
|
|
|
# fit model
|
|
trainer = Trainer(**trainer_options)
|
|
result = trainer.fit(model)
|
|
|
|
# correct result and ok accuracy
|
|
assert result == 1, 'amp + ddp model failed to complete'
|
|
|
|
# test model loading
|
|
pretrained_model = load_model(trainer.logger,
|
|
trainer.checkpoint_callback.filepath,
|
|
path_expt=trainer_options.get('default_save_path'))
|
|
|
|
# test new model accuracy
|
|
for dataloader in model.test_dataloader():
|
|
run_prediction(dataloader, pretrained_model, min_acc=min_acc)
|
|
|
|
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()
|
|
|
|
|
|
def run_model_test(trainer_options, model, on_gpu=True):
|
|
save_dir = trainer_options['default_save_path']
|
|
|
|
# logger file to get meta
|
|
logger = get_test_tube_logger(save_dir, False)
|
|
|
|
# logger file to get weights
|
|
checkpoint = 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)
|
|
result = trainer.fit(model)
|
|
|
|
# correct result and ok accuracy
|
|
assert result == 1, 'amp + ddp model failed to complete'
|
|
|
|
# test model loading
|
|
pretrained_model = load_model(logger, trainer.checkpoint_callback.filepath)
|
|
|
|
# test new model accuracy
|
|
[run_prediction(dataloader, pretrained_model) for dataloader in model.test_dataloader()]
|
|
|
|
if trainer.use_ddp or trainer.use_ddp2:
|
|
# 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=on_gpu)
|
|
|
|
|
|
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 get_model(use_test_model=False, lbfgs=False):
|
|
# set up model with these hyperparams
|
|
hparams = get_hparams()
|
|
if lbfgs:
|
|
setattr(hparams, 'optimizer_name', 'lbfgs')
|
|
setattr(hparams, 'learning_rate', 0.002)
|
|
|
|
if use_test_model:
|
|
model = LightningTestModel(hparams)
|
|
else:
|
|
model = LightningTemplateModel(hparams)
|
|
|
|
return model, hparams
|
|
|
|
|
|
def get_test_tube_logger(save_dir, debug=True, version=None):
|
|
# set up logger object without actually saving logs
|
|
logger = TestTubeLogger(save_dir, name='lightning_logs', debug=debug, version=version)
|
|
return logger
|
|
|
|
|
|
def get_data_path(expt_logger, path_dir=None):
|
|
# some calls contain only experiment not complete logger
|
|
expt = expt_logger.experiment if hasattr(expt_logger, 'experiment') else expt_logger
|
|
# each logger has to have these attributes
|
|
name, version = expt_logger.name, expt_logger.version
|
|
# only the test-tube experiment has such attribute
|
|
if hasattr(expt, 'get_data_path'):
|
|
return expt.get_data_path(name, version)
|
|
# the other experiments...
|
|
if not path_dir:
|
|
path_dir = ROOT_PATH
|
|
path_expt = os.path.join(path_dir, name, 'version_%s' % version)
|
|
# try if the new sub-folder exists, typical case for test-tube
|
|
if not os.path.isdir(path_expt):
|
|
path_expt = path_dir
|
|
return path_expt
|
|
|
|
|
|
def load_model(exp, root_weights_dir, module_class=LightningTemplateModel, path_expt=None):
|
|
# load trained model
|
|
path_expt_dir = get_data_path(exp, path_dir=path_expt)
|
|
tags_path = os.path.join(path_expt_dir, TensorBoardLogger.NAME_CSV_TAGS)
|
|
|
|
checkpoints = [x for x in os.listdir(root_weights_dir) if '.ckpt' in x]
|
|
weights_dir = os.path.join(root_weights_dir, checkpoints[0])
|
|
|
|
trained_model = module_class.load_from_metrics(weights_path=weights_dir,
|
|
tags_csv=tags_path)
|
|
|
|
assert trained_model is not None, 'loading model failed'
|
|
|
|
return trained_model
|
|
|
|
|
|
def run_prediction(dataloader, trained_model, dp=False, min_acc=0.50):
|
|
# run prediction on 1 batch
|
|
for batch in dataloader:
|
|
break
|
|
|
|
x, y = batch
|
|
x = x.view(x.size(0), -1)
|
|
|
|
if dp:
|
|
output = trained_model(batch, 0)
|
|
acc = output['val_acc']
|
|
acc = torch.mean(acc).item()
|
|
|
|
else:
|
|
y_hat = trained_model(x)
|
|
|
|
# acc
|
|
labels_hat = torch.argmax(y_hat, dim=1)
|
|
acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
|
|
acc = torch.tensor(acc)
|
|
acc = acc.item()
|
|
|
|
assert acc > min_acc, f'this model is expected to get > {min_acc} in test set (it got {acc})'
|
|
|
|
|
|
def assert_ok_model_acc(trainer, key='test_acc', thr=0.4):
|
|
# this model should get 0.80+ acc
|
|
acc = trainer.training_tqdm_dict[key]
|
|
assert acc > thr, f'Model failed to get expected {thr} accuracy. {key} = {acc}'
|
|
|
|
|
|
def can_run_gpu_test():
|
|
if not torch.cuda.is_available():
|
|
warnings.warn('test_multi_gpu_model_ddp cannot run.'
|
|
' Rerun on a GPU node to run this test')
|
|
return False
|
|
if not torch.cuda.device_count() > 1:
|
|
warnings.warn('test_multi_gpu_model_ddp cannot run.'
|
|
' Rerun on a node with 2+ GPUs to run this test')
|
|
return False
|
|
return True
|
|
|
|
|
|
def reset_seed():
|
|
seed = RANDOM_SEEDS.pop()
|
|
torch.manual_seed(seed)
|
|
np.random.seed(seed)
|
|
|
|
|
|
def set_random_master_port():
|
|
port = RANDOM_PORTS.pop()
|
|
os.environ['MASTER_PORT'] = str(port)
|
|
|
|
|
|
def init_checkpoint_callback(logger, path_dir=None):
|
|
exp_path = get_data_path(logger, path_dir=path_dir)
|
|
ckpt_dir = os.path.join(exp_path, 'checkpoints')
|
|
checkpoint = ModelCheckpoint(ckpt_dir)
|
|
return checkpoint
|