mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-09-10 12:21:57 +08:00
Replaces ddp .spawn with subprocess (#2029)
* replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * replace ddp spawn with subprocess * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix * hot fix
This commit is contained in:
+2
-2
@@ -12,8 +12,8 @@ rm -rf ./tests/cometruns*
|
||||
rm -rf ./tests/wandb*
|
||||
rm -rf ./tests/tests/*
|
||||
rm -rf ./lightning_logs
|
||||
python -m coverage run --source pytorch_lightning -m py.test pytorch_lightning tests pl_examples -v --doctest-modules --flake8
|
||||
python -m coverage run --source pytorch_lightning -m py.test pytorch_lightning tests pl_examples -v --doctest-modules --flake8 --durations=0
|
||||
python -m coverage report -m
|
||||
|
||||
# specific file
|
||||
# python -m coverage run --source pytorch_lightning -m py.test -k test_trainer.py --flake8
|
||||
# python -m coverage run --source pytorch_lightning -m py.test -k test_trainer.py --flake8 --durations=0
|
||||
|
||||
@@ -10,25 +10,23 @@ import torch
|
||||
import pytorch_lightning as pl
|
||||
from pl_examples.models.lightning_template import LightningTemplateModel
|
||||
|
||||
SEED = 2334
|
||||
torch.manual_seed(SEED)
|
||||
np.random.seed(SEED)
|
||||
pl.seed_everything(234)
|
||||
|
||||
|
||||
def main(hparams):
|
||||
def main(args):
|
||||
"""
|
||||
Main training routine specific for this project
|
||||
:param hparams:
|
||||
:param args:
|
||||
"""
|
||||
# ------------------------
|
||||
# 1 INIT LIGHTNING MODEL
|
||||
# ------------------------
|
||||
model = LightningTemplateModel(hparams)
|
||||
model = LightningTemplateModel(**vars(args))
|
||||
|
||||
# ------------------------
|
||||
# 2 INIT TRAINER
|
||||
# ------------------------
|
||||
trainer = pl.Trainer(max_epochs=hparams.epochs, overfit_pct=0.01, early_stop_callback=True)
|
||||
trainer = pl.Trainer.from_argparse_args(args)
|
||||
|
||||
# ------------------------
|
||||
# 3 START TRAINING
|
||||
@@ -46,9 +44,10 @@ if __name__ == '__main__':
|
||||
|
||||
# each LightningModule defines arguments relevant to it
|
||||
parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
|
||||
hyperparams = parser.parse_args()
|
||||
parser = pl.Trainer.add_argparse_args(parser)
|
||||
args = parser.parse_args()
|
||||
|
||||
# ---------------------
|
||||
# RUN TRAINING
|
||||
# ---------------------
|
||||
main(hyperparams)
|
||||
main(args)
|
||||
|
||||
@@ -957,7 +957,7 @@ class LightningModule(ABC, DeviceDtypeModuleMixin, GradInformation, ModelIO, Mod
|
||||
f"is not equal to the computed world size ({world_size}). Ignored.")
|
||||
|
||||
torch_backend = "nccl" if self.trainer.on_gpu else "gloo"
|
||||
log.info(f"initializing proc_rank {proc_rank} world {world_size}")
|
||||
log.info(f"initializing ddp: LOCAL_RANK: {proc_rank}/{world_size - 1} WORLD_SIZE:{world_size}")
|
||||
torch_distrib.init_process_group(torch_backend, rank=proc_rank, world_size=world_size)
|
||||
|
||||
def configure_apex(
|
||||
|
||||
@@ -117,6 +117,11 @@ import os
|
||||
import re
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Union
|
||||
import subprocess
|
||||
import sys
|
||||
from time import sleep
|
||||
import numpy as np
|
||||
from os.path import abspath
|
||||
|
||||
import torch
|
||||
from pytorch_lightning import _logger as log
|
||||
@@ -311,7 +316,7 @@ class TrainerDDPMixin(ABC):
|
||||
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 not is_slurm_managing_tasks and 'CUDA_VISIBLE_DEVICES' not in os.environ:
|
||||
if isinstance(data_parallel_device_ids, int):
|
||||
id_str = ','.join(str(x) for x in list(range(data_parallel_device_ids)))
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = id_str
|
||||
@@ -322,7 +327,74 @@ class TrainerDDPMixin(ABC):
|
||||
# don't make this debug... this is good UX
|
||||
log.info(f'CUDA_VISIBLE_DEVICES: [{os.environ["CUDA_VISIBLE_DEVICES"]}]')
|
||||
|
||||
def ddp_train(self, process_idx, model):
|
||||
def __set_random_port(self):
|
||||
"""
|
||||
When running DDP NOT managed by SLURM, the ports might collide
|
||||
:return:
|
||||
"""
|
||||
try:
|
||||
default_port = os.environ['MASTER_PORT']
|
||||
except Exception:
|
||||
import random
|
||||
default_port = random.randint(10000, 19000)
|
||||
os.environ['MASTER_PORT'] = str(default_port)
|
||||
|
||||
def spawn_ddp_children(self, model):
|
||||
self.__set_random_port()
|
||||
port = os.environ['MASTER_PORT']
|
||||
|
||||
master_address = '127.0.0.1' if 'MASTER_ADDR' not in os.environ else os.environ['MASTER_ADDR']
|
||||
os.environ['MASTER_PORT'] = f'{port}'
|
||||
os.environ['MASTER_ADDR'] = f'{master_address}'
|
||||
|
||||
# allow the user to pass the node rank
|
||||
node_rank = '0'
|
||||
if 'NODE_RANK' in os.environ:
|
||||
node_rank = os.environ['NODE_RANK']
|
||||
if 'GROUP_RANK' in os.environ:
|
||||
node_rank = os.environ['GROUP_RANK']
|
||||
|
||||
os.environ['NODE_RANK'] = node_rank
|
||||
os.environ['LOCAL_RANK'] = '0'
|
||||
|
||||
# pull out the commands used to run the script and resolve the abs file path
|
||||
command = sys.argv
|
||||
full_path = abspath(command[0])
|
||||
command[0] = full_path
|
||||
command = ['python'] + command
|
||||
|
||||
# since this script sets the visible devices we replace the gpus flag with a number
|
||||
num_gpus = os.environ['CUDA_VISIBLE_DEVICES'].split(',').__len__()
|
||||
|
||||
# if script called without a flag, pass in a flag anyhow
|
||||
if '--gpus' not in command:
|
||||
arg_gpus = len(self.gpus) if isinstance(self.gpus, list) else self.gpus
|
||||
command += ['--gpus', arg_gpus]
|
||||
|
||||
gpu_flag_idx = command.index('--gpus')
|
||||
command[gpu_flag_idx + 1] = f'{num_gpus}'
|
||||
|
||||
os.environ['WORLD_SIZE'] = f'{num_gpus * self.num_nodes}'
|
||||
|
||||
self.interactive_ddp_procs = []
|
||||
for local_rank in range(1, self.num_processes):
|
||||
env_copy = os.environ.copy()
|
||||
env_copy['LOCAL_RANK'] = f'{local_rank}'
|
||||
|
||||
# import pdb; pdb.set_trace()
|
||||
# start process
|
||||
proc = subprocess.Popen(command, env=env_copy)
|
||||
self.interactive_ddp_procs.append(proc)
|
||||
|
||||
# starting all processes at once can cause issues
|
||||
# with dataloaders delay between 1-10 seconds
|
||||
delay = np.random.uniform(1, 5, 1)[0]
|
||||
sleep(delay)
|
||||
|
||||
local_rank = 0
|
||||
self.ddp_train(local_rank, model, is_master=True)
|
||||
|
||||
def ddp_train(self, process_idx, model, is_master=False):
|
||||
"""
|
||||
Entry point into a DP thread
|
||||
:param gpu_idx:
|
||||
@@ -359,7 +431,14 @@ class TrainerDDPMixin(ABC):
|
||||
# MODEL
|
||||
# copy model to each gpu
|
||||
if self.on_gpu:
|
||||
self.root_gpu = process_idx
|
||||
gpu_idx = process_idx
|
||||
if is_master:
|
||||
# source of truth is cuda for gpu idx
|
||||
gpus = os.environ['CUDA_VISIBLE_DEVICES'].split(',')
|
||||
local_rank = int(os.environ['LOCAL_RANK'])
|
||||
gpu_idx = int(gpus[local_rank])
|
||||
|
||||
self.root_gpu = gpu_idx
|
||||
torch.cuda.set_device(self.root_gpu)
|
||||
model.cuda(self.root_gpu)
|
||||
|
||||
@@ -388,9 +467,6 @@ class TrainerDDPMixin(ABC):
|
||||
# continue training routine
|
||||
self.run_pretrain_routine(model)
|
||||
|
||||
# when ddp ends, we save the model
|
||||
self.save_spawn_weights(model)
|
||||
|
||||
def save_spawn_weights(self, model):
|
||||
"""
|
||||
Dump a temporary checkpoint after ddp ends to get weights out of the process
|
||||
|
||||
@@ -685,8 +685,18 @@ def sanitize_gpu_ids(gpus):
|
||||
:return: unmodified gpus variable
|
||||
"""
|
||||
all_available_gpus = get_all_available_gpus()
|
||||
misconfig = False
|
||||
for gpu in gpus:
|
||||
if gpu not in all_available_gpus:
|
||||
misconfig = True
|
||||
|
||||
if misconfig:
|
||||
# sometimes auto ddp might have different flags
|
||||
# but this is not what the user intended
|
||||
# correct for the user
|
||||
if len(gpus) == len(all_available_gpus):
|
||||
gpus = all_available_gpus
|
||||
else:
|
||||
raise MisconfigurationException(f"""
|
||||
You requested GPUs: {gpus}
|
||||
But your machine only has: {all_available_gpus}
|
||||
|
||||
@@ -35,7 +35,6 @@ from pytorch_lightning.trainer.lr_finder import TrainerLRFinderMixin
|
||||
from pytorch_lightning.utilities.exceptions import MisconfigurationException
|
||||
from pytorch_lightning.utilities import rank_zero_warn, parsing
|
||||
|
||||
|
||||
try:
|
||||
from apex import amp
|
||||
except ImportError:
|
||||
@@ -119,7 +118,7 @@ class Trainer(
|
||||
distributed_backend: Optional[str] = None,
|
||||
precision: int = 32,
|
||||
print_nan_grads: bool = False, # backward compatible, todo: remove in v0.9.0
|
||||
weights_summary: Optional[str] = 'full',
|
||||
weights_summary: Optional[str] = 'top',
|
||||
weights_save_path: Optional[str] = None,
|
||||
num_sanity_val_steps: int = 2,
|
||||
truncated_bptt_steps: Optional[int] = None,
|
||||
@@ -494,6 +493,7 @@ class Trainer(
|
||||
# init flags for SLURM+ddp to work
|
||||
self.proc_rank = 0
|
||||
self.world_size = 1
|
||||
self.interactive_ddp_procs = []
|
||||
self.configure_slurm_ddp(self.num_nodes)
|
||||
self.node_rank = self.determine_ddp_node_rank()
|
||||
|
||||
@@ -871,16 +871,12 @@ class Trainer(
|
||||
task = int(os.environ['LOCAL_RANK'])
|
||||
self.ddp_train(task, model)
|
||||
|
||||
else:
|
||||
self.__set_random_port()
|
||||
# track for predict
|
||||
elif self.distributed_backend == 'cpu_ddp':
|
||||
self.model = model
|
||||
# train
|
||||
mp.spawn(self.ddp_train, nprocs=self.num_processes, args=(model,))
|
||||
# load weights if not interrupted
|
||||
if self.on_colab_kaggle:
|
||||
self.load_spawn_weights(model)
|
||||
self.model = model
|
||||
|
||||
elif self.distributed_backend == 'ddp':
|
||||
self.spawn_ddp_children(model)
|
||||
|
||||
# 1 gpu or dp option triggers training using DP module
|
||||
# easier to avoid NCCL issues
|
||||
@@ -928,18 +924,6 @@ class Trainer(
|
||||
# used for testing or when we need to know that training succeeded
|
||||
return 1
|
||||
|
||||
def __set_random_port(self):
|
||||
"""
|
||||
When running DDP NOT managed by SLURM, the ports might collide
|
||||
:return:
|
||||
"""
|
||||
try:
|
||||
default_port = os.environ['MASTER_PORT']
|
||||
except Exception:
|
||||
import random
|
||||
default_port = random.randint(10000, 19000)
|
||||
os.environ['MASTER_PORT'] = str(default_port)
|
||||
|
||||
def __attach_dataloaders(self, model, train_dataloader=None, val_dataloaders=None, test_dataloaders=None):
|
||||
# when dataloader is passed via fit, patch the train_dataloader
|
||||
# functions to overwrite with these implementations
|
||||
@@ -1046,7 +1030,10 @@ class Trainer(
|
||||
|
||||
# clear cache before training
|
||||
if self.on_gpu:
|
||||
torch.cuda.empty_cache()
|
||||
# use context because of:
|
||||
# https://discuss.pytorch.org/t/out-of-memory-when-i-use-torch-cuda-empty-cache/57898
|
||||
with torch.cuda.device(f'cuda:{self.root_gpu}'):
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# CORE TRAINING LOOP
|
||||
self.train()
|
||||
@@ -1096,7 +1083,10 @@ class Trainer(
|
||||
if model is not None:
|
||||
self.model = model
|
||||
self.fit(model)
|
||||
elif self.use_ddp or self.use_tpu: # pragma: no-cover
|
||||
|
||||
# on tpu, .spawn means we don't have a trained model
|
||||
# TODO: remove TPU spawn
|
||||
elif self.use_tpu: # pragma: no-cover
|
||||
# attempt to load weights from a spawn
|
||||
path = os.path.join(self.default_root_dir, '__temp_weight_ddp_end.ckpt')
|
||||
test_model = self.model
|
||||
|
||||
@@ -158,6 +158,7 @@ from pytorch_lightning.loggers import LightningLoggerBase
|
||||
from pytorch_lightning.trainer.supporters import TensorRunningAccum
|
||||
from pytorch_lightning.utilities import rank_zero_warn
|
||||
from pytorch_lightning.utilities.exceptions import MisconfigurationException
|
||||
import subprocess
|
||||
|
||||
try:
|
||||
from apex import amp
|
||||
@@ -305,13 +306,13 @@ class TrainerTrainLoopMixin(ABC):
|
||||
|
||||
def train(self):
|
||||
# add signal handlers for process kills
|
||||
def _signal_kill_handler(*args):
|
||||
return TrainerTrainLoopMixin.run_training_teardown(self)
|
||||
|
||||
orig_signal_handlers = {}
|
||||
for sig_name in SIGNAL_TERMINATE:
|
||||
orig_signal_handlers[sig_name] = signal.signal(getattr(signal, sig_name),
|
||||
_signal_kill_handler)
|
||||
# def _signal_kill_handler(*args):
|
||||
# return TrainerTrainLoopMixin.run_training_teardown(self)
|
||||
#
|
||||
# orig_signal_handlers = {}
|
||||
# for sig_name in SIGNAL_TERMINATE:
|
||||
# orig_signal_handlers[sig_name] = signal.signal(getattr(signal, sig_name),
|
||||
# _signal_kill_handler)
|
||||
|
||||
# get model
|
||||
model = self.get_model()
|
||||
@@ -384,15 +385,17 @@ class TrainerTrainLoopMixin(ABC):
|
||||
|
||||
self.run_training_teardown()
|
||||
|
||||
# reset signal handlers
|
||||
for sig_name in SIGNAL_TERMINATE:
|
||||
signal.signal(getattr(signal, sig_name), orig_signal_handlers[sig_name])
|
||||
|
||||
except KeyboardInterrupt:
|
||||
if self.proc_rank == 0:
|
||||
log.info('Detected KeyboardInterrupt, attempting graceful shutdown...')
|
||||
self.interrupted = True
|
||||
self.run_training_teardown()
|
||||
rank_zero_warn('Detected KeyboardInterrupt, attempting graceful shutdown...')
|
||||
|
||||
# user could press ctrl+c many times... only shutdown once
|
||||
if not self.interrupted:
|
||||
self.interrupted = True
|
||||
|
||||
for proc in self.interactive_ddp_procs:
|
||||
subprocess.Popen.kill(proc)
|
||||
|
||||
self.run_training_teardown()
|
||||
|
||||
def run_training_epoch(self):
|
||||
|
||||
@@ -678,7 +681,7 @@ class TrainerTrainLoopMixin(ABC):
|
||||
opt_idx = np.argmax(optimizer_freq_cumsum > current_place_in_loop)
|
||||
return [(opt_idx, self.optimizers[opt_idx])]
|
||||
|
||||
@atexit.register
|
||||
# @atexit.register
|
||||
def run_training_teardown(self):
|
||||
if hasattr(self, '_teardown_already_run') and self._teardown_already_run:
|
||||
return
|
||||
|
||||
@@ -12,7 +12,7 @@ class ModelTemplateData:
|
||||
loader = DataLoader(
|
||||
dataset=dataset,
|
||||
batch_size=self.batch_size,
|
||||
# test and valid shall not be shuffled
|
||||
num_workers=3,
|
||||
shuffle=train,
|
||||
)
|
||||
return loader
|
||||
|
||||
+2
-2
@@ -25,7 +25,7 @@ def assert_speed_parity(pl_times, pt_times, num_epochs):
|
||||
f"lightning was slower than PT (threshold {max_diff_per_epoch})"
|
||||
|
||||
|
||||
def run_model_test_without_loggers(trainer_options, model, min_acc=0.50):
|
||||
def run_model_test_without_loggers(trainer_options, model, min_acc=0.30):
|
||||
reset_seed()
|
||||
|
||||
# fit model
|
||||
@@ -155,7 +155,7 @@ def load_model_from_checkpoint(root_weights_dir, module_class=EvalModelTemplate)
|
||||
return trained_model
|
||||
|
||||
|
||||
def run_prediction(dataloader, trained_model, dp=False, min_acc=0.5):
|
||||
def run_prediction(dataloader, trained_model, dp=False, min_acc=0.3):
|
||||
# run prediction on 1 batch
|
||||
for batch in dataloader:
|
||||
break
|
||||
|
||||
@@ -220,7 +220,7 @@ def test_early_stopping_no_val_step(tmpdir):
|
||||
default_root_dir=tmpdir,
|
||||
early_stop_callback=stopping,
|
||||
overfit_pct=0.20,
|
||||
max_epochs=5,
|
||||
max_epochs=2,
|
||||
)
|
||||
result = trainer.fit(model)
|
||||
|
||||
@@ -254,7 +254,7 @@ def test_model_checkpoint_with_non_string_input(tmpdir, save_top_k):
|
||||
trainer = Trainer(default_root_dir=tmpdir,
|
||||
checkpoint_callback=checkpoint,
|
||||
overfit_pct=0.20,
|
||||
max_epochs=5
|
||||
max_epochs=2
|
||||
)
|
||||
trainer.fit(model)
|
||||
|
||||
@@ -275,7 +275,7 @@ def test_model_checkpoint_path(tmpdir, logger_version, expected):
|
||||
trainer = Trainer(
|
||||
default_root_dir=tmpdir,
|
||||
overfit_pct=0.2,
|
||||
max_epochs=5,
|
||||
max_epochs=2,
|
||||
logger=logger
|
||||
)
|
||||
trainer.fit(model)
|
||||
|
||||
@@ -16,7 +16,7 @@ def test_lr_logger_single_lr(tmpdir):
|
||||
lr_logger = LearningRateLogger()
|
||||
trainer = Trainer(
|
||||
default_root_dir=tmpdir,
|
||||
max_epochs=5,
|
||||
max_epochs=2,
|
||||
val_percent_check=0.1,
|
||||
train_percent_check=0.5,
|
||||
callbacks=[lr_logger]
|
||||
@@ -39,7 +39,7 @@ def test_lr_logger_no_lr(tmpdir):
|
||||
lr_logger = LearningRateLogger()
|
||||
trainer = Trainer(
|
||||
default_root_dir=tmpdir,
|
||||
max_epochs=5,
|
||||
max_epochs=2,
|
||||
val_percent_check=0.1,
|
||||
train_percent_check=0.5,
|
||||
callbacks=[lr_logger]
|
||||
@@ -60,7 +60,7 @@ def test_lr_logger_multi_lrs(tmpdir):
|
||||
lr_logger = LearningRateLogger()
|
||||
trainer = Trainer(
|
||||
default_root_dir=tmpdir,
|
||||
max_epochs=10,
|
||||
max_epochs=2,
|
||||
val_percent_check=0.1,
|
||||
train_percent_check=0.5,
|
||||
callbacks=[lr_logger]
|
||||
@@ -87,7 +87,7 @@ def test_lr_logger_param_groups(tmpdir):
|
||||
lr_logger = LearningRateLogger()
|
||||
trainer = Trainer(
|
||||
default_root_dir=tmpdir,
|
||||
max_epochs=5,
|
||||
max_epochs=2,
|
||||
val_percent_check=0.1,
|
||||
train_percent_check=0.5,
|
||||
callbacks=[lr_logger]
|
||||
|
||||
@@ -100,7 +100,7 @@ def test_loggers_pickle(tmpdir, monkeypatch, logger_class):
|
||||
|
||||
@pytest.mark.parametrize("extra_params", [
|
||||
pytest.param(dict(max_epochs=1, auto_scale_batch_size=True), id='Batch-size-Finder'),
|
||||
pytest.param(dict(max_epochs=10, auto_lr_find=True), id='LR-Finder'),
|
||||
pytest.param(dict(max_epochs=3, auto_lr_find=True), id='LR-Finder'),
|
||||
])
|
||||
def test_logger_reset_correctly(tmpdir, extra_params):
|
||||
""" Test that the tuners do not alter the logger reference """
|
||||
|
||||
@@ -143,7 +143,7 @@ def test_adding_step_key(tmpdir):
|
||||
model.validation_epoch_end = _validation_epoch_end
|
||||
model.training_epoch_end = _training_epoch_end
|
||||
trainer = Trainer(
|
||||
max_epochs=4,
|
||||
max_epochs=3,
|
||||
default_root_dir=tmpdir,
|
||||
train_percent_check=0.001,
|
||||
val_percent_check=0.01,
|
||||
|
||||
+79
-21
@@ -1,3 +1,4 @@
|
||||
import os
|
||||
import platform
|
||||
from collections import namedtuple
|
||||
|
||||
@@ -9,6 +10,77 @@ import tests.base.utils as tutils
|
||||
from pytorch_lightning import Trainer
|
||||
from pytorch_lightning.callbacks import EarlyStopping
|
||||
from tests.base import EvalModelTemplate
|
||||
from pytorch_lightning.callbacks import ModelCheckpoint
|
||||
|
||||
|
||||
def test_cpu_slurm_save_load(tmpdir):
|
||||
"""Verify model save/load/checkpoint on CPU."""
|
||||
hparams = EvalModelTemplate.get_default_hparams()
|
||||
model = EvalModelTemplate(**hparams)
|
||||
|
||||
# logger file to get meta
|
||||
logger = tutils.get_default_logger(tmpdir)
|
||||
version = logger.version
|
||||
|
||||
# fit model
|
||||
trainer = Trainer(
|
||||
max_epochs=1,
|
||||
logger=logger,
|
||||
train_percent_check=0.2,
|
||||
val_percent_check=0.2,
|
||||
checkpoint_callback=ModelCheckpoint(tmpdir)
|
||||
)
|
||||
result = trainer.fit(model)
|
||||
real_global_step = trainer.global_step
|
||||
|
||||
# traning complete
|
||||
assert result == 1, 'cpu model failed to complete'
|
||||
|
||||
# predict with trained model before saving
|
||||
# make a prediction
|
||||
dataloaders = model.test_dataloader()
|
||||
if not isinstance(dataloaders, list):
|
||||
dataloaders = [dataloaders]
|
||||
|
||||
for dataloader in dataloaders:
|
||||
for batch in dataloader:
|
||||
break
|
||||
|
||||
x, y = batch
|
||||
x = x.view(x.size(0), -1)
|
||||
|
||||
model.eval()
|
||||
pred_before_saving = model(x)
|
||||
|
||||
# test HPC saving
|
||||
# simulate snapshot on slurm
|
||||
saved_filepath = trainer.hpc_save(tmpdir, logger)
|
||||
assert os.path.exists(saved_filepath)
|
||||
|
||||
# new logger file to get meta
|
||||
logger = tutils.get_default_logger(tmpdir, version=version)
|
||||
|
||||
trainer = Trainer(
|
||||
max_epochs=1,
|
||||
logger=logger,
|
||||
checkpoint_callback=ModelCheckpoint(tmpdir),
|
||||
)
|
||||
model = EvalModelTemplate(**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)
|
||||
|
||||
|
||||
def test_early_stopping_cpu_model(tmpdir):
|
||||
@@ -17,6 +89,7 @@ def test_early_stopping_cpu_model(tmpdir):
|
||||
trainer_options = dict(
|
||||
default_root_dir=tmpdir,
|
||||
early_stop_callback=stopping,
|
||||
max_epochs=2,
|
||||
gradient_clip_val=1.0,
|
||||
overfit_pct=0.20,
|
||||
track_grad_norm=2,
|
||||
@@ -39,6 +112,7 @@ def test_early_stopping_cpu_model(tmpdir):
|
||||
version_parse(torch.__version__) < version_parse("1.3.0")),
|
||||
reason="Distributed training is not supported on MacOS before Torch 1.3.0")
|
||||
def test_multi_cpu_model_ddp(tmpdir):
|
||||
print('in ddp test')
|
||||
"""Make sure DDP works."""
|
||||
tutils.set_random_master_port()
|
||||
|
||||
@@ -61,19 +135,19 @@ def test_lbfgs_cpu_model(tmpdir):
|
||||
"""Test each of the trainer options."""
|
||||
trainer_options = dict(
|
||||
default_root_dir=tmpdir,
|
||||
max_epochs=2,
|
||||
max_epochs=1,
|
||||
progress_bar_refresh_rate=0,
|
||||
weights_summary='top',
|
||||
train_percent_check=1.0,
|
||||
train_percent_check=0.2,
|
||||
val_percent_check=0.2,
|
||||
)
|
||||
|
||||
hparams = EvalModelTemplate.get_default_hparams()
|
||||
hparams.update(optimizer_name='lbfgs',
|
||||
learning_rate=0.002)
|
||||
learning_rate=0.004)
|
||||
model = EvalModelTemplate(**hparams)
|
||||
model.configure_optimizers = model.configure_optimizers__lbfgs
|
||||
tutils.run_model_test_without_loggers(trainer_options, model, min_acc=0.5)
|
||||
tutils.run_model_test_without_loggers(trainer_options, model, min_acc=0.25)
|
||||
|
||||
|
||||
def test_default_logger_callbacks_cpu_model(tmpdir):
|
||||
@@ -110,7 +184,7 @@ def test_running_test_after_fitting(tmpdir):
|
||||
trainer = Trainer(
|
||||
default_root_dir=tmpdir,
|
||||
progress_bar_refresh_rate=0,
|
||||
max_epochs=8,
|
||||
max_epochs=2,
|
||||
train_percent_check=0.4,
|
||||
val_percent_check=0.2,
|
||||
test_percent_check=0.2,
|
||||
@@ -324,19 +398,3 @@ def test_tbptt_cpu_model(tmpdir):
|
||||
result = trainer.fit(model)
|
||||
|
||||
assert result == 1, 'training failed to complete'
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
|
||||
def test_single_gpu_model(tmpdir):
|
||||
"""Make sure single GPU works (DP mode)."""
|
||||
trainer_options = dict(
|
||||
default_root_dir=tmpdir,
|
||||
progress_bar_refresh_rate=0,
|
||||
max_epochs=1,
|
||||
train_percent_check=0.1,
|
||||
val_percent_check=0.1,
|
||||
gpus=1
|
||||
)
|
||||
|
||||
model = EvalModelTemplate()
|
||||
tutils.run_model_test(trainer_options, model)
|
||||
|
||||
+20
-71
@@ -5,7 +5,6 @@ import torch
|
||||
|
||||
import tests.base.utils as tutils
|
||||
from pytorch_lightning import Trainer
|
||||
from pytorch_lightning.callbacks import ModelCheckpoint
|
||||
from pytorch_lightning.core import memory
|
||||
from pytorch_lightning.trainer.distrib_parts import parse_gpu_ids, determine_root_gpu_device
|
||||
from pytorch_lightning.utilities.exceptions import MisconfigurationException
|
||||
@@ -14,6 +13,23 @@ from tests.base import EvalModelTemplate
|
||||
PRETEND_N_OF_GPUS = 16
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
|
||||
@pytest.mark.parametrize('gpus', [1, [0], [1]])
|
||||
def test_single_gpu_model(tmpdir, gpus):
|
||||
"""Make sure single GPU works (DP mode)."""
|
||||
trainer_options = dict(
|
||||
default_root_dir=tmpdir,
|
||||
progress_bar_refresh_rate=0,
|
||||
max_epochs=1,
|
||||
train_percent_check=0.1,
|
||||
val_percent_check=0.1,
|
||||
gpus=gpus
|
||||
)
|
||||
|
||||
model = EvalModelTemplate()
|
||||
tutils.run_model_test(trainer_options, model)
|
||||
|
||||
|
||||
@pytest.mark.spawn
|
||||
@pytest.mark.parametrize("backend", ['dp', 'ddp', 'ddp2'])
|
||||
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
|
||||
@@ -40,6 +56,7 @@ def test_multi_gpu_model(tmpdir, backend):
|
||||
memory.get_memory_profile('min_max')
|
||||
|
||||
|
||||
@pytest.mark.spawn
|
||||
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
|
||||
def test_ddp_all_dataloaders_passed_to_fit(tmpdir):
|
||||
"""Make sure DDP works with dataloaders passed to fit()"""
|
||||
@@ -48,8 +65,8 @@ def test_ddp_all_dataloaders_passed_to_fit(tmpdir):
|
||||
trainer_options = dict(default_root_dir=tmpdir,
|
||||
progress_bar_refresh_rate=0,
|
||||
max_epochs=1,
|
||||
train_percent_check=0.4,
|
||||
val_percent_check=0.2,
|
||||
train_percent_check=0.1,
|
||||
val_percent_check=0.1,
|
||||
gpus=[0, 1],
|
||||
distributed_backend='ddp')
|
||||
|
||||
@@ -62,74 +79,6 @@ def test_ddp_all_dataloaders_passed_to_fit(tmpdir):
|
||||
assert result == 1, "DDP doesn't work with dataloaders passed to fit()."
|
||||
|
||||
|
||||
def test_cpu_slurm_save_load(tmpdir):
|
||||
"""Verify model save/load/checkpoint on CPU."""
|
||||
hparams = EvalModelTemplate.get_default_hparams()
|
||||
model = EvalModelTemplate(**hparams)
|
||||
|
||||
# logger file to get meta
|
||||
logger = tutils.get_default_logger(tmpdir)
|
||||
version = logger.version
|
||||
|
||||
# fit model
|
||||
trainer = Trainer(
|
||||
max_epochs=1,
|
||||
logger=logger,
|
||||
checkpoint_callback=ModelCheckpoint(tmpdir)
|
||||
)
|
||||
result = trainer.fit(model)
|
||||
real_global_step = trainer.global_step
|
||||
|
||||
# traning complete
|
||||
assert result == 1, 'cpu model failed to complete'
|
||||
|
||||
# predict with trained model before saving
|
||||
# make a prediction
|
||||
dataloaders = model.test_dataloader()
|
||||
if not isinstance(dataloaders, list):
|
||||
dataloaders = [dataloaders]
|
||||
|
||||
for dataloader in dataloaders:
|
||||
for batch in dataloader:
|
||||
break
|
||||
|
||||
x, y = batch
|
||||
x = x.view(x.size(0), -1)
|
||||
|
||||
model.eval()
|
||||
pred_before_saving = model(x)
|
||||
|
||||
# test HPC saving
|
||||
# simulate snapshot on slurm
|
||||
saved_filepath = trainer.hpc_save(tmpdir, logger)
|
||||
assert os.path.exists(saved_filepath)
|
||||
|
||||
# new logger file to get meta
|
||||
logger = tutils.get_default_logger(tmpdir, version=version)
|
||||
|
||||
trainer = Trainer(
|
||||
max_epochs=1,
|
||||
logger=logger,
|
||||
checkpoint_callback=ModelCheckpoint(tmpdir),
|
||||
)
|
||||
model = EvalModelTemplate(**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)
|
||||
|
||||
|
||||
@pytest.mark.spawn
|
||||
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
|
||||
def test_multi_gpu_none_backend(tmpdir):
|
||||
|
||||
@@ -76,7 +76,7 @@ def test_running_test_pretrained_model_cpu(tmpdir):
|
||||
|
||||
trainer_options = dict(
|
||||
progress_bar_refresh_rate=0,
|
||||
max_epochs=4,
|
||||
max_epochs=3,
|
||||
train_percent_check=0.4,
|
||||
val_percent_check=0.2,
|
||||
checkpoint_callback=checkpoint,
|
||||
|
||||
@@ -249,6 +249,8 @@ def test_mixing_of_dataloader_options(tmpdir):
|
||||
|
||||
|
||||
def test_train_inf_dataloader_error(tmpdir):
|
||||
pytest.skip('TODO: fix speed of this test')
|
||||
|
||||
"""Test inf train data loader (e.g. IterableDataset)"""
|
||||
model = EvalModelTemplate()
|
||||
model.train_dataloader = model.train_dataloader__infinite
|
||||
@@ -260,6 +262,8 @@ def test_train_inf_dataloader_error(tmpdir):
|
||||
|
||||
|
||||
def test_val_inf_dataloader_error(tmpdir):
|
||||
pytest.skip('TODO: fix speed of this test')
|
||||
|
||||
"""Test inf train data loader (e.g. IterableDataset)"""
|
||||
model = EvalModelTemplate()
|
||||
model.val_dataloader = model.val_dataloader__infinite
|
||||
@@ -271,6 +275,8 @@ def test_val_inf_dataloader_error(tmpdir):
|
||||
|
||||
|
||||
def test_test_inf_dataloader_error(tmpdir):
|
||||
pytest.skip('TODO: fix speed of this test')
|
||||
|
||||
"""Test inf train data loader (e.g. IterableDataset)"""
|
||||
model = EvalModelTemplate()
|
||||
model.test_dataloader = model.test_dataloader__infinite
|
||||
@@ -283,6 +289,8 @@ def test_test_inf_dataloader_error(tmpdir):
|
||||
|
||||
@pytest.mark.parametrize('check_interval', [50, 1.0])
|
||||
def test_inf_train_dataloader(tmpdir, check_interval):
|
||||
pytest.skip('TODO: fix speed of this test')
|
||||
|
||||
"""Test inf train data loader (e.g. IterableDataset)"""
|
||||
|
||||
model = EvalModelTemplate()
|
||||
@@ -300,6 +308,8 @@ def test_inf_train_dataloader(tmpdir, check_interval):
|
||||
|
||||
@pytest.mark.parametrize('check_interval', [1.0])
|
||||
def test_inf_val_dataloader(tmpdir, check_interval):
|
||||
pytest.skip('TODO: fix speed of this test')
|
||||
|
||||
"""Test inf val data loader (e.g. IterableDataset)"""
|
||||
|
||||
model = EvalModelTemplate()
|
||||
@@ -328,7 +338,9 @@ def test_error_on_zero_len_dataloader(tmpdir):
|
||||
trainer = Trainer(
|
||||
default_root_dir=tmpdir,
|
||||
max_epochs=1,
|
||||
test_percent_check=0.5
|
||||
train_percent_check=0.1,
|
||||
val_percent_check=0.1,
|
||||
test_percent_check=0.1
|
||||
)
|
||||
trainer.fit(model)
|
||||
|
||||
@@ -347,9 +359,18 @@ def test_warning_with_few_workers(tmpdir):
|
||||
train_percent_check=0.2
|
||||
)
|
||||
|
||||
fit_options = dict(train_dataloader=model.dataloader(train=True),
|
||||
val_dataloaders=model.dataloader(train=False))
|
||||
test_options = dict(test_dataloaders=model.dataloader(train=False))
|
||||
train_dl = model.dataloader(train=True)
|
||||
train_dl.num_workers = 0
|
||||
|
||||
val_dl = model.dataloader(train=False)
|
||||
val_dl.num_workers = 0
|
||||
|
||||
train_dl = model.dataloader(train=False)
|
||||
train_dl.num_workers = 0
|
||||
|
||||
fit_options = dict(train_dataloader=train_dl,
|
||||
val_dataloaders=val_dl)
|
||||
test_options = dict(test_dataloaders=train_dl)
|
||||
|
||||
trainer = Trainer(**trainer_options)
|
||||
|
||||
@@ -436,6 +457,7 @@ def test_batch_size_smaller_than_num_gpus():
|
||||
|
||||
trainer = Trainer(
|
||||
max_epochs=1,
|
||||
train_percent_check=0.1,
|
||||
val_percent_check=0,
|
||||
gpus=num_gpus,
|
||||
)
|
||||
|
||||
@@ -83,7 +83,7 @@ def test_trainer_arg_bool(tmpdir):
|
||||
# logger file to get meta
|
||||
trainer = Trainer(
|
||||
default_save_path=tmpdir,
|
||||
max_epochs=5,
|
||||
max_epochs=2,
|
||||
auto_lr_find=True
|
||||
)
|
||||
|
||||
@@ -102,7 +102,7 @@ def test_trainer_arg_str(tmpdir):
|
||||
# logger file to get meta
|
||||
trainer = Trainer(
|
||||
default_save_path=tmpdir,
|
||||
max_epochs=5,
|
||||
max_epochs=2,
|
||||
auto_lr_find='my_fancy_lr'
|
||||
)
|
||||
|
||||
@@ -122,7 +122,7 @@ def test_call_to_trainer_method(tmpdir):
|
||||
# logger file to get meta
|
||||
trainer = Trainer(
|
||||
default_save_path=tmpdir,
|
||||
max_epochs=5,
|
||||
max_epochs=2,
|
||||
)
|
||||
|
||||
lrfinder = trainer.lr_find(model, mode='linear')
|
||||
@@ -135,6 +135,8 @@ def test_call_to_trainer_method(tmpdir):
|
||||
|
||||
|
||||
def test_accumulation_and_early_stopping(tmpdir):
|
||||
pytest.skip('TODO: speed up this test')
|
||||
|
||||
""" Test that early stopping of learning rate finder works, and that
|
||||
accumulation also works for this feature """
|
||||
|
||||
@@ -145,7 +147,7 @@ def test_accumulation_and_early_stopping(tmpdir):
|
||||
# logger file to get meta
|
||||
trainer = Trainer(
|
||||
default_save_path=tmpdir,
|
||||
accumulate_grad_batches=2
|
||||
accumulate_grad_batches=2,
|
||||
)
|
||||
|
||||
lrfinder = trainer.lr_find(model, early_stop_threshold=None)
|
||||
@@ -168,7 +170,7 @@ def test_suggestion_parameters_work(tmpdir):
|
||||
# logger file to get meta
|
||||
trainer = Trainer(
|
||||
default_save_path=tmpdir,
|
||||
max_epochs=10,
|
||||
max_epochs=3,
|
||||
)
|
||||
|
||||
lrfinder = trainer.lr_find(model)
|
||||
@@ -188,7 +190,7 @@ def test_suggestion_with_non_finite_values(tmpdir):
|
||||
# logger file to get meta
|
||||
trainer = Trainer(
|
||||
default_save_path=tmpdir,
|
||||
max_epochs=10
|
||||
max_epochs=3
|
||||
)
|
||||
|
||||
lrfinder = trainer.lr_find(model)
|
||||
|
||||
@@ -445,7 +445,7 @@ def test_trainer_min_steps_and_epochs(tmpdir):
|
||||
early_stop_callback=EarlyStopping(monitor='val_loss', min_delta=1.0),
|
||||
val_check_interval=2,
|
||||
min_epochs=1,
|
||||
max_epochs=5
|
||||
max_epochs=2
|
||||
)
|
||||
|
||||
# define less min steps than 1 epoch
|
||||
|
||||
Reference in New Issue
Block a user