165 lines
4.6 KiB
Python
165 lines
4.6 KiB
Python
import os
|
|
import torch
|
|
import math
|
|
|
|
from pytorch_lightning.root_module.memory import ModelSummary
|
|
from pytorch_lightning.root_module.grads import GradInformation
|
|
from pytorch_lightning.root_module.model_saving import ModelIO, load_hparams_from_tags_csv
|
|
from pytorch_lightning.root_module.optimization import OptimizerConfig
|
|
from pytorch_lightning.root_module.hooks import ModelHooks
|
|
|
|
|
|
class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
|
|
|
|
def __init__(self, hparams):
|
|
super(LightningModule, self).__init__()
|
|
self.hparams = hparams
|
|
|
|
self.dtype = torch.FloatTensor
|
|
self.exp_save_path = None
|
|
self.current_epoch = 0
|
|
self.global_step = 0
|
|
self.loaded_optimizer_states_dict = {}
|
|
self.trainer = None
|
|
self.experiment = None
|
|
|
|
# track if gpu was requested for checkpointing
|
|
self.on_gpu = False
|
|
|
|
# computed vars for the dataloaders
|
|
self._tng_dataloader = None
|
|
self._val_dataloader = None
|
|
self._test_dataloader = None
|
|
|
|
def forward(self, *args, **kwargs):
|
|
"""
|
|
Expand model in into whatever you need.
|
|
Also need to return the target
|
|
:param x:
|
|
:return:
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def validation_step(self, data_batch, batch_nb):
|
|
"""
|
|
return whatever outputs will need to be aggregated in validation_end
|
|
:param data_batch:
|
|
:return:
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def validation_end(self, outputs):
|
|
"""
|
|
Outputs has the appended output after each validation step
|
|
:param outputs:
|
|
:return: dic_with_metrics for tqdm
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def training_step(self, data_batch, batch_nb):
|
|
"""
|
|
return loss, dict with metrics for tqdm
|
|
:param data_batch:
|
|
:return:
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def configure_optimizers(self):
|
|
"""
|
|
Return array of optimizers
|
|
:return:
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def update_tng_log_metrics(self, logs):
|
|
"""
|
|
Chance to update metrics to be logged for training step.
|
|
For example, add music, images, etc... to log
|
|
:param logs:
|
|
:return:
|
|
"""
|
|
return logs
|
|
|
|
def loss(self, *args, **kwargs):
|
|
"""
|
|
Expand model_out into your components
|
|
:param model_out:
|
|
:return:
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def summarize(self):
|
|
model_summary = ModelSummary(self)
|
|
print(model_summary)
|
|
|
|
|
|
def freeze(self):
|
|
for param in self.parameters():
|
|
param.requires_grad = False
|
|
|
|
def unfreeze(self):
|
|
for param in self.parameters():
|
|
param.requires_grad = True
|
|
|
|
@property
|
|
def tng_dataloader(self):
|
|
"""
|
|
Implement a function to load an h5py of this data
|
|
:return:
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
@property
|
|
def test_dataloader(self):
|
|
"""
|
|
Implement a function to load an h5py of this data
|
|
:return:
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
@property
|
|
def val_dataloader(self):
|
|
"""
|
|
Implement a function to load an h5py of this data
|
|
:return:
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
@staticmethod
|
|
def get_process_position(gpus):
|
|
try:
|
|
current_gpu = os.environ["CUDA_VISIBLE_DEVICES"]
|
|
gpu_ids = gpus.split(',')
|
|
process_position = gpu_ids.index(current_gpu)
|
|
return process_position, current_gpu
|
|
except Exception as e:
|
|
return 0, 0
|
|
|
|
@classmethod
|
|
def load_from_metrics(cls, weights_path, tags_csv, on_gpu, map_location=None):
|
|
"""
|
|
Primary way of loading model from csv weights path
|
|
:param weights_path:
|
|
:param tags_csv:
|
|
:param on_gpu:
|
|
:param map_location: dic for mapping storage {'cuda:1':'cuda:0'}
|
|
:return:
|
|
"""
|
|
hparams = load_hparams_from_tags_csv(tags_csv)
|
|
hparams.__setattr__('on_gpu', on_gpu)
|
|
|
|
if on_gpu:
|
|
if map_location is not None:
|
|
checkpoint = torch.load(weights_path, map_location=map_location)
|
|
else:
|
|
checkpoint = torch.load(weights_path)
|
|
else:
|
|
checkpoint = torch.load(weights_path, map_location=lambda storage, loc: storage)
|
|
|
|
model = cls(hparams)
|
|
|
|
# allow model to load
|
|
model.load_model_specific(checkpoint)
|
|
model.load_state_dict(checkpoint['state_dict'], strict=False)
|
|
return model
|