diff --git a/CHANGELOG.md b/CHANGELOG.md index fa568479..1c35dcdc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,8 +8,12 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). ### Added +- Added type hints in `Trainer.fit()` and `Trainer.test()` to reflect that also a list of dataloaders can be passed in ([#1723](https://github.com/PyTorchLightning/pytorch-lightning/pull/1723)). + ### Changed +- Allow user to select individual TPU core to train on ([#1729](https://github.com/PyTorchLightning/pytorch-lightning/pull/1729)) + ### Deprecated ### Removed diff --git a/README.md b/README.md index 40fd93ea..02ca78f6 100644 --- a/README.md +++ b/README.md @@ -287,7 +287,11 @@ trainer = Trainer(max_epochs=1, gpus=8, num_nodes=32) Or TPUs ```python -trainer = Trainer(num_tpu_cores=8) +# Distributes TPU core training +trainer = Trainer(tpu_cores=8) + +# Single TPU core training +trainer = Trainer(tpu_cores=[1]) ``` When you're done training, run the test accuracy diff --git a/docs/source/apex.rst b/docs/source/apex.rst index f705e040..8cfee78e 100644 --- a/docs/source/apex.rst +++ b/docs/source/apex.rst @@ -58,7 +58,7 @@ TPU 16-bit .. testcode:: # DEFAULT - trainer = Trainer(num_tpu_cores=8, precision=32) + trainer = Trainer(tpu_cores=8, precision=32) # turn on 16-bit - trainer = Trainer(num_tpu_cores=8, precision=16) + trainer = Trainer(tpu_cores=8, precision=16) diff --git a/docs/source/introduction_guide.rst b/docs/source/introduction_guide.rst index 5d262784..210dec4d 100644 --- a/docs/source/introduction_guide.rst +++ b/docs/source/introduction_guide.rst @@ -185,7 +185,7 @@ EXACTLY the same as you would a PyTorch Module. Out: - .. code-block:: none + .. code-block:: python torch.Size([1, 10]) @@ -519,50 +519,8 @@ First, change the runtime to TPU (and reinstall lightning). Next, install the required xla library (adds support for PyTorch on TPUs) -.. code-block:: python - - import collections - from datetime import datetime, timedelta - import os - import requests - import threading - - _VersionConfig = collections.namedtuple('_VersionConfig', 'wheels,server') - VERSION = "torch_xla==nightly" #@param ["xrt==1.15.0", "torch_xla==nightly"] - CONFIG = { - 'xrt==1.15.0': _VersionConfig('1.15', '1.15.0'), - 'torch_xla==nightly': _VersionConfig('nightly', 'XRT-dev{}'.format( - (datetime.today() - timedelta(1)).strftime('%Y%m%d'))), - }[VERSION] - DIST_BUCKET = 'gs://tpu-pytorch/wheels' - TORCH_WHEEL = 'torch-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels) - TORCH_XLA_WHEEL = 'torch_xla-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels) - TORCHVISION_WHEEL = 'torchvision-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels) - - # Update TPU XRT version - def update_server_xrt(): - print('Updating server-side XRT to {} ...'.format(CONFIG.server)) - url = 'http://{TPU_ADDRESS}:8475/requestversion/{XRT_VERSION}'.format( - TPU_ADDRESS=os.environ['COLAB_TPU_ADDR'].split(':')[0], - XRT_VERSION=CONFIG.server, - ) - print('Done updating server-side XRT: {}'.format(requests.post(url))) - - update = threading.Thread(target=update_server_xrt) - update.start() - -.. code-block:: - - # Install Colab TPU compat PyTorch/TPU wheels and dependencies - !pip uninstall -y torch torchvision - !gsutil cp "$DIST_BUCKET/$TORCH_WHEEL" . - !gsutil cp "$DIST_BUCKET/$TORCH_XLA_WHEEL" . - !gsutil cp "$DIST_BUCKET/$TORCHVISION_WHEEL" . - !pip install "$TORCH_WHEEL" - !pip install "$TORCH_XLA_WHEEL" - !pip install "$TORCHVISION_WHEEL" - !sudo apt-get install libomp5 - update.join() + !curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py + !python pytorch-xla-env-setup.py --version nightly --apt-packages libomp5 libopenblas-dev In distributed training (multiple GPUs and multiple TPU cores) each GPU or TPU core will run a copy of this program. This means that without taking any care you will download the dataset N times which @@ -609,7 +567,7 @@ Now we can train the LightningModule on a TPU without doing anything else! .. code-block:: python model = LitMNIST() - trainer = Trainer(num_tpu_cores=8) + trainer = Trainer(tpu_cores=8) trainer.fit(model) You'll now see the TPU cores booting up. @@ -696,7 +654,7 @@ while checking the validation set. from pytorch_lightning import Trainer model = LitMNIST() - trainer = Trainer(num_tpu_cores=8) + trainer = Trainer(tpu_cores=8) trainer.fit(model) You may have noticed the words `Validation sanity check` logged. This is because Lightning runs 5 batches @@ -747,7 +705,7 @@ Once you train your model simply call `.test()`. from pytorch_lightning import Trainer model = LitMNIST() - trainer = Trainer(num_tpu_cores=8) + trainer = Trainer(tpu_cores=8) trainer.fit(model) # run test set @@ -769,7 +727,7 @@ You can also run the test from a saved lightning model .. code-block:: python model = LitMNIST.load_from_checkpoint(PATH) - trainer = Trainer(num_tpu_cores=8) + trainer = Trainer(tpu_cores=8) trainer.test(model) .. note:: Lightning disables gradients, puts model in eval mode and does everything needed for testing. diff --git a/docs/source/multi_gpu.rst b/docs/source/multi_gpu.rst index a094a636..21655227 100644 --- a/docs/source/multi_gpu.rst +++ b/docs/source/multi_gpu.rst @@ -130,7 +130,7 @@ Lightning allows multiple ways of training - DistributedDataParallel (`distributed_backend='ddp'`) (multiple-gpus across many machines). - DistributedDataParallel2 (`distributed_backend='ddp2'`) (dp in a machine, ddp across machines). - Horovod (`distributed_backend='horovod'`) (multi-machine, multi-gpu, configured at runtime) -- TPUs (`num_tpu_cores=8|x`) (tpu or TPU pod) +- TPUs (`tpu_cores=8|x`) (tpu or TPU pod) .. note:: If you request multiple GPUs without setting a mode, ddp will be automatically used. diff --git a/docs/source/new-project.rst b/docs/source/new-project.rst index 24b11412..e1efb898 100644 --- a/docs/source/new-project.rst +++ b/docs/source/new-project.rst @@ -189,7 +189,7 @@ However, this time you need to specifically call test (this is done so you don't # OPTION 2: # test after loading weights model = LitModel.load_from_checkpoint(PATH) - trainer = Trainer(num_tpu_cores=1) + trainer = Trainer(tpu_cores=1) trainer.test() Again, under the hood, lightning does the following in (pseudocode): @@ -236,7 +236,7 @@ Without changing a SINGLE line of your code, you can now do the following with t # train on TPUs using 16 bit precision with early stopping # using only half the training data and checking validation every quarter of a training epoch trainer = Trainer( - nb_tpu_cores=8, + tpu_cores=8, precision=16, early_stop_checkpoint=True, train_percent_check=0.5, diff --git a/docs/source/tpu.rst b/docs/source/tpu.rst index b2fb6e85..774af763 100644 --- a/docs/source/tpu.rst +++ b/docs/source/tpu.rst @@ -1,8 +1,8 @@ TPU support =========== -Lightning supports running on TPUs. At this moment, TPUs are only available -on Google Cloud (GCP). For more information on TPUs +Lightning supports running on TPUs. At this moment, TPUs are available +on Google Cloud (GCP), Google Colab and Kaggle Environments. For more information on TPUs `watch this video `_. --------------- @@ -31,6 +31,7 @@ To access TPUs there are two main ways. 1. Using google colab. 2. Using Google Cloud (GCP). +3. Using Kaggle. --------------- @@ -51,50 +52,10 @@ To get a TPU on colab, follow these steps: 4. Next, insert this code into the first cell and execute. This will install the xla library that interfaces between PyTorch and the TPU. - .. code-block:: python - - import collections - from datetime import datetime, timedelta - import os - import requests - import threading - - _VersionConfig = collections.namedtuple('_VersionConfig', 'wheels,server') - VERSION = "xrt==1.15.0" #@param ["xrt==1.15.0", "torch_xla==nightly"] - CONFIG = { - 'xrt==1.15.0': _VersionConfig('1.15', '1.15.0'), - 'torch_xla==nightly': _VersionConfig('nightly', 'XRT-dev{}'.format( - (datetime.today() - timedelta(1)).strftime('%Y%m%d'))), - }[VERSION] - DIST_BUCKET = 'gs://tpu-pytorch/wheels' - TORCH_WHEEL = 'torch-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels) - TORCH_XLA_WHEEL = 'torch_xla-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels) - TORCHVISION_WHEEL = 'torchvision-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels) - - # Update TPU XRT version - def update_server_xrt(): - print('Updating server-side XRT to {} ...'.format(CONFIG.server)) - url = 'http://{TPU_ADDRESS}:8475/requestversion/{XRT_VERSION}'.format( - TPU_ADDRESS=os.environ['COLAB_TPU_ADDR'].split(':')[0], - XRT_VERSION=CONFIG.server, - ) - print('Done updating server-side XRT: {}'.format(requests.post(url))) - - update = threading.Thread(target=update_server_xrt) - update.start() - .. code-block:: - # Install Colab TPU compat PyTorch/TPU wheels and dependencies - !pip uninstall -y torch torchvision - !gsutil cp "$DIST_BUCKET/$TORCH_WHEEL" . - !gsutil cp "$DIST_BUCKET/$TORCH_XLA_WHEEL" . - !gsutil cp "$DIST_BUCKET/$TORCHVISION_WHEEL" . - !pip install "$TORCH_WHEEL" - !pip install "$TORCH_XLA_WHEEL" - !pip install "$TORCHVISION_WHEEL" - !sudo apt-get install libomp5 - update.join() + !curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py + !python pytorch-xla-env-setup.py --version nightly --apt-packages libomp5 libopenblas-dev 5. Once the above is done, install PyTorch Lightning (v 0.7.0+). @@ -156,13 +117,23 @@ To use a full TPU pod skip to the TPU pod section. import pytorch_lightning as pl my_model = MyLightningModule() - trainer = pl.Trainer(num_tpu_cores=8) + trainer = pl.Trainer(tpu_cores=8) trainer.fit(my_model) That's it! Your model will train on all 8 TPU cores. --------------- +Single TPU core training +---------------------------- +Lightning supports training on a single TPU core. Just pass the TPU core ID [1-8] in a list. + +.. code-block:: python + + trainer = pl.Trainer(tpu_cores=[1]) + +--------------- + Distributed Backend with TPU ---------------------------- The ```distributed_backend``` option used for GPUs does not apply to TPUs. @@ -195,7 +166,7 @@ set the 16-bit flag. import pytorch_lightning as pl my_model = MyLightningModule() - trainer = pl.Trainer(num_tpu_cores=8, precision=16) + trainer = pl.Trainer(tpu_cores=8, precision=16) trainer.fit(my_model) Under the hood the xla library will use the `bfloat16 type `_. diff --git a/pytorch_lightning/trainer/__init__.py b/pytorch_lightning/trainer/__init__.py index 98d2d8a7..3465d49b 100644 --- a/pytorch_lightning/trainer/__init__.py +++ b/pytorch_lightning/trainer/__init__.py @@ -598,7 +598,22 @@ nb_sanity_val_steps: num_tpu_cores ^^^^^^^^^^^^^ -How many TPU cores to train on (1 or 8). +.. warning:: .. deprecated:: 0.7.6 + + Use `tpu_cores` instead. Will remove 0.9.0. + +Example:: + + python -m torch_xla.distributed.xla_dist + --tpu=$TPU_POD_NAME + --conda-env=torch-xla-nightly + --env=XLA_USE_BF16=1 + -- python your_trainer_file.py + +tpu_cores +^^^^^^^^^ +- How many TPU cores to train on (1 or 8). +- Which TPU core to train on [1-8] A single TPU v2 or v3 has 8 cores. A TPU pod has up to 2048 cores. A slice of a POD means you get as many cores @@ -615,21 +630,21 @@ Example:: # your_trainer_file.py # default used by the Trainer (ie: train on CPU) - trainer = Trainer(num_tpu_cores=None) + trainer = Trainer(tpu_cores=None) # int: train on a single core - trainer = Trainer(num_tpu_cores=1) + trainer = Trainer(tpu_cores=1) + + # list: train on a single selected core + trainer = Trainer(tpu_cores=[2]) # int: train on all cores few cores - trainer = Trainer(num_tpu_cores=8) + trainer = Trainer(tpu_cores=8) # for 8+ cores must submit via xla script with # a max of 8 cores specified. The XLA script # will duplicate script onto each TPU in the POD - trainer = Trainer(num_tpu_cores=8) - - # -1: train on all available TPUs - trainer = Trainer(num_tpu_cores=-1) + trainer = Trainer(tpu_cores=8) To train on more than 8 cores (ie: a POD), submit this script using the xla_dist script. diff --git a/pytorch_lightning/trainer/deprecated_api.py b/pytorch_lightning/trainer/deprecated_api.py index 2705c4f1..5b615eba 100644 --- a/pytorch_lightning/trainer/deprecated_api.py +++ b/pytorch_lightning/trainer/deprecated_api.py @@ -135,3 +135,9 @@ class TrainerDeprecatedAPITillVer0_9(ABC): rank_zero_warn("`training_tqdm_dict` was renamed to `progress_bar_dict` in v0.7.3" " and this method will be removed in v0.9.0", DeprecationWarning) return self.progress_bar_dict + + @property + def num_tpu_cores(self): + """Back compatibility, will be removed in v0.9.0""" + rank_zero_warn("Argument `num_tpu_cores` is now set by `tpu_cores` since v0.7.6" + " and this argument will be removed in v0.9.0", DeprecationWarning) diff --git a/pytorch_lightning/trainer/distrib_parts.py b/pytorch_lightning/trainer/distrib_parts.py index 31865fc6..a08c7a1c 100644 --- a/pytorch_lightning/trainer/distrib_parts.py +++ b/pytorch_lightning/trainer/distrib_parts.py @@ -389,7 +389,6 @@ class TrainerDPMixin(ABC): root_gpu: ... amp_level: str precision: ... - current_tpu_idx: ... proc_rank: int tpu_local_core_rank: int tpu_global_core_rank: int @@ -398,6 +397,7 @@ class TrainerDPMixin(ABC): data_parallel_device_ids: ... logger: Union[LightningLoggerBase, bool] progress_bar_callback: ... + tpu_id: int @property @abstractmethod @@ -442,7 +442,8 @@ class TrainerDPMixin(ABC): if device == 'tpu' and XLA_AVAILABLE: # base case: object can be directly moved using `to` if callable(getattr(batch, 'to', None)): - return batch.to(xm.xla_device()) + xla_device = xm.xla_device(self.tpu_id) if self.tpu_id is not None else xm.xla_device() + return batch.to(xla_device) if device == 'gpu': # base case: object can be directly moved using `cuda` or `to` @@ -501,7 +502,8 @@ class TrainerDPMixin(ABC): def tpu_train(self, tpu_core_idx, model): # put model on tpu - model.to(xm.xla_device()) + self._device = xm.xla_device(self.tpu_id) if self.tpu_id is not None else xm.xla_device() + model.to(self._device) # get the appropriate tpu ranks self.tpu_local_core_rank = xm.get_local_ordinal() @@ -511,8 +513,6 @@ class TrainerDPMixin(ABC): if self.tpu_global_core_rank != 0 and self.progress_bar_callback is not None: self.progress_bar_callback.disable() - # track current tpu - self.current_tpu_idx = tpu_core_idx self.proc_rank = self.tpu_local_core_rank rank_zero_only.rank = self.proc_rank diff --git a/pytorch_lightning/trainer/evaluation_loop.py b/pytorch_lightning/trainer/evaluation_loop.py index b056b85a..58d16632 100644 --- a/pytorch_lightning/trainer/evaluation_loop.py +++ b/pytorch_lightning/trainer/evaluation_loop.py @@ -174,6 +174,7 @@ class TrainerEvaluationLoopMixin(ABC): val_dataloaders: DataLoader use_tpu: bool reload_dataloaders_every_epoch: ... + tpu_id: int # Callback system on_validation_batch_start: Callable @@ -248,7 +249,7 @@ class TrainerEvaluationLoopMixin(ABC): dl_outputs = [] # on TPU we have to wrap it under the ParallelLoader - if self.use_tpu: + if self.use_tpu and self.tpu_id is None: device = xm.xla_device() dataloader = xla_pl.ParallelLoader(dataloader, [device]) dataloader = dataloader.per_device_loader(device) diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py index bed2aa05..e4eae0de 100644 --- a/pytorch_lightning/trainer/trainer.py +++ b/pytorch_lightning/trainer/trainer.py @@ -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: @@ -82,7 +81,7 @@ class Trainer( 'gradient_clip', 'nb_gpu_nodes', 'max_nb_epochs', 'min_nb_epochs', 'add_row_log_interval', 'nb_sanity_val_steps', 'tng_tqdm_dic', ) - DEPRECATED_IN_0_9 = ('use_amp', 'show_progress_bar', 'training_tqdm_dict') + DEPRECATED_IN_0_9 = ('use_amp', 'show_progress_bar', 'training_tqdm_dict', 'num_tpu_cores') def __init__( self, @@ -97,7 +96,7 @@ class Trainer( num_processes: int = 1, gpus: Optional[Union[List[int], str, int]] = None, auto_select_gpus: bool = False, - num_tpu_cores: Optional[int] = None, + tpu_cores: Optional[Union[List[int], int]] = None, log_gpu_memory: Optional[str] = None, progress_bar_refresh_rate: int = 1, overfit_pct: float = 0.0, @@ -133,6 +132,7 @@ class Trainer( progress_bar_callback: Optional[Union[ProgressBarBase, bool]] = True, terminate_on_nan: bool = False, auto_scale_batch_size: Union[str, bool] = False, + num_tpu_cores: Optional[int] = None, # backward compatible, todo: remove in v0.9.0 amp_level: str = 'O1', # backward compatible, todo: remove in v0.8.0 default_save_path=None, # backward compatible, todo: remove in v0.8.0 gradient_clip=None, # backward compatible, todo: remove in v0.8.0 @@ -188,7 +188,10 @@ class Trainer( GPUs are configured to be in "exclusive mode", such that only one process at a time can access them. - num_tpu_cores: How many TPU cores to train on (1 or 8). + tpu_cores: How many TPU cores to train on (1 or 8) / Single TPU to train on [1] + + num_tpu_cores: How many TPU cores to train on (1 or 8) + .. warning:: .. deprecated:: 0.7.6. Will remove 0.9.0. log_gpu_memory: None, 'min_max', 'all'. Might slow performance @@ -342,9 +345,19 @@ class Trainer( self.on_gpu = True if (gpus and torch.cuda.is_available()) else False # tpu config - self.on_tpu = num_tpu_cores is not None - self.num_tpu_cores = num_tpu_cores - assert num_tpu_cores in [1, 8, None], 'num_tpu_cores can only be 1 or 8' + if num_tpu_cores is not None: + rank_zero_warn("Argument `num_tpu_cores` is now set by `tpu_cores` since v0.7.6" + " and this argument will be removed in v0.9.0", DeprecationWarning) + + if tpu_cores is None: + tpu_cores = num_tpu_cores + self.on_tpu = tpu_cores is not None + self.tpu_cores = tpu_cores + assert self.tpu_cores in (1, 8, None) or ( + isinstance(self.tpu_cores, (list, tuple, set)) and len(self.tpu_cores) == 1 + ), '`tpu_cores` can only be 1, 8 or [<1-8>]' + + self.tpu_id = tpu_cores[0] if isinstance(tpu_cores, list) else None if num_processes != 1 and distributed_backend != "ddp_cpu": rank_zero_warn("num_processes is only used for distributed_backend=\"ddp_cpu\". Ignoring it.") @@ -477,7 +490,6 @@ class Trainer( # override dist backend when using tpus if self.on_tpu: self.init_tpu() - self.current_tpu_idx = None # init flags for SLURM+ddp to work self.proc_rank = 0 @@ -858,7 +870,7 @@ class Trainer( self.single_gpu_train(model) elif self.use_tpu: # pragma: no-cover - log.info(f'training on {self.num_tpu_cores} TPU cores') + log.info(f'training on {self.tpu_cores} TPU cores') # COLAB_GPU is an env var available by default in Colab environments. start_method = 'fork' if self.on_colab_kaggle else 'spawn' @@ -867,7 +879,10 @@ class Trainer( self.model = model # train - xmp.spawn(self.tpu_train, args=(model,), nprocs=self.num_tpu_cores, start_method=start_method) + if self.tpu_id is not None: + self.tpu_train(self.tpu_id, model) + else: + xmp.spawn(self.tpu_train, args=(model,), nprocs=self.tpu_cores, start_method=start_method) # load weights if not interrupted self.load_spawn_weights(model) diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py index 2eb033c0..cbe11864 100644 --- a/pytorch_lightning/trainer/training_loop.py +++ b/pytorch_lightning/trainer/training_loop.py @@ -231,6 +231,7 @@ class TrainerTrainLoopMixin(ABC): total_batch_idx: int checkpoint_callback: ... terminate_on_nan: bool + tpu_id: int # Callback system callbacks: List[Callback] @@ -393,7 +394,7 @@ class TrainerTrainLoopMixin(ABC): train_dataloader = self.train_dataloader # on TPU we have to wrap it under the ParallelLoader - if self.use_tpu: + if self.use_tpu and self.tpu_id is None: device = xm.xla_device() train_dataloader = xla_pl.ParallelLoader(train_dataloader, [device]) train_dataloader = train_dataloader.per_device_loader(device) diff --git a/tests/test_deprecated.py b/tests/test_deprecated.py index 54113b1e..df541b62 100644 --- a/tests/test_deprecated.py +++ b/tests/test_deprecated.py @@ -96,6 +96,9 @@ def test_tbd_remove_in_v0_9_0_trainer(): trainer = Trainer(progress_bar_refresh_rate=50, show_progress_bar=False) assert getattr(trainer, 'show_progress_bar') + with pytest.deprecated_call(match='v0.9.0'): + _ = Trainer(num_tpu_cores=8) + def test_tbd_remove_in_v0_9_0_module_imports(): _soft_unimport_module("pytorch_lightning.core.decorators")