mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-09-12 12:40:20 +08:00
Refactor: name modules (#548)
* refactor: rename some modules * add deprecation warnings * fix paths
This commit is contained in:
Jirka Borovec
committed by
William Falcon
William Falcon
parent
fea7cc87f6
commit
9785a3e78e
@@ -0,0 +1,7 @@
|
||||
import warnings
|
||||
|
||||
warnings.warn("`root_module` package has been renamed to `core` since v0.5.3"
|
||||
" and will be removed in v0.8.0", DeprecationWarning)
|
||||
|
||||
from pytorch_lightning.core import ( # noqa: E402
|
||||
decorators, grads, hooks, root_module, memory, model_saving)
|
||||
|
||||
@@ -1,33 +0,0 @@
|
||||
import traceback
|
||||
|
||||
|
||||
def data_loader(fn):
|
||||
"""
|
||||
Decorator to make any fx with this use the lazy property
|
||||
:param fn:
|
||||
:return:
|
||||
"""
|
||||
|
||||
attr_name = '_lazy_' + fn.__name__
|
||||
|
||||
def _get_data_loader(self):
|
||||
try:
|
||||
value = getattr(self, attr_name)
|
||||
except AttributeError:
|
||||
try:
|
||||
value = fn(self) # Lazy evaluation, done only once.
|
||||
if (
|
||||
value is not None and
|
||||
not isinstance(value, list) and
|
||||
fn.__name__ in ['test_dataloader', 'val_dataloader']
|
||||
):
|
||||
value = [value]
|
||||
except AttributeError as e:
|
||||
# Guard against AttributeError suppression. (Issue #142)
|
||||
traceback.print_exc()
|
||||
error = f'{fn.__name__}: An AttributeError was encountered: ' + str(e)
|
||||
raise RuntimeError(error) from e
|
||||
setattr(self, attr_name, value) # Memoize evaluation.
|
||||
return value
|
||||
|
||||
return _get_data_loader
|
||||
@@ -1,29 +0,0 @@
|
||||
"""
|
||||
Module to describe gradients
|
||||
"""
|
||||
|
||||
from torch import nn
|
||||
|
||||
|
||||
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)
|
||||
|
||||
grad = round(norm.data.cpu().numpy().flatten()[0], 3)
|
||||
results['grad_{}_norm_{}'.format(norm_type, i)] = grad
|
||||
except Exception:
|
||||
# this param had no grad
|
||||
pass
|
||||
|
||||
total_norm = total_norm ** (1. / norm_type)
|
||||
grad = round(total_norm.data.cpu().numpy().flatten()[0], 3)
|
||||
results['grad_{}_norm_total'.format(norm_type)] = grad
|
||||
return results
|
||||
@@ -1,72 +0,0 @@
|
||||
import torch
|
||||
|
||||
|
||||
try:
|
||||
from apex import amp
|
||||
|
||||
APEX_AVAILABLE = True
|
||||
except ImportError:
|
||||
APEX_AVAILABLE = False
|
||||
|
||||
|
||||
class ModelHooks(torch.nn.Module):
|
||||
|
||||
def on_sanity_check_start(self):
|
||||
"""
|
||||
Called before starting evaluate
|
||||
:return:
|
||||
"""
|
||||
pass
|
||||
|
||||
def on_batch_start(self, batch):
|
||||
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
|
||||
|
||||
def on_before_zero_grad(self, optimizer):
|
||||
"""
|
||||
Called after optimizer.step() and before optimizer.zero_grad()
|
||||
|
||||
for optimizer in optimizers:
|
||||
optimizer.step()
|
||||
model.on_before_zero_grad(optimizer) # < ---- called here
|
||||
optimizer.zero_grad
|
||||
|
||||
:param optimizer:
|
||||
:return:
|
||||
"""
|
||||
pass
|
||||
|
||||
def on_after_backward(self):
|
||||
"""
|
||||
Called after loss.backward() and before optimizers do anything
|
||||
:return:
|
||||
"""
|
||||
pass
|
||||
|
||||
def backward(self, use_amp, loss, optimizer):
|
||||
"""
|
||||
Override backward with your own implementation if you need to
|
||||
:param use_amp: Whether amp was requested or not
|
||||
:param loss: Loss is already scaled by accumulated grads
|
||||
:param optimizer: Current optimizer being used
|
||||
:return:
|
||||
"""
|
||||
if use_amp:
|
||||
with amp.scale_loss(loss, optimizer) as scaled_loss:
|
||||
scaled_loss.backward()
|
||||
else:
|
||||
loss.backward()
|
||||
@@ -1,257 +0,0 @@
|
||||
'''
|
||||
Generates a summary of a model's layers and dimensionality
|
||||
'''
|
||||
|
||||
import gc
|
||||
import os
|
||||
import subprocess
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import torch
|
||||
import logging
|
||||
|
||||
|
||||
class ModelSummary(object):
|
||||
|
||||
def __init__(self, model, mode='full'):
|
||||
'''
|
||||
Generates summaries of model layers and dimensions.
|
||||
'''
|
||||
self.model = model
|
||||
self.mode = mode
|
||||
self.in_sizes = []
|
||||
self.out_sizes = []
|
||||
|
||||
self.summarize()
|
||||
|
||||
def __str__(self):
|
||||
return self.summary.__str__()
|
||||
|
||||
def __repr__(self):
|
||||
return self.summary.__str__()
|
||||
|
||||
def named_modules(self):
|
||||
if self.mode == 'full':
|
||||
mods = self.model.named_modules()
|
||||
mods = list(mods)[1:] # do not include root module (LightningModule)
|
||||
elif self.mode == 'top':
|
||||
# the children are the top-level modules
|
||||
mods = self.model.named_children()
|
||||
else:
|
||||
mods = []
|
||||
return list(mods)
|
||||
|
||||
def get_variable_sizes(self):
|
||||
'''Run sample input through each layer to get output sizes'''
|
||||
mods = self.named_modules()
|
||||
in_sizes = []
|
||||
out_sizes = []
|
||||
input_ = self.model.example_input_array
|
||||
|
||||
if self.model.on_gpu:
|
||||
input_ = input_.cuda(0)
|
||||
|
||||
if self.model.trainer.use_amp:
|
||||
input_ = input_.half()
|
||||
|
||||
with torch.no_grad():
|
||||
|
||||
for _, m in mods:
|
||||
if type(input_) is list or type(input_) is tuple: # pragma: no cover
|
||||
out = m(*input_)
|
||||
else:
|
||||
out = m(input_)
|
||||
|
||||
if type(input_) is tuple or type(input_) is list: # pragma: no cover
|
||||
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: # pragma: no cover
|
||||
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
|
||||
assert len(in_sizes) == len(out_sizes)
|
||||
return
|
||||
|
||||
def get_layer_names(self):
|
||||
'''Collect Layer Names'''
|
||||
mods = self.named_modules()
|
||||
names = []
|
||||
layers = []
|
||||
for name, m in mods:
|
||||
names += [name]
|
||||
layers += [str(m.__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 = self.named_modules()
|
||||
sizes = []
|
||||
for _, m in mods:
|
||||
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
|
||||
'''
|
||||
|
||||
cols = ['Name', 'Type', 'Params']
|
||||
if self.model.example_input_array is not None:
|
||||
cols.extend(['In_sizes', 'Out_sizes'])
|
||||
|
||||
df = pd.DataFrame(np.zeros((len(self.layer_names), len(cols))))
|
||||
df.columns = cols
|
||||
|
||||
df['Name'] = self.layer_names
|
||||
df['Type'] = self.layer_types
|
||||
df['Params'] = self.param_nums
|
||||
df['Params'] = df['Params'].map(get_human_readable_count)
|
||||
|
||||
if self.model.example_input_array is not None:
|
||||
df['In_sizes'] = self.in_sizes
|
||||
df['Out_sizes'] = self.out_sizes
|
||||
|
||||
self.summary = df
|
||||
return
|
||||
|
||||
def summarize(self):
|
||||
self.get_layer_names()
|
||||
self.get_parameter_sizes()
|
||||
self.get_parameter_nums()
|
||||
|
||||
if self.model.example_input_array is not None:
|
||||
self.get_variable_sizes()
|
||||
self.make_summary()
|
||||
|
||||
|
||||
def print_mem_stack(): # pragma: no cover
|
||||
for obj in gc.get_objects():
|
||||
try:
|
||||
if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):
|
||||
logging.info(type(obj), obj.size())
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def count_mem_items(): # pragma: no cover
|
||||
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:
|
||||
pass
|
||||
|
||||
return nb_params, nb_tensors
|
||||
|
||||
|
||||
def get_memory_profile(mode):
|
||||
"""
|
||||
'all' means return memory for all gpus
|
||||
'min_max' means return memory for max and min
|
||||
:param mode:
|
||||
:return:
|
||||
"""
|
||||
memory_map = get_gpu_memory_map()
|
||||
|
||||
if mode == 'min_max':
|
||||
min_index, min_memory = min(memory_map.items(), key=lambda item: item[1])
|
||||
max_index, max_memory = max(memory_map.items(), key=lambda item: item[1])
|
||||
|
||||
memory_map = {min_index: min_memory, max_index: max_memory}
|
||||
|
||||
return memory_map
|
||||
|
||||
|
||||
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.run(
|
||||
[
|
||||
'nvidia-smi',
|
||||
'--query-gpu=memory.used',
|
||||
'--format=csv,nounits,noheader',
|
||||
],
|
||||
encoding='utf-8',
|
||||
capture_output=True,
|
||||
check=True)
|
||||
# Convert lines into a dictionary
|
||||
gpu_memory = [int(x) for x in result.stdout.strip().split(os.linesep)]
|
||||
gpu_memory_map = {f'gpu_{index}': memory for index, memory in enumerate(gpu_memory)}
|
||||
return gpu_memory_map
|
||||
|
||||
|
||||
def get_human_readable_count(number):
|
||||
"""
|
||||
Abbreviates an integer number with K, M, B, T for thousands, millions,
|
||||
billions and trillions, respectively.
|
||||
Examples:
|
||||
123 -> 123
|
||||
1234 -> 1 K (one thousand)
|
||||
2e6 -> 2 M (two million)
|
||||
3e9 -> 3 B (three billion)
|
||||
4e12 -> 4 T (four trillion)
|
||||
5e15 -> 5,000 T
|
||||
:param number: a positive integer number
|
||||
:returns a string formatted according to the pattern described above.
|
||||
"""
|
||||
assert number >= 0
|
||||
labels = [' ', 'K', 'M', 'B', 'T']
|
||||
num_digits = int(np.floor(np.log10(number)) + 1 if number > 0 else 1)
|
||||
num_groups = int(np.ceil(num_digits / 3))
|
||||
num_groups = min(num_groups, len(labels)) # don't abbreviate beyond trillions
|
||||
shift = -3 * (num_groups - 1)
|
||||
number = number * (10 ** shift)
|
||||
index = num_groups - 1
|
||||
return f'{int(number):,d} {labels[index]}'
|
||||
@@ -1,34 +0,0 @@
|
||||
class ModelIO(object):
|
||||
|
||||
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:
|
||||
"""
|
||||
pass
|
||||
|
||||
def on_save_checkpoint(self, checkpoint):
|
||||
"""
|
||||
Give the model a chance to add something to the checkpoint.
|
||||
state_dict is already there
|
||||
"""
|
||||
pass
|
||||
|
||||
# -------------------------
|
||||
# OPTIONAL HOOKS
|
||||
# -------------------------
|
||||
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, checkpoint):
|
||||
"""
|
||||
Hook to do whatever you need right before Slurm manager loads the model
|
||||
:return:
|
||||
"""
|
||||
pass
|
||||
@@ -1,334 +0,0 @@
|
||||
import os
|
||||
import warnings
|
||||
import collections
|
||||
from argparse import Namespace
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from pytorch_lightning.root_module.decorators import data_loader
|
||||
from pytorch_lightning.root_module.grads import GradInformation
|
||||
from pytorch_lightning.root_module.hooks import ModelHooks
|
||||
from pytorch_lightning.root_module.memory import ModelSummary
|
||||
from pytorch_lightning.root_module.model_saving import ModelIO
|
||||
from pytorch_lightning.trainer.trainer_io import load_hparams_from_tags_csv
|
||||
import logging
|
||||
from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel
|
||||
|
||||
|
||||
class LightningModule(GradInformation, ModelIO, ModelHooks):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super(LightningModule, self).__init__(*args, **kwargs)
|
||||
|
||||
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.logger = None
|
||||
self.example_input_array = None
|
||||
|
||||
# track if gpu was requested for checkpointing
|
||||
self.on_gpu = False
|
||||
self.use_dp = False
|
||||
self.use_ddp = False
|
||||
self.use_ddp2 = False
|
||||
self.use_amp = False
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
"""
|
||||
Expand model in into whatever you need.
|
||||
Also need to return the target
|
||||
:param x:
|
||||
:return:
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def training_step(self, *args, **kwargs):
|
||||
"""
|
||||
return loss, dict with metrics for tqdm
|
||||
:param called with batch, batch_nb
|
||||
additional: optimizer_i if multiple optimizers used
|
||||
:return: dict with loss key and optional log, progress keys
|
||||
if implementing training_step, return whatever you need in that step
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def training_end(self, *args, **kwargs):
|
||||
"""
|
||||
return loss, dict with metrics for tqdm
|
||||
:param called with outputs of training_step
|
||||
:return: dict with loss key and optional log, progress keys
|
||||
"""
|
||||
pass
|
||||
|
||||
def validation_step(self, *args, **kwargs):
|
||||
"""
|
||||
return whatever outputs will need to be aggregated in validation_end
|
||||
OPTIONAL
|
||||
:param called with batch, batch_nb
|
||||
additional: dataset_i if multiple val datasets used
|
||||
:return:
|
||||
"""
|
||||
pass
|
||||
|
||||
def test_step(self, *args, **kwargs):
|
||||
"""
|
||||
return whatever outputs will need to be aggregated in test_end
|
||||
OPTIONAL
|
||||
:param called with batch, batch_nb
|
||||
additional: dataset_i if multiple val datasets used
|
||||
:return:
|
||||
"""
|
||||
pass
|
||||
|
||||
def validation_end(self, outputs):
|
||||
"""
|
||||
Outputs has the appended output after each validation step
|
||||
OPTIONAL
|
||||
:param outputs:
|
||||
:return: dic_with_metrics for tqdm
|
||||
"""
|
||||
pass
|
||||
|
||||
def test_end(self, outputs):
|
||||
"""
|
||||
Outputs has the appended output after each test step
|
||||
OPTIONAL
|
||||
:param outputs:
|
||||
:return: dic_with_metrics for tqdm
|
||||
"""
|
||||
pass
|
||||
|
||||
def configure_ddp(self, model, device_ids):
|
||||
"""
|
||||
Override to init DDP in a different way or use your own wrapper.
|
||||
Must return model.
|
||||
:param model:
|
||||
:param device_ids:
|
||||
:return: DDP wrapped model
|
||||
"""
|
||||
model = LightningDistributedDataParallel(
|
||||
model,
|
||||
device_ids=device_ids,
|
||||
find_unused_parameters=True
|
||||
)
|
||||
return model
|
||||
|
||||
def init_ddp_connection(self, proc_rank, world_size):
|
||||
"""
|
||||
Connect all procs in the world using the env:// init
|
||||
Use the first node as the root address
|
||||
"""
|
||||
|
||||
# use slurm job id for the port number
|
||||
# guarantees unique ports across jobs from same grid search
|
||||
try:
|
||||
# use the last 4 numbers in the job id as the id
|
||||
default_port = os.environ['SLURM_JOB_ID']
|
||||
default_port = default_port[-4:]
|
||||
|
||||
# all ports should be in the 10k+ range
|
||||
default_port = int(default_port) + 15000
|
||||
|
||||
except Exception as e:
|
||||
default_port = 12910
|
||||
|
||||
# if user gave a port number, use that one instead
|
||||
try:
|
||||
default_port = os.environ['MASTER_PORT']
|
||||
except Exception:
|
||||
os.environ['MASTER_PORT'] = str(default_port)
|
||||
|
||||
# figure out the root node addr
|
||||
try:
|
||||
root_node = os.environ['SLURM_NODELIST'].split(' ')[0]
|
||||
except Exception:
|
||||
root_node = '127.0.0.2'
|
||||
|
||||
root_node = self.trainer.resolve_root_node_address(root_node)
|
||||
os.environ['MASTER_ADDR'] = root_node
|
||||
dist.init_process_group('nccl', rank=proc_rank, world_size=world_size)
|
||||
|
||||
def configure_apex(self, amp, model, optimizers, amp_level):
|
||||
"""
|
||||
Override to init AMP your own way
|
||||
Must return a model and list of optimizers
|
||||
:param amp:
|
||||
:param model:
|
||||
:param optimizers:
|
||||
:param amp_level:
|
||||
:return: Apex wrapped model and optimizers
|
||||
"""
|
||||
model, optimizers = amp.initialize(
|
||||
model, optimizers, opt_level=amp_level,
|
||||
)
|
||||
|
||||
return model, optimizers
|
||||
|
||||
def configure_optimizers(self):
|
||||
"""
|
||||
Return a list of optimizers and a list of schedulers (could be empty)
|
||||
:return:
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def optimizer_step(self, epoch_nb, batch_nb, optimizer, optimizer_i, second_order_closure=None):
|
||||
"""
|
||||
Do something instead of the standard optimizer behavior
|
||||
:param epoch_nb:
|
||||
:param batch_nb:
|
||||
:param optimizer:
|
||||
:param optimizer_i:
|
||||
:param second_order_closure: closure for second order methods
|
||||
:return:
|
||||
"""
|
||||
if isinstance(optimizer, torch.optim.LBFGS):
|
||||
optimizer.step(second_order_closure)
|
||||
else:
|
||||
optimizer.step()
|
||||
|
||||
# clear gradients
|
||||
optimizer.zero_grad()
|
||||
|
||||
def tbptt_split_batch(self, batch, split_size):
|
||||
"""
|
||||
Return list of batch splits. Each split will be passed to forward_step to enable truncated
|
||||
back propagation through time. The default implementation splits root level Tensors and
|
||||
Sequences at dim=1 (i.e. time dim). It assumes that each time dim is the same length.
|
||||
:return:
|
||||
"""
|
||||
time_dims = [len(x[0]) for x in batch if isinstance(
|
||||
x, torch.Tensor) or isinstance(x, collections.Sequence)]
|
||||
assert len(time_dims) >= 1, "Unable to determine batch time dimension"
|
||||
assert all(x == time_dims[0] for x in time_dims), "Batch time dimension length is ambiguous"
|
||||
|
||||
splits = []
|
||||
for t in range(0, time_dims[0], split_size):
|
||||
batch_split = []
|
||||
for i, x in enumerate(batch):
|
||||
if isinstance(x, torch.Tensor):
|
||||
split_x = x[:, t:t + split_size]
|
||||
elif isinstance(x, collections.Sequence):
|
||||
split_x = [None] * len(x)
|
||||
for batch_idx in range(len(x)):
|
||||
split_x[batch_idx] = x[batch_idx][t:t + split_size]
|
||||
|
||||
batch_split.append(split_x)
|
||||
|
||||
splits.append(batch_split)
|
||||
|
||||
return splits
|
||||
|
||||
@data_loader
|
||||
def tng_dataloader(self):
|
||||
"""
|
||||
Implement a PyTorch DataLoader
|
||||
* Deprecated in v0.5.0. use train_dataloader instead. *
|
||||
:return:
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@data_loader
|
||||
def train_dataloader(self):
|
||||
"""
|
||||
Implement a PyTorch DataLoader
|
||||
:return:
|
||||
"""
|
||||
#
|
||||
try:
|
||||
output = self.tng_dataloader()
|
||||
warnings.warn("tng_dataloader has been renamed to train_dataloader since v0.5.0",
|
||||
DeprecationWarning)
|
||||
return output
|
||||
except NotImplementedError:
|
||||
raise NotImplementedError
|
||||
|
||||
@data_loader
|
||||
def test_dataloader(self):
|
||||
"""
|
||||
Implement a PyTorch DataLoader
|
||||
:return:
|
||||
"""
|
||||
return None
|
||||
|
||||
@data_loader
|
||||
def val_dataloader(self):
|
||||
"""
|
||||
Implement a PyTorch DataLoader
|
||||
:return:
|
||||
"""
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def load_from_metrics(cls, weights_path, tags_csv):
|
||||
"""
|
||||
Primary way of loading model from csv weights path
|
||||
:param weights_path:
|
||||
:param tags_csv:
|
||||
:param map_location: dic for mapping storage {'cuda:1':'cuda:0'}
|
||||
:return:
|
||||
"""
|
||||
hparams = load_hparams_from_tags_csv(tags_csv)
|
||||
hparams.__setattr__('on_gpu', False)
|
||||
|
||||
# load on CPU only to avoid OOM issues
|
||||
# then its up to user to put back on GPUs
|
||||
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)
|
||||
|
||||
return model
|
||||
|
||||
@classmethod
|
||||
def load_from_checkpoint(cls, checkpoint_path):
|
||||
"""
|
||||
Primary way of loading model from a checkpoint
|
||||
:param checkpoint_path:
|
||||
:param map_location: dic for mapping storage {'cuda:1':'cuda:0'}
|
||||
:return:
|
||||
"""
|
||||
|
||||
# load on CPU only to avoid OOM issues
|
||||
# then its up to user to put back on GPUs
|
||||
checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)
|
||||
try:
|
||||
ckpt_hparams = checkpoint['hparams']
|
||||
except KeyError:
|
||||
raise IOError(
|
||||
"Checkpoint does not contain hyperparameters. Are your model hyperparameters stored"
|
||||
"in self.hparams?"
|
||||
)
|
||||
hparams = Namespace(**ckpt_hparams)
|
||||
|
||||
# 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)
|
||||
|
||||
return model
|
||||
|
||||
def summarize(self, mode):
|
||||
model_summary = ModelSummary(self, mode=mode)
|
||||
logging.info(model_summary)
|
||||
|
||||
def freeze(self):
|
||||
for param in self.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
self.eval()
|
||||
|
||||
def unfreeze(self):
|
||||
for param in self.parameters():
|
||||
param.requires_grad = True
|
||||
|
||||
self.train()
|
||||
Reference in New Issue
Block a user