mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-09-10 12:21:57 +08:00
Refactor (#407)
* moved dp, ddp outside of trainer * added main mixins * finished major mixin refactor * flake8 * finished major mixin refactor * finished major mixin refactor * finished major mixin refactor * finished major mixin refactor * finished major mixin refactor * finished major mixin refactor * finished major mixin refactor
This commit is contained in:
William Falcon
@@ -0,0 +1,23 @@
|
||||
try:
|
||||
from apex import amp
|
||||
APEX_AVAILABLE = True
|
||||
except ImportError:
|
||||
APEX_AVAILABLE = False
|
||||
|
||||
|
||||
class TrainerAMPMixin(object):
|
||||
|
||||
def init_amp(self, use_amp):
|
||||
self.use_amp = use_amp and APEX_AVAILABLE
|
||||
if self.use_amp:
|
||||
print('using 16bit precision')
|
||||
|
||||
if use_amp and not APEX_AVAILABLE: # pragma: no cover
|
||||
msg = """
|
||||
You set use_amp=True but do not have apex installed.
|
||||
Install apex first using this guide and rerun with use_amp=True:
|
||||
https://github.com/NVIDIA/apex#linux
|
||||
|
||||
this run will NOT use 16 bit precision
|
||||
"""
|
||||
raise ModuleNotFoundError(msg)
|
||||
@@ -0,0 +1,70 @@
|
||||
from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
|
||||
from pytorch_lightning.logging import TestTubeLogger
|
||||
|
||||
|
||||
class TrainerCallbackConfigMixin(object):
|
||||
def configure_checkpoint_callback(self):
|
||||
"""
|
||||
Weight path set in this priority:
|
||||
Checkpoint_callback's path (if passed in).
|
||||
User provided weights_saved_path
|
||||
Otherwise use os.getcwd()
|
||||
"""
|
||||
if self.checkpoint_callback is True:
|
||||
# init a default one
|
||||
if isinstance(self.logger, TestTubeLogger):
|
||||
ckpt_path = '{}/{}/version_{}/{}'.format(
|
||||
self.default_save_path,
|
||||
self.logger.experiment.name,
|
||||
self.logger.experiment.version,
|
||||
'checkpoints')
|
||||
else:
|
||||
ckpt_path = self.default_save_path
|
||||
|
||||
self.checkpoint_callback = ModelCheckpoint(
|
||||
filepath=ckpt_path
|
||||
)
|
||||
elif self.checkpoint_callback is False:
|
||||
self.checkpoint_callback = None
|
||||
|
||||
if self.checkpoint_callback:
|
||||
# set the path for the callbacks
|
||||
self.checkpoint_callback.save_function = self.save_checkpoint
|
||||
|
||||
# if checkpoint callback used, then override the weights path
|
||||
self.weights_save_path = self.checkpoint_callback.filepath
|
||||
|
||||
# if weights_save_path is still none here, set to current working dir
|
||||
if self.weights_save_path is None:
|
||||
self.weights_save_path = self.default_save_path
|
||||
|
||||
def configure_early_stopping(self, early_stop_callback, logger):
|
||||
if early_stop_callback is True:
|
||||
self.early_stop_callback = EarlyStopping(
|
||||
monitor='val_loss',
|
||||
patience=3,
|
||||
verbose=True,
|
||||
mode='min'
|
||||
)
|
||||
self.enable_early_stop = True
|
||||
elif not early_stop_callback:
|
||||
self.early_stop_callback = None
|
||||
self.enable_early_stop = False
|
||||
else:
|
||||
self.early_stop_callback = early_stop_callback
|
||||
self.enable_early_stop = True
|
||||
|
||||
# configure logger
|
||||
if logger is True:
|
||||
# default logger
|
||||
self.logger = TestTubeLogger(
|
||||
save_dir=self.default_save_path,
|
||||
version=self.slurm_job_id,
|
||||
name='lightning_logs'
|
||||
)
|
||||
self.logger.rank = 0
|
||||
elif logger is False:
|
||||
self.logger = None
|
||||
else:
|
||||
self.logger = logger
|
||||
self.logger.rank = 0
|
||||
@@ -0,0 +1,141 @@
|
||||
import warnings
|
||||
|
||||
from torch.utils.data.distributed import DistributedSampler
|
||||
import torch.distributed as dist
|
||||
|
||||
try:
|
||||
from apex import amp
|
||||
APEX_AVAILABLE = True
|
||||
except ImportError:
|
||||
APEX_AVAILABLE = False
|
||||
|
||||
|
||||
class TrainerDataLoadingMixin(object):
|
||||
|
||||
def layout_bookeeping(self):
|
||||
|
||||
# determine number of training batches
|
||||
self.nb_training_batches = len(self.get_train_dataloader())
|
||||
self.nb_training_batches = int(self.nb_training_batches * self.train_percent_check)
|
||||
|
||||
# determine number of validation batches
|
||||
# val datasets could be none, 1 or 2+
|
||||
if self.get_val_dataloaders() is not None:
|
||||
self.nb_val_batches = sum(len(dataloader) for dataloader in self.get_val_dataloaders())
|
||||
self.nb_val_batches = int(self.nb_val_batches * self.val_percent_check)
|
||||
self.nb_val_batches = max(1, self.nb_val_batches)
|
||||
|
||||
# determine number of test batches
|
||||
if self.get_test_dataloaders() is not None:
|
||||
self.nb_test_batches = sum(
|
||||
len(dataloader) for dataloader in self.get_test_dataloaders()
|
||||
)
|
||||
self.nb_test_batches = int(self.nb_test_batches * self.test_percent_check)
|
||||
self.nb_test_batches = max(1, self.nb_test_batches)
|
||||
|
||||
# determine when to check validation
|
||||
self.val_check_batch = int(self.nb_training_batches * self.val_check_interval)
|
||||
self.val_check_batch = max(1, self.val_check_batch)
|
||||
|
||||
def get_dataloaders(self, model):
|
||||
"""
|
||||
Dataloaders are provided by the model
|
||||
:param model:
|
||||
:return:
|
||||
"""
|
||||
self.get_train_dataloader = model.train_dataloader
|
||||
self.get_test_dataloaders = model.test_dataloader
|
||||
self.get_val_dataloaders = model.val_dataloader
|
||||
|
||||
# call warnings from proc zero only which triggers dataloaders
|
||||
# if those have to download data it will only happen on proc 0
|
||||
if self.proc_rank == 0:
|
||||
on_ddp = self.use_ddp or self.use_ddp2
|
||||
if on_ddp and not isinstance(self.get_train_dataloader().sampler, DistributedSampler):
|
||||
msg = """
|
||||
You're using multiple gpus and multiple nodes without using a DistributedSampler
|
||||
to assign a subset of your data to each process. To silence this warning, pass a
|
||||
DistributedSampler to your DataLoader.
|
||||
|
||||
ie: this:
|
||||
dataset = myDataset()
|
||||
dataloader = Dataloader(dataset)
|
||||
|
||||
becomes:
|
||||
dataset = myDataset()
|
||||
dist_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
|
||||
dataloader = Dataloader(dataset, sampler=dist_sampler)
|
||||
|
||||
If you want each process to load the full dataset, ignore this warning.
|
||||
"""
|
||||
warnings.warn(msg)
|
||||
|
||||
if on_ddp and self.get_val_dataloaders() is not None:
|
||||
for dataloader in self.get_val_dataloaders():
|
||||
if not isinstance(dataloader.sampler, DistributedSampler):
|
||||
msg = """
|
||||
Your val_dataloader(s) don't use DistributedSampler.
|
||||
|
||||
You're using multiple gpus and multiple nodes without using a
|
||||
DistributedSampler to assign a subset of your data to each process.
|
||||
To silence this warning, pass a DistributedSampler to your DataLoader.
|
||||
|
||||
ie: this:
|
||||
dataset = myDataset()
|
||||
dataloader = Dataloader(dataset)
|
||||
|
||||
becomes:
|
||||
dataset = myDataset()
|
||||
dist_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
|
||||
dataloader = Dataloader(dataset, sampler=dist_sampler)
|
||||
|
||||
If you want each process to load the full dataset, ignore this warning.
|
||||
"""
|
||||
warnings.warn(msg)
|
||||
break
|
||||
|
||||
if on_ddp and self.get_test_dataloaders() is not None:
|
||||
for dataloader in self.get_test_dataloaders():
|
||||
if not isinstance(dataloader.sampler, DistributedSampler):
|
||||
msg = """
|
||||
Your test_dataloader(s) don't use DistributedSampler.
|
||||
|
||||
You're using multiple gpus and multiple nodes without using a
|
||||
DistributedSampler to assign a subset of your data to each process.
|
||||
To silence this warning, pass a DistributedSampler to your DataLoader.
|
||||
|
||||
ie: this:
|
||||
dataset = myDataset()
|
||||
dataloader = Dataloader(dataset)
|
||||
|
||||
becomes:
|
||||
dataset = myDataset()
|
||||
dist_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
|
||||
dataloader = Dataloader(dataset, sampler=dist_sampler)
|
||||
|
||||
If you want each process to load the full dataset, ignore this warning.
|
||||
"""
|
||||
warnings.warn(msg)
|
||||
break
|
||||
|
||||
if self.use_ddp or self.use_ddp2:
|
||||
# wait for all processes to catch up
|
||||
dist.barrier()
|
||||
|
||||
# load each dataloader
|
||||
self.get_train_dataloader()
|
||||
self.get_test_dataloaders()
|
||||
self.get_val_dataloaders()
|
||||
|
||||
def determine_data_use_amount(self, train_percent_check, val_percent_check,
|
||||
test_percent_check, overfit_pct):
|
||||
"""
|
||||
Use less data for debugging purposes
|
||||
"""
|
||||
self.train_percent_check = train_percent_check
|
||||
self.val_percent_check = val_percent_check
|
||||
self.test_percent_check = test_percent_check
|
||||
if overfit_pct > 0:
|
||||
self.train_percent_check = overfit_pct
|
||||
self.val_percent_check = overfit_pct
|
||||
self.test_percent_check = overfit_pct
|
||||
@@ -0,0 +1,236 @@
|
||||
import os
|
||||
import re
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel
|
||||
from pytorch_lightning.utilities.debugging import MisconfigurationException
|
||||
|
||||
|
||||
try:
|
||||
from apex import amp
|
||||
APEX_AVAILABLE = True
|
||||
except ImportError:
|
||||
APEX_AVAILABLE = False
|
||||
|
||||
|
||||
class TrainerDDPMixin(object):
|
||||
def set_distributed_mode(self, distributed_backend, nb_gpu_nodes):
|
||||
# skip for CPU
|
||||
if self.num_gpus == 0:
|
||||
return
|
||||
|
||||
# single GPU case
|
||||
# in single gpu case we allow ddp so we can train on multiple
|
||||
# nodes, 1 gpu per node
|
||||
if self.num_gpus == 1:
|
||||
self.single_gpu = True
|
||||
|
||||
if distributed_backend is not None:
|
||||
self.use_dp = distributed_backend == 'dp'
|
||||
self.use_ddp = distributed_backend == 'ddp'
|
||||
self.use_ddp2 = distributed_backend == 'ddp2'
|
||||
|
||||
# disable single gpu when using ddp2
|
||||
if self.use_ddp2:
|
||||
self.single_gpu = False
|
||||
|
||||
# multiple GPU case
|
||||
elif self.num_gpus > 1:
|
||||
if distributed_backend is not None:
|
||||
# DP, DDP case
|
||||
self.use_dp = distributed_backend == 'dp'
|
||||
self.use_ddp = distributed_backend == 'ddp'
|
||||
self.use_ddp2 = distributed_backend == 'ddp2'
|
||||
|
||||
elif distributed_backend is None:
|
||||
m = 'When using multiple GPUs set ' \
|
||||
'Trainer(distributed_backend=dp) (or ddp)'
|
||||
raise MisconfigurationException(m)
|
||||
|
||||
# use ddp automatically if nb_gpu_nodes > 1
|
||||
if nb_gpu_nodes > 1 and self.use_dp: # pragma: no cover
|
||||
self.use_ddp = True
|
||||
self.use_dp = False
|
||||
w = 'DataParallel does not support nb_gpu_nodes > 1. ' \
|
||||
'Switching to DistributedDataParallel for you. ' \
|
||||
'To silence this warning set distributed_backend=ddp'
|
||||
warnings.warn(w)
|
||||
|
||||
print('gpu available: {}, used: {}'.format(torch.cuda.is_available(), self.on_gpu))
|
||||
|
||||
def configure_slurm_ddp(self, nb_gpu_nodes):
|
||||
self.is_slurm_managing_tasks = False
|
||||
|
||||
# extract SLURM flag vars
|
||||
# whenever we have the correct number of tasks, we let slurm manage processes
|
||||
# otherwise we launch the required number of processes
|
||||
if self.use_ddp:
|
||||
self.nb_requested_gpus = self.num_gpus * nb_gpu_nodes
|
||||
self.nb_slurm_tasks = 0
|
||||
try:
|
||||
self.nb_slurm_tasks = int(os.environ['SLURM_NTASKS'])
|
||||
self.is_slurm_managing_tasks = self.nb_slurm_tasks == self.nb_requested_gpus
|
||||
|
||||
# in interactive mode we don't manage tasks
|
||||
job_name = os.environ['SLURM_JOB_NAME']
|
||||
if job_name == 'bash':
|
||||
self.is_slurm_managing_tasks = False
|
||||
|
||||
except Exception:
|
||||
# likely not on slurm, so set the slurm managed flag to false
|
||||
self.is_slurm_managing_tasks = False
|
||||
|
||||
# used for tests only, set this flag to simulate slurm managing a task
|
||||
try:
|
||||
should_fake = int(os.environ['FAKE_SLURM_MANAGING_TASKS'])
|
||||
if should_fake:
|
||||
self.is_slurm_managing_tasks = True
|
||||
except Exception as e:
|
||||
pass
|
||||
|
||||
def set_nvidia_flags(self, is_slurm_managing_tasks, data_parallel_device_ids):
|
||||
if data_parallel_device_ids is None:
|
||||
return
|
||||
|
||||
# set the correct cuda visible devices (using pci order)
|
||||
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
|
||||
|
||||
# when slurm is managing the task it sets the visible devices
|
||||
if not is_slurm_managing_tasks:
|
||||
if type(data_parallel_device_ids) is int:
|
||||
id_str = ','.join(str(x) for x in list(range(data_parallel_device_ids)))
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = id_str
|
||||
else:
|
||||
gpu_str = ','.join([str(x) for x in data_parallel_device_ids])
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_str
|
||||
|
||||
print(f'VISIBLE GPUS: {os.environ["CUDA_VISIBLE_DEVICES"]}')
|
||||
|
||||
def ddp_train(self, gpu_nb, model):
|
||||
"""
|
||||
Entry point into a DP thread
|
||||
:param gpu_nb:
|
||||
:param model:
|
||||
:param cluster_obj:
|
||||
:return:
|
||||
"""
|
||||
# node rank using relative slurm id
|
||||
# otherwise default to node rank 0
|
||||
try:
|
||||
node_id = os.environ['SLURM_NODEID']
|
||||
self.node_rank = int(node_id)
|
||||
except Exception:
|
||||
self.node_rank = 0
|
||||
|
||||
# show progressbar only on progress_rank 0
|
||||
self.show_progress_bar = self.show_progress_bar and self.node_rank == 0 and gpu_nb == 0
|
||||
|
||||
# determine which process we are and world size
|
||||
if self.use_ddp:
|
||||
self.proc_rank = self.node_rank * self.num_gpus + gpu_nb
|
||||
self.world_size = self.nb_gpu_nodes * self.num_gpus
|
||||
|
||||
elif self.use_ddp2:
|
||||
self.proc_rank = self.node_rank
|
||||
self.world_size = self.nb_gpu_nodes
|
||||
|
||||
# let the exp know the rank to avoid overwriting logs
|
||||
if self.logger is not None:
|
||||
self.logger.rank = self.proc_rank
|
||||
|
||||
# set up server using proc 0's ip address
|
||||
# try to init for 20 times at max in case ports are taken
|
||||
# where to store ip_table
|
||||
self.__init_tcp_connection()
|
||||
|
||||
# CHOOSE OPTIMIZER
|
||||
# allow for lr schedulers as well
|
||||
self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers())
|
||||
|
||||
# MODEL
|
||||
# copy model to each gpu
|
||||
if self.distributed_backend == 'ddp':
|
||||
torch.cuda.set_device(gpu_nb)
|
||||
model.cuda(gpu_nb)
|
||||
|
||||
# set model properties before going into wrapper
|
||||
self.copy_trainer_model_properties(model)
|
||||
|
||||
# override root GPU
|
||||
self.root_gpu = gpu_nb
|
||||
|
||||
# AMP
|
||||
# run through amp wrapper before going to distributed DP
|
||||
if self.use_amp:
|
||||
# An example
|
||||
model, optimizers = amp.initialize(
|
||||
model, self.optimizers, opt_level=self.amp_level,
|
||||
)
|
||||
self.optimizers = optimizers
|
||||
|
||||
# DDP2 uses all GPUs on the machine
|
||||
if self.distributed_backend == 'ddp':
|
||||
device_ids = [gpu_nb]
|
||||
elif self.use_ddp2:
|
||||
device_ids = None
|
||||
|
||||
model = LightningDistributedDataParallel(
|
||||
model,
|
||||
device_ids=device_ids,
|
||||
find_unused_parameters=True
|
||||
)
|
||||
|
||||
# continue training routine
|
||||
self.run_pretrain_routine(model)
|
||||
|
||||
def __init_tcp_connection(self):
|
||||
"""
|
||||
Connect all procs in the world using the env:// init
|
||||
Use the first node as the root address
|
||||
:param port:
|
||||
:param tries:
|
||||
:return:
|
||||
"""
|
||||
|
||||
# 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.resolve_root_node_address(root_node)
|
||||
os.environ['MASTER_ADDR'] = root_node
|
||||
dist.init_process_group("nccl", rank=self.proc_rank, world_size=self.world_size)
|
||||
|
||||
def resolve_root_node_address(self, root_node):
|
||||
if '[' in root_node:
|
||||
name = root_node.split('[')[0]
|
||||
number = root_node.split(',')[0]
|
||||
if '-' in number:
|
||||
number = number.split('-')[0]
|
||||
|
||||
number = re.sub('[^0-9]', '', number)
|
||||
root_node = name + number
|
||||
|
||||
return root_node
|
||||
@@ -0,0 +1,111 @@
|
||||
import os
|
||||
import re
|
||||
import signal
|
||||
import pdb
|
||||
from subprocess import call
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from pytorch_lightning.pt_overrides.override_data_parallel import (
|
||||
LightningDistributedDataParallel, LightningDataParallel)
|
||||
from pytorch_lightning.utilities.debugging import MisconfigurationException
|
||||
|
||||
try:
|
||||
from apex import amp
|
||||
APEX_AVAILABLE = True
|
||||
except ImportError:
|
||||
APEX_AVAILABLE = False
|
||||
|
||||
|
||||
class TrainerDPMixin(object):
|
||||
def copy_trainer_model_properties(self, model):
|
||||
if isinstance(model, LightningDataParallel):
|
||||
ref_model = model.module
|
||||
elif isinstance(model, LightningDistributedDataParallel):
|
||||
ref_model = model.module
|
||||
else:
|
||||
ref_model = model
|
||||
|
||||
for m in [model, ref_model]:
|
||||
m.trainer = self
|
||||
m.on_gpu = self.on_gpu
|
||||
m.use_dp = self.use_dp
|
||||
m.use_ddp2 = self.use_ddp2
|
||||
m.use_ddp = self.use_ddp
|
||||
m.use_amp = self.use_amp
|
||||
m.testing = self.testing
|
||||
m.single_gpu = self.single_gpu
|
||||
|
||||
def transfer_batch_to_gpu(self, batch, gpu_id):
|
||||
# base case: object can be directly moved using `cuda` or `to`
|
||||
if callable(getattr(batch, 'cuda', None)):
|
||||
return batch.cuda(gpu_id)
|
||||
|
||||
elif callable(getattr(batch, 'to', None)):
|
||||
return batch.to(torch.device('cuda', gpu_id))
|
||||
|
||||
# when list
|
||||
elif isinstance(batch, list):
|
||||
for i, x in enumerate(batch):
|
||||
batch[i] = self.transfer_batch_to_gpu(x, gpu_id)
|
||||
return batch
|
||||
|
||||
# when tuple
|
||||
elif isinstance(batch, tuple):
|
||||
batch = list(batch)
|
||||
for i, x in enumerate(batch):
|
||||
batch[i] = self.transfer_batch_to_gpu(x, gpu_id)
|
||||
return tuple(batch)
|
||||
|
||||
# when dict
|
||||
elif isinstance(batch, dict):
|
||||
for k, v in batch.items():
|
||||
batch[k] = self.transfer_batch_to_gpu(v, gpu_id)
|
||||
|
||||
return batch
|
||||
|
||||
# nothing matches, return the value as is without transform
|
||||
return batch
|
||||
|
||||
def single_gpu_train(self, model):
|
||||
# CHOOSE OPTIMIZER
|
||||
# allow for lr schedulers as well
|
||||
self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers())
|
||||
|
||||
model.cuda(self.root_gpu)
|
||||
|
||||
if self.use_amp:
|
||||
# An example
|
||||
model, optimizers = amp.initialize(
|
||||
model, self.optimizers, opt_level=self.amp_level,
|
||||
)
|
||||
self.optimizers = optimizers
|
||||
|
||||
self.run_pretrain_routine(model)
|
||||
|
||||
def dp_train(self, model):
|
||||
|
||||
# CHOOSE OPTIMIZER
|
||||
# allow for lr schedulers as well
|
||||
self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers())
|
||||
|
||||
model.cuda(self.root_gpu)
|
||||
|
||||
# check for this bug (amp + dp + !01 doesn't work)
|
||||
# https://github.com/NVIDIA/apex/issues/227
|
||||
if self.use_dp and self.use_amp:
|
||||
m = f"""
|
||||
Amp level {self.amp_level} with DataParallel is not supported.
|
||||
See this note from NVIDIA for more info: https://github.com/NVIDIA/apex/issues/227.
|
||||
We recommend you switch to ddp if you want to use amp
|
||||
"""
|
||||
raise MisconfigurationException(m)
|
||||
|
||||
# create list of device ids
|
||||
device_ids = self.data_parallel_device_ids
|
||||
if type(device_ids) is int:
|
||||
device_ids = list(range(device_ids))
|
||||
|
||||
model = LightningDataParallel(model, device_ids=device_ids)
|
||||
|
||||
self.run_pretrain_routine(model)
|
||||
@@ -0,0 +1,171 @@
|
||||
import torch
|
||||
from pytorch_lightning.utilities.debugging import MisconfigurationException
|
||||
|
||||
|
||||
class TrainerEvaluationLoopMixin(object):
|
||||
|
||||
def evaluate(self, model, dataloaders, max_batches, test=False):
|
||||
"""
|
||||
Run evaluation code
|
||||
:param model: PT model
|
||||
:param dataloaders: list of PT dataloaders
|
||||
:param max_batches: Scalar
|
||||
:param test: boolean
|
||||
:return:
|
||||
"""
|
||||
# enable eval mode
|
||||
model.zero_grad()
|
||||
model.eval()
|
||||
|
||||
# copy properties for forward overrides
|
||||
self.copy_trainer_model_properties(model)
|
||||
|
||||
# disable gradients to save memory
|
||||
torch.set_grad_enabled(False)
|
||||
|
||||
# bookkeeping
|
||||
outputs = []
|
||||
|
||||
# run training
|
||||
for dataloader_idx, dataloader in enumerate(dataloaders):
|
||||
dl_outputs = []
|
||||
for batch_idx, batch in enumerate(dataloader):
|
||||
|
||||
if batch is None: # pragma: no cover
|
||||
continue
|
||||
|
||||
# stop short when on fast_dev_run (sets max_batch=1)
|
||||
if batch_idx >= max_batches:
|
||||
break
|
||||
|
||||
# -----------------
|
||||
# RUN EVALUATION STEP
|
||||
# -----------------
|
||||
output = self.evaluation_forward(model,
|
||||
batch,
|
||||
batch_idx,
|
||||
dataloader_idx,
|
||||
test)
|
||||
|
||||
# track outputs for collation
|
||||
dl_outputs.append(output)
|
||||
|
||||
# batch done
|
||||
if self.show_progress_bar:
|
||||
self.progress_bar.update(1)
|
||||
outputs.append(dl_outputs)
|
||||
|
||||
eval_results = {}
|
||||
|
||||
# with a single dataloader don't pass an array
|
||||
if len(dataloaders) == 1:
|
||||
outputs = outputs[0]
|
||||
|
||||
# give model a chance to do something with the outputs (and method defined)
|
||||
model = self.get_model()
|
||||
if test and self.is_overriden('test_end'):
|
||||
eval_results = model.test_end(outputs)
|
||||
elif self.is_overriden('validation_end'):
|
||||
eval_results = model.validation_end(outputs)
|
||||
|
||||
# enable train mode again
|
||||
model.train()
|
||||
|
||||
# enable gradients to save memory
|
||||
torch.set_grad_enabled(True)
|
||||
|
||||
return eval_results
|
||||
|
||||
def run_evaluation(self, test=False):
|
||||
# when testing make sure user defined a test step
|
||||
can_run_test_step = False
|
||||
if test:
|
||||
can_run_test_step = self.is_overriden('test_step') and self.is_overriden('test_end')
|
||||
if not can_run_test_step:
|
||||
m = '''You called .test() without defining a test step or test_end.
|
||||
Please define and try again'''
|
||||
raise MisconfigurationException(m)
|
||||
|
||||
# validate only if model has validation_step defined
|
||||
# test only if test_step or validation_step are defined
|
||||
run_val_step = self.is_overriden('validation_step')
|
||||
|
||||
if run_val_step or can_run_test_step:
|
||||
|
||||
# hook
|
||||
model = self.get_model()
|
||||
model.on_pre_performance_check()
|
||||
|
||||
# select dataloaders
|
||||
if test:
|
||||
dataloaders = self.get_test_dataloaders()
|
||||
max_batches = self.nb_test_batches
|
||||
else:
|
||||
# val
|
||||
dataloaders = self.get_val_dataloaders()
|
||||
max_batches = self.nb_val_batches
|
||||
|
||||
# cap max batches to 1 when using fast_dev_run
|
||||
if self.fast_dev_run:
|
||||
max_batches = 1
|
||||
|
||||
# run evaluation
|
||||
eval_results = self.evaluate(self.model,
|
||||
dataloaders,
|
||||
max_batches,
|
||||
test)
|
||||
_, prog_bar_metrics, log_metrics, callback_metrics = self.process_output(eval_results)
|
||||
|
||||
# add metrics to prog bar
|
||||
self.add_tqdm_metrics(prog_bar_metrics)
|
||||
|
||||
# log metrics
|
||||
self.log_metrics(log_metrics, {})
|
||||
|
||||
# track metrics for callbacks
|
||||
self.callback_metrics = callback_metrics
|
||||
|
||||
# hook
|
||||
model.on_post_performance_check()
|
||||
|
||||
if self.show_progress_bar:
|
||||
# add model specific metrics
|
||||
tqdm_metrics = self.training_tqdm_dict
|
||||
self.progress_bar.set_postfix(**tqdm_metrics)
|
||||
|
||||
# model checkpointing
|
||||
if self.proc_rank == 0 and self.checkpoint_callback is not None and not test:
|
||||
self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch,
|
||||
logs=self.callback_metrics)
|
||||
|
||||
def evaluation_forward(self, model, batch, batch_idx, dataloader_idx, test=False):
|
||||
# make dataloader_idx arg in validation_step optional
|
||||
args = [batch, batch_idx]
|
||||
|
||||
if test and len(self.get_test_dataloaders()) > 1:
|
||||
args.append(dataloader_idx)
|
||||
|
||||
elif not test and len(self.get_val_dataloaders()) > 1:
|
||||
args.append(dataloader_idx)
|
||||
|
||||
# handle DP, DDP forward
|
||||
if self.use_ddp or self.use_dp or self.use_ddp2:
|
||||
output = model(*args)
|
||||
return output
|
||||
|
||||
# single GPU
|
||||
if self.single_gpu:
|
||||
# for single GPU put inputs on gpu manually
|
||||
root_gpu = 0
|
||||
if type(self.data_parallel_device_ids) is list:
|
||||
root_gpu = self.data_parallel_device_ids[0]
|
||||
batch = self.transfer_batch_to_gpu(batch, root_gpu)
|
||||
args[0] = batch
|
||||
|
||||
# CPU
|
||||
if test:
|
||||
output = model.test_step(*args)
|
||||
else:
|
||||
output = model.validation_step(*args)
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,157 @@
|
||||
import torch
|
||||
from pytorch_lightning.root_module import memory
|
||||
|
||||
|
||||
class TrainerLoggingMixin(object):
|
||||
|
||||
def log_metrics(self, metrics, grad_norm_dic):
|
||||
"""
|
||||
Logs the metric dict passed in
|
||||
:param metrics:
|
||||
:param grad_norm_dic:
|
||||
:return:
|
||||
"""
|
||||
# added metrics by Lightning for convenience
|
||||
metrics['epoch'] = self.current_epoch
|
||||
|
||||
# add gpu memory
|
||||
if self.on_gpu and self.log_gpu_memory:
|
||||
mem_map = memory.get_memory_profile(self.log_gpu_memory)
|
||||
metrics.update(mem_map)
|
||||
|
||||
# add norms
|
||||
metrics.update(grad_norm_dic)
|
||||
|
||||
# turn all tensors to scalars
|
||||
scalar_metrics = self.metrics_to_scalars(metrics)
|
||||
|
||||
# log actual metrics
|
||||
if self.proc_rank == 0 and self.logger is not None:
|
||||
self.logger.log_metrics(scalar_metrics, step_num=self.global_step)
|
||||
self.logger.save()
|
||||
|
||||
def add_tqdm_metrics(self, metrics):
|
||||
for k, v in metrics.items():
|
||||
if type(v) is torch.Tensor:
|
||||
v = v.item()
|
||||
|
||||
self.tqdm_metrics[k] = v
|
||||
|
||||
def metrics_to_scalars(self, metrics):
|
||||
new_metrics = {}
|
||||
for k, v in metrics.items():
|
||||
if isinstance(v, torch.Tensor):
|
||||
v = v.item()
|
||||
|
||||
if type(v) is dict:
|
||||
v = self.metrics_to_scalars(v)
|
||||
|
||||
new_metrics[k] = v
|
||||
|
||||
return new_metrics
|
||||
|
||||
def process_output(self, output, train=False):
|
||||
"""
|
||||
Reduces output according to the training mode.
|
||||
Separates loss from logging and tqdm metrics
|
||||
:param output:
|
||||
:return:
|
||||
"""
|
||||
# ---------------
|
||||
# EXTRACT CALLBACK KEYS
|
||||
# ---------------
|
||||
# all keys not progress_bar or log are candidates for callbacks
|
||||
callback_metrics = {}
|
||||
for k, v in output.items():
|
||||
if k not in ['progress_bar', 'log']:
|
||||
callback_metrics[k] = v
|
||||
|
||||
if train and (self.use_dp or self.use_ddp2):
|
||||
nb_gpus = self.num_gpus
|
||||
callback_metrics = self.reduce_distributed_output(callback_metrics, nb_gpus)
|
||||
|
||||
for k, v in callback_metrics.items():
|
||||
callback_metrics[k] = v.item()
|
||||
|
||||
# ---------------
|
||||
# EXTRACT PROGRESS BAR KEYS
|
||||
# ---------------
|
||||
try:
|
||||
progress_output = output['progress_bar']
|
||||
|
||||
# reduce progress metrics for tqdm when using dp
|
||||
if train and (self.use_dp or self.use_ddp2):
|
||||
nb_gpus = self.num_gpus
|
||||
progress_output = self.reduce_distributed_output(progress_output, nb_gpus)
|
||||
|
||||
progress_bar_metrics = progress_output
|
||||
except Exception:
|
||||
progress_bar_metrics = {}
|
||||
|
||||
# ---------------
|
||||
# EXTRACT LOGGING KEYS
|
||||
# ---------------
|
||||
# extract metrics to log to experiment
|
||||
try:
|
||||
log_output = output['log']
|
||||
|
||||
# reduce progress metrics for tqdm when using dp
|
||||
if train and (self.use_dp or self.use_ddp2):
|
||||
nb_gpus = self.num_gpus
|
||||
log_output = self.reduce_distributed_output(log_output, nb_gpus)
|
||||
|
||||
log_metrics = log_output
|
||||
except Exception:
|
||||
log_metrics = {}
|
||||
|
||||
# ---------------
|
||||
# EXTRACT LOSS
|
||||
# ---------------
|
||||
# if output dict doesn't have the keyword loss
|
||||
# then assume the output=loss if scalar
|
||||
loss = None
|
||||
if train:
|
||||
try:
|
||||
loss = output['loss']
|
||||
except Exception:
|
||||
if type(output) is torch.Tensor:
|
||||
loss = output
|
||||
else:
|
||||
raise RuntimeError(
|
||||
'No `loss` value in the dictionary returned from `model.training_step()`.'
|
||||
)
|
||||
|
||||
# when using dp need to reduce the loss
|
||||
if self.use_dp or self.use_ddp2:
|
||||
loss = self.reduce_distributed_output(loss, self.num_gpus)
|
||||
|
||||
# use every metric passed in as a candidate for callback
|
||||
callback_metrics.update(progress_bar_metrics)
|
||||
callback_metrics.update(log_metrics)
|
||||
|
||||
# convert tensors to numpy
|
||||
for k, v in callback_metrics.items():
|
||||
if isinstance(v, torch.Tensor):
|
||||
callback_metrics[k] = v.item()
|
||||
|
||||
return loss, progress_bar_metrics, log_metrics, callback_metrics
|
||||
|
||||
def reduce_distributed_output(self, output, nb_gpus):
|
||||
if nb_gpus <= 1:
|
||||
return output
|
||||
|
||||
# when using DP, we get one output per gpu
|
||||
# average outputs and return
|
||||
if type(output) is torch.Tensor:
|
||||
return output.mean()
|
||||
|
||||
for k, v in output.items():
|
||||
# recurse on nested dics
|
||||
if isinstance(output[k], dict):
|
||||
output[k] = self.reduce_distributed_output(output[k], nb_gpus)
|
||||
|
||||
# reduce only metrics that have the same nb of gpus
|
||||
elif output[k].size(0) == nb_gpus:
|
||||
reduced = torch.mean(output[k])
|
||||
output[k] = reduced
|
||||
return output
|
||||
@@ -0,0 +1,17 @@
|
||||
from pytorch_lightning.root_module.root_module import LightningModule
|
||||
|
||||
|
||||
class TrainerModelHooksMixin(object):
|
||||
|
||||
def is_function_implemented(self, f_name):
|
||||
model = self.get_model()
|
||||
f_op = getattr(model, f_name, None)
|
||||
return callable(f_op)
|
||||
|
||||
def is_overriden(self, f_name):
|
||||
model = self.get_model()
|
||||
super_object = LightningModule
|
||||
|
||||
# when code pointers are different, it was overriden
|
||||
is_overriden = getattr(model, f_name).__code__ is not getattr(super_object, f_name).__code__
|
||||
return is_overriden
|
||||
@@ -0,0 +1,266 @@
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
from apex import amp
|
||||
APEX_AVAILABLE = True
|
||||
except ImportError:
|
||||
APEX_AVAILABLE = False
|
||||
|
||||
|
||||
class TrainerTrainLoopMixin(object):
|
||||
|
||||
def train(self):
|
||||
# run all epochs
|
||||
for epoch_nb in range(self.current_epoch, self.max_nb_epochs):
|
||||
# set seed for distributed sampler (enables shuffling for each epoch)
|
||||
if self.use_ddp and hasattr(self.get_train_dataloader().sampler, 'set_epoch'):
|
||||
self.get_train_dataloader().sampler.set_epoch(epoch_nb)
|
||||
|
||||
# get model
|
||||
model = self.get_model()
|
||||
|
||||
# update training progress in trainer and model
|
||||
model.current_epoch = epoch_nb
|
||||
self.current_epoch = epoch_nb
|
||||
self.total_batches = self.nb_training_batches + self.nb_val_batches
|
||||
self.batch_loss_value = 0 # accumulated grads
|
||||
|
||||
# limit the number of batches to 1 in fast_dev_run
|
||||
if self.fast_dev_run:
|
||||
self.total_batches = 1
|
||||
|
||||
# init progress_bar when requested
|
||||
if self.show_progress_bar:
|
||||
self.progress_bar.reset(self.total_batches)
|
||||
|
||||
# changing gradient according accumulation_scheduler
|
||||
self.accumulation_scheduler.on_epoch_begin(epoch_nb, self)
|
||||
|
||||
# -----------------
|
||||
# RUN TNG EPOCH
|
||||
# -----------------
|
||||
self.run_training_epoch()
|
||||
|
||||
# update LR schedulers
|
||||
if self.lr_schedulers is not None:
|
||||
for lr_scheduler in self.lr_schedulers:
|
||||
lr_scheduler.step(self.current_epoch)
|
||||
|
||||
# early stopping
|
||||
met_min_epochs = epoch_nb > self.min_nb_epochs
|
||||
if self.enable_early_stop and (met_min_epochs or self.fast_dev_run):
|
||||
should_stop = self.early_stop_callback.on_epoch_end(epoch=epoch_nb,
|
||||
logs=self.callback_metrics)
|
||||
# stop training
|
||||
stop = should_stop and met_min_epochs
|
||||
if stop:
|
||||
return
|
||||
|
||||
if self.logger is not None:
|
||||
self.logger.finalize("success")
|
||||
|
||||
def run_training_epoch(self):
|
||||
# before epoch hook
|
||||
if self.is_function_implemented('on_epoch_start'):
|
||||
model = self.get_model()
|
||||
model.on_epoch_start()
|
||||
|
||||
# run epoch
|
||||
for batch_nb, batch in enumerate(self.get_train_dataloader()):
|
||||
self.batch_nb = batch_nb
|
||||
self.global_step += 1
|
||||
|
||||
model = self.get_model()
|
||||
model.global_step = self.global_step
|
||||
|
||||
# stop when the flag is changed or we've gone past the amount
|
||||
# requested in the batches
|
||||
self.total_batch_nb += 1
|
||||
met_batch_limit = batch_nb >= self.nb_training_batches
|
||||
if met_batch_limit:
|
||||
break
|
||||
|
||||
# ---------------
|
||||
# RUN TRAIN STEP
|
||||
# ---------------
|
||||
output = self.run_training_batch(batch, batch_nb)
|
||||
batch_result, grad_norm_dic, batch_step_metrics = output
|
||||
early_stop_epoch = batch_result == -1
|
||||
|
||||
# ---------------
|
||||
# RUN VAL STEP
|
||||
# ---------------
|
||||
is_val_check_batch = (batch_nb + 1) % self.val_check_batch == 0
|
||||
can_check_epoch = (self.current_epoch + 1) % self.check_val_every_n_epoch == 0
|
||||
should_check_val = ((is_val_check_batch or early_stop_epoch) and can_check_epoch)
|
||||
|
||||
# fast_dev_run always forces val checking after train batch
|
||||
if self.fast_dev_run or should_check_val:
|
||||
self.run_evaluation(test=self.testing)
|
||||
|
||||
# when logs should be saved
|
||||
should_save_log = (batch_nb + 1) % self.log_save_interval == 0 or early_stop_epoch
|
||||
if should_save_log or self.fast_dev_run:
|
||||
if self.proc_rank == 0 and self.logger is not None:
|
||||
self.logger.save()
|
||||
|
||||
# when metrics should be logged
|
||||
should_log_metrics = batch_nb % self.row_log_interval == 0 or early_stop_epoch
|
||||
if should_log_metrics or self.fast_dev_run:
|
||||
|
||||
# logs user requested information to logger
|
||||
self.log_metrics(batch_step_metrics, grad_norm_dic)
|
||||
|
||||
# end epoch early
|
||||
if early_stop_epoch or self.fast_dev_run:
|
||||
break
|
||||
|
||||
# epoch end hook
|
||||
if self.is_function_implemented('on_epoch_end'):
|
||||
model = self.get_model()
|
||||
model.on_epoch_end()
|
||||
|
||||
def run_training_batch(self, batch, batch_nb):
|
||||
# track grad norms
|
||||
grad_norm_dic = {}
|
||||
|
||||
# track all metrics for callbacks
|
||||
all_callback_metrics = []
|
||||
|
||||
# track metrics to log
|
||||
all_log_metrics = []
|
||||
|
||||
if batch is None:
|
||||
return 0, grad_norm_dic
|
||||
|
||||
# hook
|
||||
if self.is_function_implemented('on_batch_start'):
|
||||
model_ref = self.get_model()
|
||||
response = model_ref.on_batch_start(batch)
|
||||
|
||||
if response == -1:
|
||||
return -1, grad_norm_dic
|
||||
|
||||
if self.show_progress_bar:
|
||||
self.progress_bar.update(1)
|
||||
|
||||
# call training_step once per optimizer
|
||||
for opt_idx, optimizer in enumerate(self.optimizers):
|
||||
|
||||
# wrap the forward step in a closure so second order methods work
|
||||
def optimizer_closure():
|
||||
# forward pass
|
||||
output = self.training_forward(batch, batch_nb, opt_idx)
|
||||
closure_loss, progress_bar_metrics, log_metrics, callback_metrics = output
|
||||
|
||||
# track metrics for callbacks
|
||||
all_callback_metrics.append(callback_metrics)
|
||||
|
||||
# track progress bar metrics
|
||||
self.add_tqdm_metrics(progress_bar_metrics)
|
||||
all_log_metrics.append(log_metrics)
|
||||
|
||||
# accumulate loss
|
||||
# (if accumulate_grad_batches = 1 no effect)
|
||||
closure_loss = closure_loss / self.accumulate_grad_batches
|
||||
|
||||
# backward pass
|
||||
if self.use_amp:
|
||||
with amp.scale_loss(closure_loss, optimizer) as scaled_loss:
|
||||
scaled_loss.backward()
|
||||
else:
|
||||
closure_loss.backward()
|
||||
|
||||
# insert after step hook
|
||||
if self.is_function_implemented('on_after_backward'):
|
||||
model_ref = self.get_model()
|
||||
model_ref.on_after_backward()
|
||||
|
||||
return closure_loss
|
||||
|
||||
# calculate loss
|
||||
loss = optimizer_closure()
|
||||
|
||||
# nan grads
|
||||
if self.print_nan_grads:
|
||||
self.print_nan_gradients()
|
||||
|
||||
# track total loss for logging (avoid mem leaks)
|
||||
self.batch_loss_value += loss.item()
|
||||
|
||||
# gradient update with accumulated gradients
|
||||
if (self.batch_nb + 1) % self.accumulate_grad_batches == 0:
|
||||
|
||||
# track gradient norms when requested
|
||||
if batch_nb % self.row_log_interval == 0:
|
||||
if self.track_grad_norm > 0:
|
||||
model = self.get_model()
|
||||
grad_norm_dic = model.grad_norm(self.track_grad_norm)
|
||||
|
||||
# clip gradients
|
||||
self.clip_gradients()
|
||||
|
||||
# calls .step(), .zero_grad()
|
||||
# override function to modify this behavior
|
||||
model = self.get_model()
|
||||
model.optimizer_step(self.current_epoch, batch_nb,
|
||||
optimizer, opt_idx, optimizer_closure)
|
||||
|
||||
# calculate running loss for display
|
||||
self.running_loss.append(self.batch_loss_value)
|
||||
self.batch_loss_value = 0
|
||||
self.avg_loss = np.mean(self.running_loss[-100:])
|
||||
|
||||
# update progress bar
|
||||
if self.show_progress_bar:
|
||||
# add model specific metrics
|
||||
tqdm_metrics = self.training_tqdm_dict
|
||||
self.progress_bar.set_postfix(**tqdm_metrics)
|
||||
|
||||
# activate batch end hook
|
||||
if self.is_function_implemented('on_batch_end'):
|
||||
model = self.get_model()
|
||||
model.on_batch_end()
|
||||
|
||||
# collapse all metrics into one dict
|
||||
all_log_metrics = {k: v for d in all_log_metrics for k, v in d.items()}
|
||||
|
||||
# track all metrics for callbacks
|
||||
self.callback_metrics = {k: v for d in all_callback_metrics for k, v in d.items()}
|
||||
|
||||
return 0, grad_norm_dic, all_log_metrics
|
||||
|
||||
def training_forward(self, batch, batch_nb, opt_idx):
|
||||
"""
|
||||
Handle forward for each training case (distributed, single gpu, etc...)
|
||||
:param batch:
|
||||
:param batch_nb:
|
||||
:return:
|
||||
"""
|
||||
# ---------------
|
||||
# FORWARD
|
||||
# ---------------
|
||||
# enable not needing to add opt_idx to training_step
|
||||
args = [batch, batch_nb]
|
||||
if len(self.optimizers) > 1:
|
||||
args.append(opt_idx)
|
||||
|
||||
if self.use_ddp or self.use_ddp2:
|
||||
output = self.model(*args)
|
||||
elif self.use_dp:
|
||||
output = self.model(*args)
|
||||
elif self.single_gpu:
|
||||
gpu_id = 0
|
||||
if type(self.data_parallel_device_ids) is list:
|
||||
gpu_id = self.data_parallel_device_ids[0]
|
||||
batch = self.transfer_batch_to_gpu(batch, gpu_id)
|
||||
args[0] = batch
|
||||
output = self.model.training_step(*args)
|
||||
|
||||
else:
|
||||
output = self.model.training_step(*args)
|
||||
|
||||
# format and reduce outputs accordingly
|
||||
output = self.process_output(output, train=True)
|
||||
loss, progress_bar_metrics, log_metrics, callback_metrics = output
|
||||
return loss, progress_bar_metrics, log_metrics, callback_metrics
|
||||
File diff suppressed because it is too large
Load Diff
@@ -12,7 +12,7 @@ from pytorch_lightning.pt_overrides.override_data_parallel import (
|
||||
|
||||
class TrainerIOMixin(object):
|
||||
|
||||
def __get_model(self):
|
||||
def get_model(self):
|
||||
is_dp_module = isinstance(self.model, (LightningDistributedDataParallel,
|
||||
LightningDataParallel))
|
||||
model = self.model.module if is_dp_module else self.model
|
||||
@@ -134,7 +134,7 @@ class TrainerIOMixin(object):
|
||||
checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)
|
||||
|
||||
# load model state
|
||||
model = self.__get_model()
|
||||
model = self.get_model()
|
||||
|
||||
# load the state_dict on the model automatically
|
||||
model.load_state_dict(checkpoint['state_dict'])
|
||||
@@ -173,7 +173,7 @@ class TrainerIOMixin(object):
|
||||
checkpoint['lr_schedulers'] = lr_schedulers
|
||||
|
||||
# add the state_dict from the model
|
||||
model = self.__get_model()
|
||||
model = self.get_model()
|
||||
checkpoint['state_dict'] = model.state_dict()
|
||||
|
||||
# give the model a chance to add a few things
|
||||
@@ -252,7 +252,7 @@ class TrainerIOMixin(object):
|
||||
filepath = '{}/hpc_ckpt_{}.ckpt'.format(folderpath, ckpt_number)
|
||||
|
||||
# give model a chance to do something on hpc_save
|
||||
model = self.__get_model()
|
||||
model = self.get_model()
|
||||
checkpoint = self.dump_checkpoint()
|
||||
|
||||
model.on_hpc_save(checkpoint)
|
||||
@@ -269,7 +269,7 @@ class TrainerIOMixin(object):
|
||||
checkpoint = torch.load(filepath, map_location=lambda storage, loc: storage)
|
||||
|
||||
# load model state
|
||||
model = self.__get_model()
|
||||
model = self.get_model()
|
||||
|
||||
# load the state_dict on the model automatically
|
||||
model.load_state_dict(checkpoint['state_dict'])
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
import torch
|
||||
|
||||
from pytorch_lightning.callbacks import GradientAccumulationScheduler
|
||||
|
||||
|
||||
class TrainerTrainingTricksMixin(object):
|
||||
|
||||
def clip_gradients(self):
|
||||
if self.gradient_clip_val > 0:
|
||||
model = self.get_model()
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), self.gradient_clip_val)
|
||||
|
||||
def print_nan_gradients(self):
|
||||
model = self.get_model()
|
||||
for param in model.parameters():
|
||||
if torch.isnan(param.grad.float()).any():
|
||||
print(param, param.grad)
|
||||
|
||||
def configure_accumulated_gradients(self, accumulate_grad_batches):
|
||||
self.accumulate_grad_batches = None
|
||||
|
||||
if isinstance(accumulate_grad_batches, dict):
|
||||
self.accumulation_scheduler = GradientAccumulationScheduler(accumulate_grad_batches)
|
||||
elif isinstance(accumulate_grad_batches, int):
|
||||
schedule = {1: accumulate_grad_batches}
|
||||
self.accumulation_scheduler = GradientAccumulationScheduler(schedule)
|
||||
else:
|
||||
raise TypeError("Gradient accumulation supports only int and dict types")
|
||||
@@ -25,11 +25,12 @@ from pytorch_lightning.callbacks import (
|
||||
)
|
||||
from pytorch_lightning.utilities.debugging import MisconfigurationException
|
||||
from pytorch_lightning.root_module import memory
|
||||
from pytorch_lightning.trainer.trainer import reduce_distributed_output
|
||||
from pytorch_lightning.root_module import model_saving
|
||||
from pytorch_lightning.trainer import trainer_io
|
||||
from pytorch_lightning.logging import TestTubeLogger
|
||||
from pl_examples import LightningTemplateModel
|
||||
from pytorch_lightning.trainer.logging_mixin import TrainerLoggingMixin
|
||||
|
||||
|
||||
# generate a list of random seeds for each test
|
||||
RANDOM_FILE_PATHS = list(np.random.randint(12000, 19000, 1000))
|
||||
@@ -996,14 +997,15 @@ def test_loading_meta_tags():
|
||||
|
||||
|
||||
def test_dp_output_reduce():
|
||||
mixin = TrainerLoggingMixin()
|
||||
reset_seed()
|
||||
|
||||
# test identity when we have a single gpu
|
||||
out = torch.rand(3, 1)
|
||||
assert reduce_distributed_output(out, nb_gpus=1) is out
|
||||
assert mixin.reduce_distributed_output(out, nb_gpus=1) is out
|
||||
|
||||
# average when we have multiples
|
||||
assert reduce_distributed_output(out, nb_gpus=2) == out.mean()
|
||||
assert mixin.reduce_distributed_output(out, nb_gpus=2) == out.mean()
|
||||
|
||||
# when we have a dict of vals
|
||||
out = {
|
||||
@@ -1012,7 +1014,7 @@ def test_dp_output_reduce():
|
||||
'c': out
|
||||
}
|
||||
}
|
||||
reduced = reduce_distributed_output(out, nb_gpus=3)
|
||||
reduced = mixin.reduce_distributed_output(out, nb_gpus=3)
|
||||
assert reduced['a'] == out['a']
|
||||
assert reduced['b']['c'] == out['b']['c']
|
||||
|
||||
|
||||
Reference in New Issue
Block a user