mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-09-09 11:32:07 +08:00
refactored model tests
This commit is contained in:
William Falcon
+91
-85
@@ -15,91 +15,9 @@ torch.manual_seed(SEED)
|
||||
np.random.seed(SEED)
|
||||
|
||||
|
||||
def get_model():
|
||||
# set up model with these hyperparams
|
||||
root_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
hparams = Namespace(**{'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})
|
||||
model = LightningTemplateModel(hparams)
|
||||
|
||||
return model, hparams
|
||||
|
||||
|
||||
def get_exp(debug=True):
|
||||
# set up exp object without actually saving logs
|
||||
root_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
exp = Experiment(debug=debug, save_dir=root_dir, name='tests_tt_dir')
|
||||
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):
|
||||
|
||||
# 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 = LightningTemplateModel.load_from_metrics(weights_path=weights_dir, tags_csv=tags_path, on_gpu=True)
|
||||
|
||||
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()
|
||||
|
||||
print(val_acc)
|
||||
|
||||
assert val_acc > 0.70, f'this model is expected to get > 0.7 in test set (it got {val_acc})'
|
||||
|
||||
|
||||
def assert_ok_acc(trainer):
|
||||
# this model should get 0.80+ acc
|
||||
acc = trainer.tng_tqdm_dic['val_acc']
|
||||
assert acc > 0.70, f'model failed to get expected 0.80 validation accuracy. Got: {acc}'
|
||||
|
||||
|
||||
# -----------------
|
||||
# TESTS
|
||||
# -----------------
|
||||
def test_cpu_model():
|
||||
"""
|
||||
Make sure model trains on CPU
|
||||
@@ -306,5 +224,93 @@ def test_amp_gpu_ddp():
|
||||
clear_save_dir()
|
||||
|
||||
|
||||
# -----------------
|
||||
# UTILS
|
||||
# -----------------
|
||||
def get_model():
|
||||
# set up model with these hyperparams
|
||||
root_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
hparams = Namespace(**{'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})
|
||||
model = LightningTemplateModel(hparams)
|
||||
|
||||
return model, hparams
|
||||
|
||||
|
||||
def get_exp(debug=True):
|
||||
# set up exp object without actually saving logs
|
||||
root_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
exp = Experiment(debug=debug, save_dir=root_dir, name='tests_tt_dir')
|
||||
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):
|
||||
|
||||
# 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 = LightningTemplateModel.load_from_metrics(weights_path=weights_dir, tags_csv=tags_path, on_gpu=True)
|
||||
|
||||
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()
|
||||
|
||||
print(val_acc)
|
||||
|
||||
assert val_acc > 0.70, f'this model is expected to get > 0.7 in test set (it got {val_acc})'
|
||||
|
||||
|
||||
def assert_ok_acc(trainer):
|
||||
# this model should get 0.80+ acc
|
||||
acc = trainer.tng_tqdm_dic['val_acc']
|
||||
assert acc > 0.70, f'model failed to get expected 0.70 validation accuracy. Got: {acc}'
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
pytest.main([__file__])
|
||||
|
||||
Reference in New Issue
Block a user