diff --git a/pytorch_lightning/accelerators/tpu_backend.py b/pytorch_lightning/accelerators/tpu_backend.py index 1fb73c568b..eaeb664ac6 100644 --- a/pytorch_lightning/accelerators/tpu_backend.py +++ b/pytorch_lightning/accelerators/tpu_backend.py @@ -13,6 +13,7 @@ # limitations under the License. import os +import re import torch import torch.multiprocessing as mp @@ -22,6 +23,7 @@ from pytorch_lightning.core import LightningModule from pytorch_lightning.utilities import rank_zero_info, rank_zero_only, rank_zero_warn, AMPType from pytorch_lightning.utilities.exceptions import MisconfigurationException from pytorch_lightning.accelerators.base_backend import Accelerator +from pytorch_lightning.utilities.cloud_io import atomic_save try: import torch_xla @@ -124,7 +126,7 @@ class TPUBackend(Accelerator): self.__save_end_of_training_weights(model, trainer) # persist info in spawn - trainer.transfer_distrib_spawn_state_on_fit_end(model, mp_queue, results) + self.transfer_distrib_spawn_state_on_fit_end(model, mp_queue, results) def training_step(self, args): batch = args[0] @@ -294,3 +296,25 @@ class TPUBackend(Accelerator): os.remove(path) return loaded_model + + def transfer_distrib_spawn_state_on_fit_end(self, model, mp_queue, results): + if self.trainer.distributed_backend.lower() not in ['ddp_spawn', 'ddp_cpu', 'tpu']: + return + + # track the best model path + best_model_path = None + if self.trainer.checkpoint_callback is not None: + best_model_path = self.trainer.checkpoint_callback.best_model_path + + if self.trainer.global_rank == 0 and mp_queue is not None: + rank_zero_warn('cleaning up ddp environment...') + # todo, pass complete checkpoint as state dictionary + mp_queue.put(best_model_path) + mp_queue.put(results) + + # save the last weights + last_path = None + if not self.trainer.testing and best_model_path is not None and len(best_model_path) > 0: + last_path = re.sub('.ckpt', '.tmp_end.ckpt', best_model_path) + atomic_save(model.state_dict(), last_path) + mp_queue.put(last_path) diff --git a/pytorch_lightning/trainer/connectors/__init__.py b/pytorch_lightning/trainer/connectors/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/pytorch_lightning/trainer/callback_connector.py b/pytorch_lightning/trainer/connectors/callback_connector.py similarity index 100% rename from pytorch_lightning/trainer/callback_connector.py rename to pytorch_lightning/trainer/connectors/callback_connector.py diff --git a/pytorch_lightning/trainer/data_connector.py b/pytorch_lightning/trainer/connectors/data_connector.py similarity index 100% rename from pytorch_lightning/trainer/data_connector.py rename to pytorch_lightning/trainer/connectors/data_connector.py diff --git a/pytorch_lightning/trainer/debugging_connector.py b/pytorch_lightning/trainer/connectors/debugging_connector.py similarity index 100% rename from pytorch_lightning/trainer/debugging_connector.py rename to pytorch_lightning/trainer/connectors/debugging_connector.py diff --git a/pytorch_lightning/trainer/logger_connector.py b/pytorch_lightning/trainer/connectors/logger_connector.py similarity index 100% rename from pytorch_lightning/trainer/logger_connector.py rename to pytorch_lightning/trainer/connectors/logger_connector.py diff --git a/pytorch_lightning/trainer/model_connector.py b/pytorch_lightning/trainer/connectors/model_connector.py similarity index 100% rename from pytorch_lightning/trainer/model_connector.py rename to pytorch_lightning/trainer/connectors/model_connector.py diff --git a/pytorch_lightning/trainer/optimizer_connector.py b/pytorch_lightning/trainer/connectors/optimizer_connector.py similarity index 100% rename from pytorch_lightning/trainer/optimizer_connector.py rename to pytorch_lightning/trainer/connectors/optimizer_connector.py diff --git a/pytorch_lightning/trainer/precision_connector.py b/pytorch_lightning/trainer/connectors/precision_connector.py similarity index 100% rename from pytorch_lightning/trainer/precision_connector.py rename to pytorch_lightning/trainer/connectors/precision_connector.py diff --git a/pytorch_lightning/trainer/profiler_connector.py b/pytorch_lightning/trainer/connectors/profiler_connector.py similarity index 100% rename from pytorch_lightning/trainer/profiler_connector.py rename to pytorch_lightning/trainer/connectors/profiler_connector.py diff --git a/pytorch_lightning/trainer/training_trick_connector.py b/pytorch_lightning/trainer/connectors/training_trick_connector.py similarity index 100% rename from pytorch_lightning/trainer/training_trick_connector.py rename to pytorch_lightning/trainer/connectors/training_trick_connector.py diff --git a/pytorch_lightning/trainer/properties.py b/pytorch_lightning/trainer/properties.py index c634a74aff..38ff690608 100644 --- a/pytorch_lightning/trainer/properties.py +++ b/pytorch_lightning/trainer/properties.py @@ -1,5 +1,5 @@ from pytorch_lightning.utilities.cloud_io import get_filesystem -from pytorch_lightning.trainer.logger_connector import LoggerConnector +from pytorch_lightning.trainer.connectors.logger_connector import LoggerConnector from pytorch_lightning.trainer.states import TrainerState from typing import List, Optional, Union from pytorch_lightning.utilities import argparse_utils @@ -10,7 +10,7 @@ import os from pytorch_lightning.utilities.model_utils import is_overridden from pytorch_lightning.core.lightning import LightningModule from pytorch_lightning.callbacks import ProgressBarBase -from pytorch_lightning.trainer.model_connector import ModelConnector +from pytorch_lightning.trainer.connectors.model_connector import ModelConnector class TrainerProperties(ABC): diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py index eedda5a8ff..e3740f4dc0 100644 --- a/pytorch_lightning/trainer/trainer.py +++ b/pytorch_lightning/trainer/trainer.py @@ -14,10 +14,9 @@ import os import warnings -from typing import Any, Dict, Iterable, List, Optional, Tuple, Union +from typing import Dict, Iterable, List, Optional, Union import torch -import torch.distributed as torch_distrib from torch.utils.data import DataLoader from pytorch_lightning.callbacks import Callback, EarlyStopping, ModelCheckpoint @@ -26,7 +25,7 @@ from pytorch_lightning.core.lightning import LightningModule from pytorch_lightning.core.memory import ModelSummary from pytorch_lightning.core.step_result import EvalResult from pytorch_lightning.loggers import LightningLoggerBase -from pytorch_lightning.profiler import BaseProfiler, PassThroughProfiler, SimpleProfiler +from pytorch_lightning.profiler import BaseProfiler from pytorch_lightning.trainer.callback_hook import TrainerCallbackHookMixin from pytorch_lightning.trainer.configuration_validator import ConfigValidator from pytorch_lightning.trainer.data_loading import TrainerDataLoadingMixin @@ -37,23 +36,23 @@ from pytorch_lightning.trainer.optimizers import TrainerOptimizersMixin from pytorch_lightning.trainer.states import TrainerState, trainer_state from pytorch_lightning.trainer.training_io import TrainerIOMixin from pytorch_lightning.trainer.training_tricks import TrainerTrainingTricksMixin -from pytorch_lightning.utilities import parsing, rank_zero_info, rank_zero_only, rank_zero_warn, AMPType +from pytorch_lightning.utilities import rank_zero_warn from pytorch_lightning.utilities.debugging import InternalDebugger from pytorch_lightning.utilities.exceptions import MisconfigurationException from pytorch_lightning.trainer.evaluation_loop import EvaluationLoop from pytorch_lightning.trainer.training_loop import TrainLoop from pytorch_lightning.accelerators.accelerator_connector import AcceleratorConnector -from pytorch_lightning.trainer.logger_connector import LoggerConnector -from pytorch_lightning.trainer.optimizer_connector import OptimizerConnector -from pytorch_lightning.trainer.training_trick_connector import TrainingTricksConnector -from pytorch_lightning.trainer.callback_connector import CallbackConnector -from pytorch_lightning.trainer.model_connector import ModelConnector -from pytorch_lightning.trainer.debugging_connector import DebuggingConnector +from pytorch_lightning.trainer.connectors.logger_connector import LoggerConnector +from pytorch_lightning.trainer.connectors.optimizer_connector import OptimizerConnector +from pytorch_lightning.trainer.connectors.training_trick_connector import TrainingTricksConnector +from pytorch_lightning.trainer.connectors.callback_connector import CallbackConnector +from pytorch_lightning.trainer.connectors.model_connector import ModelConnector +from pytorch_lightning.trainer.connectors.debugging_connector import DebuggingConnector from pytorch_lightning import _logger as log from pytorch_lightning.tuner.tuning import Tuner -from pytorch_lightning.trainer.precision_connector import PrecisionConnector -from pytorch_lightning.trainer.profiler_connector import ProfilerConnector -from pytorch_lightning.trainer.data_connector import DataConnector +from pytorch_lightning.trainer.connectors.precision_connector import PrecisionConnector +from pytorch_lightning.trainer.connectors.profiler_connector import ProfilerConnector +from pytorch_lightning.trainer.connectors.data_connector import DataConnector from pytorch_lightning.utilities.model_utils import is_overridden from pytorch_lightning.trainer import docstrings from pytorch_lightning.trainer.properties import TrainerProperties