wassname/pytorch-lightning
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-lightning.git synced 2026-09-09 11:32:07 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #22 from williamFalcon/loading

Browse Source 
Loading
This commit is contained in:
William Falcon
2019-07-26 23:08:17 -04:00
committed by GitHub
parent ff1ed9db7e 587c195298
commit cf898a6ecf
8 changed files with 165 additions and 284 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/examples/new_project_templates/lightning_module_template.py
-11
View File
@@ -154,17 +154,6 @@ class LightningTemplateModel(LightningModule):
tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
return tqdm_dic
# ---------------------
# MODEL SAVING
# ---------------------
def get_save_dict(self):
checkpoint = {'state_dict': self.state_dict()}
return checkpoint
def load_model_specific(self, checkpoint):
self.load_state_dict(checkpoint['state_dict'])
pass
# ---------------------
# TRAINING SETUP
# ---------------------
pytorch_lightning/models/sample_model_template/__init__.py
View File
pytorch_lightning/models/sample_model_template/model_template.py
-203
View File
@@ -1,203 +0,0 @@
import torch.nn as nn
import numpy as np
from pytorch_lightning import LightningModule
from test_tube import HyperOptArgumentParser
from torchvision.datasets import MNIST
import torchvision.transforms as transforms
import torch
import torch.nn.functional as F
class ExampleModel1(LightningModule):
"""
Sample model to show how to define a template
"""
def __init__(self, hparams):
# init superclass
super(ExampleModel1, self).__init__(hparams)
self.batch_size = hparams.batch_size
# build model
self.__build_model()
# ---------------------
# MODEL SETUP
# ---------------------
def __build_model(self):
"""
Layout model
:return:
"""
self.c_d1 = nn.Linear(in_features=self.hparams.in_features, out_features=self.hparams.hidden_dim)
self.c_d1_bn = nn.BatchNorm1d(self.hparams.hidden_dim)
self.c_d1_drop = nn.Dropout(self.hparams.drop_prob)
self.c_d2 = nn.Linear(in_features=self.hparams.hidden_dim, out_features=self.hparams.out_features)
# ---------------------
# TRAINING
# ---------------------
def forward(self, x):
x = self.c_d1(x)
x = F.tanh(x)
x = self.c_d1_bn(x)
x = self.c_d1_drop(x)
x = self.c_d2(x)
logits = F.log_softmax(x, dim=1)
return logits
def loss(self, labels, logits):
nll = F.nll_loss(logits, labels)
return nll
def training_step(self, data_batch):
"""
Called inside the training loop
:param data_batch:
:return:
"""
# forward pass
x, y = data_batch
x = x.view(x.size(0), -1)
y_hat = self.forward(x)
# calculate loss
loss_val = self.loss(y, y_hat)
tqdm_dic = {'jefe': 1}
return loss_val, tqdm_dic
def validation_step(self, data_batch):
"""
Called inside the validation loop
:param data_batch:
:return:
"""
x, y = data_batch
x = x.view(x.size(0), -1)
y_hat = self.forward(x)
loss_val = self.loss(y, y_hat)
# acc
labels_hat = torch.argmax(y_hat, dim=1)
val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
output = {'y_hat': y_hat, 'val_loss': loss_val.item(), 'val_acc': val_acc}
return output
def validation_end(self, outputs):
"""
Called at the end of validation to aggregate outputs
:param outputs: list of individual outputs of each validation step
:return:
"""
val_loss_mean = 0
accs = []
for output in outputs:
val_loss_mean += output['val_loss']
accs.append(output['val_acc'])
val_loss_mean /= len(outputs)
tqdm_dic = {'val_loss': val_loss_mean, 'val_acc': np.mean(accs)}
return tqdm_dic
def update_tng_log_metrics(self, logs):
return logs
# ---------------------
# MODEL SAVING
# ---------------------
def get_save_dict(self):
checkpoint = {
'state_dict': self.state_dict(),
}
return checkpoint
def load_model_specific(self, checkpoint):
self.load_state_dict(checkpoint['state_dict'])
pass
# ---------------------
# TRAINING SETUP
# ---------------------
def configure_optimizers(self):
"""
return whatever optimizers we want here
:return: list of optimizers
"""
optimizer = self.choose_optimizer(self.hparams.optimizer_name, self.parameters(), {'lr': self.hparams.learning_rate}, 'optimizer')
self.optimizers = [optimizer]
return self.optimizers
def __dataloader(self, train):
# init data generators
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
dataset = MNIST(root=self.hparams.data_root, train=train, transform=transform, download=True)
loader = torch.utils.data.DataLoader(
dataset=dataset,
batch_size=self.hparams.batch_size,
shuffle=True
)
return loader
@data_loader
def tng_dataloader(self):
if self._tng_dataloader is None:
try:
self._tng_dataloader = self.__dataloader(train=True)
except Exception as e:
print(e)
raise e
return self._tng_dataloader
@property
def val_dataloader(self):
if self._val_dataloader is None:
try:
self._val_dataloader = self.__dataloader(train=False)
except Exception as e:
print(e)
raise e
return self._val_dataloader
@property
def test_dataloader(self):
if self._test_dataloader is None:
try:
self._test_dataloader = self.__dataloader(train=False)
except Exception as e:
print(e)
raise e
return self._test_dataloader
@staticmethod
def add_model_specific_args(parent_parser):
parser = HyperOptArgumentParser(strategy=parent_parser.strategy, parents=[parent_parser])
# param overwrites
# parser.set_defaults(gradient_clip=5.0)
# network params
parser.opt_list('--drop_prob', default=0.2, options=[0.2, 0.5], type=float, tunable=False)
parser.add_argument('--in_features', default=28*28)
parser.add_argument('--hidden_dim', default=500)
parser.add_argument('--out_features', default=10)
# data
parser.add_argument('--data_root', default='/Users/williamfalcon/Developer/personal/research_lib/research_proj/datasets/mnist', type=str)
# training params (opt)
parser.opt_list('--learning_rate', default=0.001, type=float, options=[0.0001, 0.0005, 0.001, 0.005],
tunable=False)
parser.opt_list('--batch_size', default=256, type=int, options=[32, 64, 128, 256], tunable=False)
parser.opt_list('--optimizer_name', default='adam', type=str, options=['adam'], tunable=False)
return parser
pytorch_lightning/models/trainer.py
+5 -2
View File
@@ -161,7 +161,6 @@ class Trainer(TrainerIO):
self.nb_tng_batches = None
self.nb_test_batches = None
# gpus come in as a string.
# if gpus = -1 then use all available devices
# otherwise, split the string using commas
@@ -611,14 +610,18 @@ class Trainer(TrainerIO):
if self.proc_rank == 0:
self.experiment.save()
# track model now.
# if cluster resets state, the model will update with the saved weights
self.model = model
# enable cluster checkpointing
# also restores training state
if self.cluster is not None: # pragma: no cover
self.enable_auto_hpc_walltime_manager()
# ---------------------------
# CORE TRAINING LOOP
# ---------------------------
self.model = model
self.__train()
def __train(self):
pytorch_lightning/root_module/model_saving.py
+23 -19
View File
@@ -4,34 +4,36 @@ import re
import pdb
from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel, LightningDataParallel
class ModelIO(object):
def load_model_specific(self, checkpoint):
def on_load_checkpoint(self, checkpoint):
"""
Do something with the checkpoint
Gives model a chance to load something before state_dict is restored
:param checkpoint:
:return:
"""
raise NotImplementedError
pass
def get_save_dict(self):
def on_save_checkpoint(self, checkpoint):
"""
Return specific things for the model
:return:
Give the model a chance to add something to the checkpoint.
state_dict is already there
"""
raise NotImplementedError
pass
# -------------------------
# OPTIONAL HOOKS
# -------------------------
def on_hpc_save(self):
def on_hpc_save(self, checkpoint):
"""
Hook to do whatever you need right before Slurm manager saves the model
:return:
"""
pass
def on_hpc_load(self):
def on_hpc_load(self, checkpoint):
"""
Hook to do whatever you need right before Slurm manager loads the model
:return:
@@ -75,12 +77,13 @@ class TrainerIO(object):
checkpoint['optimizer_states'] = optimizer_states
# request what to save from the model
# add the state_dict from the model
model = self.__get_model()
checkpoint_dict = model.get_save_dict()
checkpoint['state_dict'] = model.state_dict()
# give the model a chance to add a few things
model.on_save_checkpoint(checkpoint)
# merge trainer and model saving items
checkpoint.update(checkpoint_dict)
return checkpoint
# --------------------
@@ -149,13 +152,12 @@ class TrainerIO(object):
# give model a chance to do something on hpc_save
model = self.__get_model()
model.on_hpc_save()
checkpoint = self.dump_checkpoint()
# request what to save from the model
checkpoint_dict = self.dump_checkpoint()
model.on_hpc_save(checkpoint)
# do the actual save
torch.save(checkpoint_dict, filepath)
torch.save(checkpoint, filepath)
return filepath
@@ -167,15 +169,17 @@ class TrainerIO(object):
else:
checkpoint = torch.load(filepath, map_location=lambda storage, loc: storage)
# load training state
# load training state (affects trainer only)
self.restore_training_state(checkpoint)
# load model state
model = self.__get_model()
model.load_model_specific(checkpoint)
# load the state_dict on the model automatically
model.load_state_dict(checkpoint['state_dict'])
# call model hook
model.on_hpc_load()
model.on_hpc_load(checkpoint)
def max_ckpt_in_folder(self, path):
files = os.listdir(path)
pytorch_lightning/root_module/root_module.py
+5 -3
View File
@@ -108,11 +108,13 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
else:
checkpoint = torch.load(weights_path, map_location=lambda storage, loc: storage)
# load the state_dict on the model automatically
model = cls(hparams)
model.load_state_dict(checkpoint['state_dict'])
# give model a chance to load something
model.on_load_checkpoint(checkpoint)
# allow model to load
model.load_model_specific(checkpoint)
model.load_state_dict(checkpoint['state_dict'], strict=False)
return model
def summarize(self):
pytorch_lightning/testing_models/lm_test_module.py
-11
View File
@@ -171,17 +171,6 @@ class LightningTestModel(LightningModule):
def on_tng_metrics(self, logs):
logs['some_tensor_to_test'] = torch.rand(1)
# ---------------------
# MODEL SAVING
# ---------------------
def get_save_dict(self):
checkpoint = {'state_dict': self.state_dict()}
return checkpoint
def load_model_specific(self, checkpoint):
self.load_state_dict(checkpoint['state_dict'])
pass
# ---------------------
# TRAINING SETUP
# ---------------------
tests/test_models.py
+132 -35
View File
@@ -24,6 +24,88 @@ np.random.seed(SEED)
# ------------------------------------------------------------------------
# TESTS
# ------------------------------------------------------------------------
def test_cpu_slurm_save_load():
"""
Verify model save/load/checkpoint on CPU
:return:
"""
hparams = get_hparams()
model = LightningTestModel(hparams)
save_dir = init_save_dir()
# exp file to get meta
exp = get_exp(False)
exp.argparse(hparams)
exp.save()
cluster_a = SlurmCluster()
trainer_options = dict(
max_nb_epochs=1,
cluster=cluster_a,
experiment=exp,
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 batch in model.test_dataloader:
break
x, y = batch
x = x.view(x.size(0), -1)
model.eval()
pred_before_saving = model(x)
# test registering a save function
trainer.enable_auto_hpc_walltime_manager()
# test HPC saving
# simulate snapshot on slurm
saved_filepath = trainer.hpc_save(save_dir, exp)
assert os.path.exists(saved_filepath)
# 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
continue_tng_hparams = get_hparams(continue_training=True, hpc_exp_number=cluster_a.hpc_exp_number)
trainer_options = dict(
max_nb_epochs=1,
cluster=SlurmCluster(continue_tng_hparams),
experiment=exp,
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)
clear_save_dir()
def test_loading_meta_tags():
hparams = get_hparams()
@@ -43,6 +125,7 @@ def test_loading_meta_tags():
clear_save_dir()
def test_dp_output_reduce():
# test identity when we have a single gpu
@@ -64,9 +147,9 @@ def test_dp_output_reduce():
assert reduced['b']['c'] == out['b']['c']
def test_cpu_slurm_saving_loading():
def test_model_saving_loading():
"""
Verify model save/load/checkpoint on CPU
Tests use case where trainer saves the model, and user loads it from tags independently
:return:
"""
hparams = get_hparams()
@@ -89,43 +172,49 @@ def test_cpu_slurm_saving_loading():
# 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'
# test saving checkpoint
ckpt_test = os.path.join(save_dir, 'test.ckpt')
trainer.save_checkpoint(ckpt_test)
# make a prediction
for batch in model.test_dataloader:
break
# test registering a save function
trainer.enable_auto_hpc_walltime_manager()
x, y = batch
x = x.view(x.size(0), -1)
# test model loading with a map_location
pretrained_model = load_model(exp, save_dir, True)
# generate preds before saving model
model.eval()
pred_before_saving = model(x)
# test model preds
run_prediction(model.test_dataloader, pretrained_model)
# save model
new_weights_path = os.path.join(save_dir, 'save_test.ckpt')
trainer.save_checkpoint(new_weights_path)
trainer.model = pretrained_model
trainer.optimizers = pretrained_model.configure_optimizers()
# load new model
tags_path = exp.get_data_path(exp.name, exp.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, on_gpu=False)
model_2.eval()
# test HPC saving
saved_filepath = trainer.hpc_save(save_dir, exp)
assert os.path.exists(saved_filepath)
# test HPC loading
trainer.global_step = 20000000
trainer.hpc_load(save_dir, on_gpu=False)
assert trainer.global_step == real_global_step and trainer.global_step != 20000000
# test freeze on gpu
model.freeze()
model.unfreeze()
# 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
clear_save_dir()
def test_model_freeze_unfreeze():
hparams = get_hparams()
model = LightningTestModel(hparams)
model.freeze()
model.unfreeze()
def test_amp_gpu_ddp_slurm_managed():
"""
Make sure DDP + AMP work
@@ -494,16 +583,24 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
clear_save_dir()
def get_hparams():
def get_hparams(continue_training=False, hpc_exp_number=0):
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})
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