mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-09-11 12:31:23 +08:00
Add src, docs and other important folders
This commit is contained in:
Shreyas Bapat
@@ -1,40 +0,0 @@
|
||||
import numpy as np
|
||||
from torch import nn
|
||||
|
||||
"""
|
||||
Module to describe gradients
|
||||
"""
|
||||
|
||||
|
||||
class GradInformation(nn.Module):
|
||||
|
||||
def grad_norm(self, norm_type):
|
||||
results = {}
|
||||
total_norm = 0
|
||||
for i, p in enumerate(self.parameters()):
|
||||
if p.requires_grad:
|
||||
try:
|
||||
param_norm = p.grad.data.norm(norm_type)
|
||||
total_norm += param_norm ** norm_type
|
||||
norm = param_norm ** (1 / norm_type)
|
||||
|
||||
results['grad_{}_norm_{}'.format(norm_type, i)] = round(norm.data.cpu().numpy().flatten()[0], 3)
|
||||
except Exception as e:
|
||||
# this param had no grad
|
||||
pass
|
||||
|
||||
total_norm = total_norm ** (1. / norm_type)
|
||||
results['grad_{}_norm_total'.format(norm_type)] = round(total_norm.data.cpu().numpy().flatten()[0], 3)
|
||||
return results
|
||||
|
||||
|
||||
def describe_grads(self):
|
||||
for p in self.parameters():
|
||||
g = p.grad.data.numpy().flatten()
|
||||
print(np.max(g), np.min(g), np.mean(g))
|
||||
|
||||
|
||||
def describe_params(self):
|
||||
for p in self.parameters():
|
||||
g = p.data.numpy().flatten()
|
||||
print(np.max(g), np.min(g), np.mean(g))
|
||||
@@ -1,20 +0,0 @@
|
||||
import torch
|
||||
|
||||
class ModelHooks(torch.nn.Module):
|
||||
def on_batch_start(self):
|
||||
pass
|
||||
|
||||
def on_batch_end(self):
|
||||
pass
|
||||
|
||||
def on_epoch_start(self):
|
||||
pass
|
||||
|
||||
def on_epoch_end(self):
|
||||
pass
|
||||
|
||||
def on_pre_performance_check(self):
|
||||
pass
|
||||
|
||||
def on_post_performance_check(self):
|
||||
pass
|
||||
@@ -1,180 +0,0 @@
|
||||
import torch
|
||||
import gc
|
||||
import subprocess
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
|
||||
'''
|
||||
Generates a summary of a model's layers and dimensionality
|
||||
'''
|
||||
|
||||
|
||||
class ModelSummary(object):
|
||||
|
||||
def __init__(self, model):
|
||||
'''
|
||||
Generates summaries of model layers and dimensions.
|
||||
'''
|
||||
self.model = model
|
||||
self.in_sizes = []
|
||||
self.out_sizes = []
|
||||
|
||||
self.summarize()
|
||||
|
||||
def __str__(self):
|
||||
return self.summary.__str__()
|
||||
|
||||
def __repr__(self):
|
||||
return self.summary.__str__()
|
||||
|
||||
def get_variable_sizes(self):
|
||||
'''Run sample input through each layer to get output sizes'''
|
||||
mods = list(self.model.modules())
|
||||
in_sizes = []
|
||||
out_sizes = []
|
||||
input_ = self.example_input_array
|
||||
for i in range(1, len(mods)):
|
||||
m = mods[i]
|
||||
if type(input_) is list or type(input_) is tuple:
|
||||
out = m(*input_)
|
||||
else:
|
||||
out = m(input_)
|
||||
|
||||
if type(input_) is tuple or type(input_) is list:
|
||||
in_size = []
|
||||
for x in input_:
|
||||
if type(x) is list:
|
||||
in_size.append(len(x))
|
||||
else:
|
||||
in_size.append(x.size())
|
||||
else:
|
||||
in_size = np.array(input_.size())
|
||||
|
||||
in_sizes.append(in_size)
|
||||
|
||||
if type(out) is tuple or type(out) is list:
|
||||
out_size = np.asarray([x.size() for x in out])
|
||||
else:
|
||||
out_size = np.array(out.size())
|
||||
|
||||
out_sizes.append(out_size)
|
||||
input_ = out
|
||||
|
||||
self.in_sizes = in_sizes
|
||||
self.out_sizes = out_sizes
|
||||
return
|
||||
|
||||
def get_layer_names(self):
|
||||
'''Collect Layer Names'''
|
||||
mods = list(self.model.named_modules())
|
||||
names = []
|
||||
layers = []
|
||||
for m in mods[1:]:
|
||||
names += [m[0]]
|
||||
layers += [str(m[1].__class__)]
|
||||
|
||||
layer_types = [x.split('.')[-1][:-2] for x in layers]
|
||||
|
||||
self.layer_names = names
|
||||
self.layer_types = layer_types
|
||||
return
|
||||
|
||||
def get_parameter_sizes(self):
|
||||
'''Get sizes of all parameters in `model`'''
|
||||
mods = list(self.model.modules())
|
||||
sizes = []
|
||||
|
||||
for i in range(1,len(mods)):
|
||||
m = mods[i]
|
||||
p = list(m.parameters())
|
||||
modsz = []
|
||||
for j in range(len(p)):
|
||||
modsz.append(np.array(p[j].size()))
|
||||
sizes.append(modsz)
|
||||
|
||||
self.param_sizes = sizes
|
||||
return
|
||||
|
||||
def get_parameter_nums(self):
|
||||
'''Get number of parameters in each layer'''
|
||||
param_nums = []
|
||||
for mod in self.param_sizes:
|
||||
all_params = 0
|
||||
for p in mod:
|
||||
all_params += np.prod(p)
|
||||
param_nums.append(all_params)
|
||||
self.param_nums = param_nums
|
||||
return
|
||||
|
||||
def make_summary(self):
|
||||
'''
|
||||
Makes a summary listing with:
|
||||
|
||||
Layer Name, Layer Type, Input Size, Output Size, Number of Parameters
|
||||
'''
|
||||
|
||||
df = pd.DataFrame( np.zeros( (len(self.layer_names), 3) ) )
|
||||
df.columns = ['Name', 'Type', 'Params']
|
||||
|
||||
df['Name'] = self.layer_names
|
||||
df['Type'] = self.layer_types
|
||||
df['Params'] = self.param_nums
|
||||
|
||||
self.summary = df
|
||||
return
|
||||
|
||||
def summarize(self):
|
||||
self.get_layer_names()
|
||||
self.get_parameter_sizes()
|
||||
self.get_parameter_nums()
|
||||
self.make_summary()
|
||||
|
||||
|
||||
def print_mem_stack():
|
||||
for obj in gc.get_objects():
|
||||
try:
|
||||
if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):
|
||||
print(type(obj), obj.size())
|
||||
except Exception as e:
|
||||
pass
|
||||
|
||||
|
||||
def count_mem_items():
|
||||
nb_params = 0
|
||||
nb_tensors = 0
|
||||
for obj in gc.get_objects():
|
||||
try:
|
||||
if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):
|
||||
obj_type = str(type(obj))
|
||||
if 'parameter' in obj_type:
|
||||
nb_params += 1
|
||||
else:
|
||||
nb_tensors += 1
|
||||
except Exception as e:
|
||||
pass
|
||||
|
||||
return nb_params, nb_tensors
|
||||
|
||||
|
||||
def get_gpu_memory_map():
|
||||
"""Get the current gpu usage.
|
||||
|
||||
Returns
|
||||
-------
|
||||
usage: dict
|
||||
Keys are device ids as integers.
|
||||
Values are memory usage as integers in MB.
|
||||
"""
|
||||
result = subprocess.check_output(
|
||||
[
|
||||
'nvidia-smi', '--query-gpu=memory.used',
|
||||
'--format=csv,nounits,noheader'
|
||||
], encoding='utf-8')
|
||||
# Convert lines into a dictionary
|
||||
gpu_memory = [int(x) for x in result.strip().split('\n')]
|
||||
gpu_memory_map = {}
|
||||
for k, v in zip(range(len(gpu_memory)), gpu_memory):
|
||||
k = f'gpu_{k}'
|
||||
gpu_memory_map[k] = v
|
||||
return gpu_memory_map
|
||||
@@ -1,168 +0,0 @@
|
||||
import torch
|
||||
import os
|
||||
import re
|
||||
|
||||
|
||||
class ModelIO(object):
|
||||
|
||||
def load_model_specific(self, checkpoint):
|
||||
"""
|
||||
Do something with the checkpoint
|
||||
:param checkpoint:
|
||||
:return:
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_save_dict(self):
|
||||
"""
|
||||
Return specific things for the model
|
||||
:return:
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class TrainerIO(object):
|
||||
|
||||
# --------------------
|
||||
# MODEL SAVE CHECKPOINT
|
||||
# --------------------
|
||||
def save_checkpoint(self, filepath):
|
||||
checkpoint = self.dump_checkpoint()
|
||||
|
||||
# do the actual save
|
||||
torch.save(checkpoint, filepath)
|
||||
|
||||
def dump_checkpoint(self):
|
||||
checkpoint = {
|
||||
'epoch': self.current_epoch,
|
||||
'checkpoint_callback_best': self.checkpoint_callback.best,
|
||||
'early_stop_callback_wait': self.early_stop_callback.wait,
|
||||
'early_stop_callback_patience': self.early_stop_callback.patience,
|
||||
'global_step': self.global_step
|
||||
}
|
||||
|
||||
optimizer_states = []
|
||||
for i, optimizer in enumerate(self.optimizers):
|
||||
optimizer_states.append(optimizer.state_dict())
|
||||
|
||||
checkpoint['optimizer_states'] = optimizer_states
|
||||
|
||||
# request what to save from the model
|
||||
checkpoint_dict = self.model.get_save_dict()
|
||||
|
||||
# merge trainer and model saving items
|
||||
checkpoint.update(checkpoint_dict)
|
||||
return checkpoint
|
||||
|
||||
# --------------------
|
||||
# HPC IO
|
||||
# --------------------
|
||||
def enable_auto_hpc_walltime_manager(self):
|
||||
if self.cluster is None:
|
||||
return
|
||||
|
||||
# allow test tube to handle model check pointing automatically
|
||||
self.cluster.set_checkpoint_save_function(
|
||||
self.hpc_save,
|
||||
kwargs={
|
||||
'folderpath': self.checkpoint_callback.filepath,
|
||||
'experiment': self.experiment
|
||||
}
|
||||
)
|
||||
self.cluster.set_checkpoint_load_function(
|
||||
self.hpc_load,
|
||||
kwargs={
|
||||
'folderpath': self.checkpoint_callback.filepath,
|
||||
'on_gpu': self.on_gpu
|
||||
}
|
||||
)
|
||||
|
||||
def restore_training_state(self, checkpoint):
|
||||
"""
|
||||
Restore trainer state.
|
||||
Model will get its change to update
|
||||
:param checkpoint:
|
||||
:return:
|
||||
"""
|
||||
self.checkpoint_callback.best = checkpoint['checkpoint_callback_best']
|
||||
self.early_stop_callback.wait = checkpoint['early_stop_callback_wait']
|
||||
self.early_stop_callback.patience = checkpoint['early_stop_callback_patience']
|
||||
self.global_step = checkpoint['global_step']
|
||||
|
||||
# restore the optimizers
|
||||
optimizer_states = checkpoint['optimizer_states']
|
||||
for optimizer, opt_state in zip(self.optimizers, optimizer_states):
|
||||
optimizer.load_state_dict(opt_state)
|
||||
|
||||
# ----------------------------------
|
||||
# PRIVATE OPS
|
||||
# ----------------------------------
|
||||
def hpc_save(self, folderpath, experiment):
|
||||
# save exp to make sure we get all the metrics
|
||||
experiment.save()
|
||||
|
||||
ckpt_number = self.max_ckpt_in_folder(folderpath) + 1
|
||||
|
||||
if not os.path.exists(folderpath):
|
||||
os.makedirs(folderpath, exist_ok=True)
|
||||
filepath = '{}/hpc_ckpt_{}.ckpt'.format(folderpath, ckpt_number)
|
||||
|
||||
# request what to save from the model
|
||||
checkpoint_dict = self.dump_checkpoint()
|
||||
|
||||
# do the actual save
|
||||
torch.save(checkpoint_dict, filepath)
|
||||
|
||||
def hpc_load(self, folderpath, on_gpu):
|
||||
filepath = '{}/hpc_ckpt_{}.ckpt'.format(folderpath, self.max_ckpt_in_folder(folderpath))
|
||||
|
||||
if on_gpu:
|
||||
checkpoint = torch.load(filepath)
|
||||
else:
|
||||
checkpoint = torch.load(filepath, map_location=lambda storage, loc: storage)
|
||||
|
||||
# load training state
|
||||
self.restore_training_state(checkpoint)
|
||||
|
||||
# load model state
|
||||
self.model.load_model_specific(checkpoint)
|
||||
|
||||
def max_ckpt_in_folder(self, path):
|
||||
files = os.listdir(path)
|
||||
ckpt_vs = []
|
||||
for name in files:
|
||||
name = name.split('ckpt_')[-1]
|
||||
name = re.sub('[^0-9]', '', name)
|
||||
ckpt_vs.append(int(name))
|
||||
|
||||
return max(ckpt_vs)
|
||||
|
||||
|
||||
def load_hparams_from_tags_csv(tags_csv):
|
||||
from argparse import Namespace
|
||||
import pandas as pd
|
||||
|
||||
tags_df = pd.read_csv(tags_csv)
|
||||
dic = tags_df.to_dict(orient='records')
|
||||
|
||||
ns_dict = {row['key']: convert(row['value']) for row in dic}
|
||||
|
||||
ns = Namespace(**ns_dict)
|
||||
return ns
|
||||
|
||||
|
||||
def convert(val):
|
||||
constructors = [int, float, str]
|
||||
|
||||
if type(val) is str:
|
||||
if val.lower() == 'true':
|
||||
return True
|
||||
if val.lower() == 'false':
|
||||
return False
|
||||
|
||||
for c in constructors:
|
||||
try:
|
||||
return c(val)
|
||||
except ValueError:
|
||||
pass
|
||||
return val
|
||||
@@ -1,22 +0,0 @@
|
||||
from torch import nn
|
||||
from torch import optim
|
||||
|
||||
|
||||
class OptimizerConfig(nn.Module):
|
||||
|
||||
def choose_optimizer(self, optimizer, params, optimizer_params, opt_name_key):
|
||||
if optimizer == 'adam':
|
||||
optimizer = optim.Adam(params, **optimizer_params)
|
||||
if optimizer == 'sparse_adam':
|
||||
optimizer = optim.SparseAdam(params, **optimizer_params)
|
||||
if optimizer == 'sgd':
|
||||
optimizer = optim.SGD(params, **optimizer_params)
|
||||
if optimizer == 'adadelta':
|
||||
optimizer = optim.Adadelta(params, **optimizer_params)
|
||||
|
||||
# transfer opt state if loaded
|
||||
if opt_name_key in self.loaded_optimizer_states_dict:
|
||||
state = self.loaded_optimizer_states_dict[opt_name_key]
|
||||
optimizer.load_state_dict(state)
|
||||
|
||||
return optimizer
|
||||
@@ -1,174 +0,0 @@
|
||||
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 RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
|
||||
|
||||
def __init__(self, hparams):
|
||||
super(RootModule, 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.fast_dev_run = hparams.fast_dev_run
|
||||
self.overfit = hparams.overfit
|
||||
self.gradient_clip = hparams.gradient_clip
|
||||
self.num = 2
|
||||
|
||||
# track if gpu was requested for checkpointing
|
||||
self.on_gpu = False
|
||||
try:
|
||||
self.on_gpu = hparams.on_gpu
|
||||
except Exception as e:
|
||||
pass
|
||||
|
||||
# computed vars for the dataloaders
|
||||
self._tng_dataloader = None
|
||||
self._val_dataloader = None
|
||||
self._test_dataloader = None
|
||||
|
||||
if self.on_gpu:
|
||||
print('running on gpu...')
|
||||
self.dtype = torch.cuda.FloatTensor
|
||||
torch.set_default_tensor_type('torch.cuda.FloatTensor')
|
||||
|
||||
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):
|
||||
"""
|
||||
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):
|
||||
"""
|
||||
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:
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
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 nb_batches(self, dataloader):
|
||||
a = math.ceil(float(len(dataloader.dataset) / self.batch_size))
|
||||
return int(a)
|
||||
|
||||
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):
|
||||
"""
|
||||
Primary way of loading model from csv weights path
|
||||
:param weights_path:
|
||||
:param tags_csv:
|
||||
:param on_gpu:
|
||||
:return:
|
||||
"""
|
||||
hparams = load_hparams_from_tags_csv(tags_csv)
|
||||
hparams.__setattr__('on_gpu', on_gpu)
|
||||
|
||||
if on_gpu:
|
||||
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
|
||||
Reference in New Issue
Block a user