From 2ab2f7d08df4e4f913e229caf92bbd92f31f6f93 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= Date: Thu, 16 Apr 2020 18:04:55 +0200 Subject: [PATCH] Improved docs for pytorch_lightning.core (continued) (#1483) * improved docs for core update links add references to hooks lifecycle wip continue with __init__.py improve docs for memory.py improve docs for saving.py simpler links fix formatting * move hooks lifecycle to top of file * fix doctest import problem * add missing hook in lifecycle --- docs/source/hooks.rst | 96 +++++---- pytorch_lightning/core/__init__.py | 327 ++++++++++++++--------------- pytorch_lightning/core/hooks.py | 92 ++++---- pytorch_lightning/core/memory.py | 49 +++-- pytorch_lightning/core/saving.py | 26 ++- 5 files changed, 320 insertions(+), 270 deletions(-) diff --git a/docs/source/hooks.rst b/docs/source/hooks.rst index 37c138fe..18bfb028 100644 --- a/docs/source/hooks.rst +++ b/docs/source/hooks.rst @@ -1,8 +1,17 @@ -Hooks -===== +Model Hooks +=========== + +There are cases when you might want to do something different at different parts of the training/validation loop. +To enable a hook, simply override the method in your LightningModule and the trainer will call it at the correct time. + +**Contributing** If there's a hook you'd like to add, simply: + +1. Fork `PyTorchLightning `_. + +2. Add the hook to :class:`pytorch_lightning.core.hooks.ModelHooks`. + +3. Add it in the correct place in :mod:`pytorch_lightning.trainer` where it should be called. -.. automodule:: pytorch_lightning.core.hooks - :noindex: Hooks lifecycle --------------- @@ -10,50 +19,59 @@ Hooks lifecycle Training set-up ^^^^^^^^^^^^^^^ -- init_ddp_connection -- init_optimizers -- configure_apex -- configure_ddp -- train_dataloader -- test_dataloaders -- val_dataloaders -- summarize -- restore_weights +- :meth:`~pytorch_lightning.core.lightning.LightningModule.init_ddp_connection` +- :meth:`~pytorch_lightning.trainer.optimizers.TrainerOptimizersMixin.init_optimizers` +- :meth:`~pytorch_lightning.core.lightning.LightningModule.configure_apex` +- :meth:`~pytorch_lightning.core.lightning.LightningModule.configure_ddp` +- :meth:`~pytorch_lightning.core.lightning.LightningModule.train_dataloader` +- :meth:`~pytorch_lightning.core.lightning.LightningModule.test_dataloader` +- :meth:`~pytorch_lightning.core.lightning.LightningModule.val_dataloader` +- :meth:`~pytorch_lightning.core.lightning.LightningModule.summarize` +- :meth:`~pytorch_lightning.trainer.training_io.TrainerIOMixin.restore_weights` Training loop ^^^^^^^^^^^^^ -- on_epoch_start -- on_batch_start -- tbptt_split_batch -- training_step -- training_step_end (optional) -- backward -- on_after_backward -- optimizer.step() -- on_batch_end -- on_epoch_end +- :meth:`~pytorch_lightning.core.hooks.ModelHooks.on_epoch_start` +- :meth:`~pytorch_lightning.core.hooks.ModelHooks.on_batch_start` +- :meth:`~pytorch_lightning.core.lightning.LightningModule.tbptt_split_batch` +- :meth:`~pytorch_lightning.core.lightning.LightningModule.training_step` +- :meth:`~pytorch_lightning.core.lightning.LightningModule.training_step_end` (optional) +- :meth:`~pytorch_lightning.core.hooks.ModelHooks.on_before_zero_grad` +- :meth:`~pytorch_lightning.core.hooks.ModelHooks.backward` +- :meth:`~pytorch_lightning.core.hooks.ModelHooks.on_after_backward` +- ``optimizer.step()`` +- :meth:`~pytorch_lightning.core.hooks.ModelHooks.on_batch_end` +- :meth:`~pytorch_lightning.core.lightning.LightningModule.training_epoch_end` +- :meth:`~pytorch_lightning.core.hooks.ModelHooks.on_epoch_end` Validation loop ^^^^^^^^^^^^^^^ -- model.zero_grad() -- model.eval() -- torch.set_grad_enabled(False) -- validation_step -- validation_end -- model.train() -- torch.set_grad_enabled(True) -- on_post_performance_check +- ``model.zero_grad()`` +- ``model.eval()`` +- ``torch.set_grad_enabled(False)`` +- :meth:`~pytorch_lightning.core.lightning.LightningModule.validation_step` +- :meth:`~pytorch_lightning.core.lightning.LightningModule.validation_step_end` +- :meth:`~pytorch_lightning.core.lightning.LightningModule.validation_epoch_end` +- ``model.train()`` +- ``torch.set_grad_enabled(True)`` +- :meth:`~pytorch_lightning.core.hooks.ModelHooks.on_post_performance_check` Test loop ^^^^^^^^^ -- model.zero_grad() -- model.eval() -- torch.set_grad_enabled(False) -- test_step -- test_end -- model.train() -- torch.set_grad_enabled(True) -- on_post_performance_check \ No newline at end of file +- ``model.zero_grad()`` +- ``model.eval()`` +- ``torch.set_grad_enabled(False)`` +- :meth:`~pytorch_lightning.core.lightning.LightningModule.test_step` +- :meth:`~pytorch_lightning.core.lightning.LightningModule.test_step_end` +- :meth:`~pytorch_lightning.core.lightning.LightningModule.test_epoch_end` +- ``model.train()`` +- ``torch.set_grad_enabled(True)`` +- :meth:`~pytorch_lightning.core.hooks.ModelHooks.on_post_performance_check` + + + +.. automodule:: pytorch_lightning.core.hooks + :noindex: \ No newline at end of file diff --git a/pytorch_lightning/core/__init__.py b/pytorch_lightning/core/__init__.py index c30f3b5e..83fff7d8 100644 --- a/pytorch_lightning/core/__init__.py +++ b/pytorch_lightning/core/__init__.py @@ -1,5 +1,5 @@ """ -A LightningModule organizes your PyTorch code into the following sections: +A :class:`~LightningModule` organizes your PyTorch code into the following sections: .. figure:: /_images/lightning_module/pt_to_pl.png :alt: Convert from PyTorch to Lightning @@ -7,48 +7,53 @@ A LightningModule organizes your PyTorch code into the following sections: Notice a few things. - 1. It's the SAME code. - 2. The PyTorch code IS NOT abstracted - just organized. - 3. All the other code that not in the LightningModule has been automated for you by the trainer - .. code-block:: python +1. It's the SAME code. +2. The PyTorch code IS NOT abstracted - just organized. +3. All the other code that's not in the :class:`~LightningModule` + has been automated for you by the trainer. - net = Net() - trainer = Trainer() - trainer.fit(net) + .. code-block:: python - 4. There are no .cuda() or .to() calls... Lightning does these for you. - .. code-block:: python + net = Net() + trainer = Trainer() + trainer.fit(net) - # don't do in lightning - x = torch.Tensor(2, 3) - x = x.cuda() - x = x.to(device) +4. There are no .cuda() or .to() calls... Lightning does these for you. - # do this instead - x = x # leave it alone! + .. code-block:: python - # or to init a new tensor - new_x = torch.Tensor(2, 3) - new_x = new_x.type_as(x.type()) + # don't do in lightning + x = torch.Tensor(2, 3) + x = x.cuda() + x = x.to(device) - 5. There are no samplers for distributed, Lightning also does this for you. - .. code-block:: python + # do this instead + x = x # leave it alone! - # Don't do in Lightning... - data = MNIST(...) - sampler = DistributedSampler(data) - DataLoader(data, sampler=sampler) + # or to init a new tensor + new_x = torch.Tensor(2, 3) + new_x = new_x.type_as(x.type()) - # do this instead - data = MNIST(...) - DataLoader(data) +5. There are no samplers for distributed, Lightning also does this for you. - 6. A LightingModule is a torch.nn.Module but with added functionality. Use it as such! - .. code-block:: python + .. code-block:: python - net = Net.load_from_checkpoint(PATH) - net.freeze() - out = net(x) + # Don't do in Lightning... + data = MNIST(...) + sampler = DistributedSampler(data) + DataLoader(data, sampler=sampler) + + # do this instead + data = MNIST(...) + DataLoader(data) + +6. A :class:`~LightningModule` is a :class:`torch.nn.Module` but with added functionality. Use it as such! + + .. code-block:: python + + net = Net.load_from_checkpoint(PATH) + net.freeze() + out = net(x) Thus, to use Lightning, you just need to organize your code which takes about 30 minutes, (and let's be real, you probably should do anyhow). @@ -62,35 +67,27 @@ Here are the only required methods. .. code-block:: python - import os - import torch - from torch.nn import functional as F - from torch.utils.data import DataLoader - from torchvision.datasets import MNIST - import torchvision.transforms as transforms - - import pytorch_lightning as pl - - class LitModel(pl.LightningModule): - - def __init__(self): - super().__init__() - self.l1 = torch.nn.Linear(28 * 28, 10) - - def forward(self, x): - return torch.relu(self.l1(x.view(x.size(0), -1))) - - def training_step(self, batch, batch_idx): - x, y = batch - y_hat = self(x) - return {'loss': F.cross_entropy(y_hat, y)} - - def train_dataloader(self): - return DataLoader(MNIST(os.getcwd(), train=True, download=True, - transform=transforms.ToTensor()), batch_size=32) - - def configure_optimizers(self): - return torch.optim.Adam(self.parameters(), lr=0.02) + >>> import pytorch_lightning as pl + >>> class LitModel(pl.LightningModule): + ... + ... def __init__(self): + ... super().__init__() + ... self.l1 = torch.nn.Linear(28 * 28, 10) + ... + ... def forward(self, x): + ... return torch.relu(self.l1(x.view(x.size(0), -1))) + ... + ... def training_step(self, batch, batch_idx): + ... x, y = batch + ... y_hat = self(x) + ... return {'loss': F.cross_entropy(y_hat, y)} + ... + ... def train_dataloader(self): + ... return DataLoader(MNIST(os.getcwd(), train=True, download=True, + ... transform=transforms.ToTensor()), batch_size=32) + ... + ... def configure_optimizers(self): + ... return torch.optim.Adam(self.parameters(), lr=0.02) Which you can train by doing: @@ -109,11 +106,11 @@ Training loop structure The general pattern is that each loop (training, validation, test loop) has 3 methods: -- ``` ___step ``` -- ``` ___step_end ``` -- ``` ___epoch_end``` +- ``___step`` +- ``___step_end`` +- ``___epoch_end`` -To show how lightning calls these, let's use the validation loop as an example +To show how Lightning calls these, let's use the validation loop as an example: .. code-block:: python @@ -127,10 +124,8 @@ To show how lightning calls these, let's use the validation loop as an example # like calculate validation set accuracy or loss validation_epoch_end(val_outs) -if we use dp or ddp2 mode, we can also define the ```XXX_step_end``` method to operate -on all parts of the batch - -.. code-block:: python +If we use dp or ddp2 mode, we can also define the ``XXX_step_end`` method to operate +on all parts of the batch:: val_outs = [] for val_batch in val_data: @@ -147,47 +142,43 @@ on all parts of the batch # like calculate validation set accuracy or loss validation_epoch_end(val_outs) -.. note:: ```training_step_end``` is not available yet but coming in the next release. Add validation loop ^^^^^^^^^^^^^^^^^^^ -Thus, if we wanted to add a validation loop you would add this to your LightningModule +Thus, if we wanted to add a validation loop you would add this to your +:class:`~LightningModule`: -.. code-block:: python - - class LitModel(pl.LightningModule): - def validation_step(self, batch, batch_idx): - x, y = batch - y_hat = self(x) - return {'val_loss': F.cross_entropy(y_hat, y)} - - def validation_epoch_end(self, outputs): - val_loss_mean = torch.stack([x['val_loss'] for x in outputs]).mean() - return {'val_loss': val_loss_mean} - - def val_dataloader(self): - # can also return a list of val dataloaders - return DataLoader(...) + >>> class LitModel(pl.LightningModule): + ... def validation_step(self, batch, batch_idx): + ... x, y = batch + ... y_hat = self(x) + ... return {'val_loss': F.cross_entropy(y_hat, y)} + ... + ... def validation_epoch_end(self, outputs): + ... val_loss_mean = torch.stack([x['val_loss'] for x in outputs]).mean() + ... return {'val_loss': val_loss_mean} + ... + ... def val_dataloader(self): + ... # can also return a list of val dataloaders + ... return DataLoader(...) Add test loop ^^^^^^^^^^^^^ -.. code-block:: python - - class LitModel(pl.LightningModule): - def test_step(self, batch, batch_idx): - x, y = batch - y_hat = self(x) - return {'test_loss': F.cross_entropy(y_hat, y)} - - def test_epoch_end(self, outputs): - test_loss_mean = torch.stack([x['test_loss'] for x in outputs]).mean() - return {'test_loss': test_loss_mean} - - def test_dataloader(self): - # can also return a list of test dataloaders - return DataLoader(...) + >>> class LitModel(pl.LightningModule): + ... def test_step(self, batch, batch_idx): + ... x, y = batch + ... y_hat = self(x) + ... return {'test_loss': F.cross_entropy(y_hat, y)} + ... + ... def test_epoch_end(self, outputs): + ... test_loss_mean = torch.stack([x['test_loss'] for x in outputs]).mean() + ... return {'test_loss': test_loss_mean} + ... + ... def test_dataloader(self): + ... # can also return a list of test dataloaders + ... return DataLoader(...) However, the test loop won't ever be called automatically to make sure you don't run your test data by accident. Instead you have to explicitly call: @@ -208,41 +199,43 @@ don't run your test data by accident. Instead you have to explicitly call: Training_step_end method ------------------------ -When using dataParallel or distributedDataParallel2, the training_step +When using :class:`~pytorch_lightning.overrides.data_parallel.LightningDataParallel` or +:class:`~pytorch_lightning.overrides.data_parallel.LightningDistributedDataParallel`, the +:meth:`~LightningModule.training_step` will be operating on a portion of the batch. This is normally ok but in special cases like calculating NCE loss using negative samples, we might want to perform a softmax across all samples in the batch. -For these types of situations, each loop has an additional ```__step_end``` method -which allows you to operate on the pieces of the batch +For these types of situations, each loop has an additional ``__step_end`` method +which allows you to operate on the pieces of the batch: .. code-block:: python - training_outs = [] - for train_batch in train_data: - # dp, ddp2 splits the batch - sub_batches = split_batches_for_dp(batch) + training_outs = [] + for train_batch in train_data: + # dp, ddp2 splits the batch + sub_batches = split_batches_for_dp(batch) - # run training_step on each piece of the batch - batch_parts_outputs = [training_step(sub_batch) for sub_batch in sub_batches] + # run training_step on each piece of the batch + batch_parts_outputs = [training_step(sub_batch) for sub_batch in sub_batches] - # do softmax with all pieces - out = training_step_end(batch_parts_outputs) - training_outs.append(out) + # do softmax with all pieces + out = training_step_end(batch_parts_outputs) + training_outs.append(out) - # do something with the outputs for all batches - # like calculate validation set accuracy or loss - training_epoch_end(val_outs) + # do something with the outputs for all batches + # like calculate validation set accuracy or loss + training_epoch_end(val_outs) ---------- Remove cuda calls ----------------- -In a LightningModule, all calls to ```.cuda()``` -and ```.to(device)``` should be removed. Lightning will do these +In a :class:`~LightningModule`, all calls to ``.cuda()`` +and ``.to(device)`` should be removed. Lightning will do these automatically. This will allow your code to work on CPUs, TPUs and GPUs. -When you init a new tensor in your code, just use type_as +When you init a new tensor in your code, just use :meth:`~torch.Tensor.type_as`: .. code-block:: python @@ -259,70 +252,74 @@ Data preparation ---------------- Data preparation in PyTorch follows 5 steps: - 1. Download - 2. Clean and (maybe) save to disk - 3. Load inside dataset - 4. Apply transforms (rotate, tokenize, etc...) - 5. Wrap inside a dataloader +1. Download +2. Clean and (maybe) save to disk +3. Load inside :class:`~torch.utils.data.Dataset` +4. Apply transforms (rotate, tokenize, etc...) +5. Wrap inside a :class:`~torch.utils.data.DataLoader` When working in distributed settings, steps 1 and 2 have to be done from a single GPU, otherwise you will overwrite these files from -every GPU. The lightningModule has the ```prepare_data``` method to -allow for this +every GPU. The :class:`~LightningModule` has the +:class:`~LightningModule.prepare_data` method to +allow for this: -.. code-block:: python + >>> class LitModel(pl.LightningModule): + ... def prepare_data(self): + ... # download + ... mnist_train = MNIST(os.getcwd(), train=True, download=True, + ... transform=transforms.ToTensor()) + ... mnist_test = MNIST(os.getcwd(), train=False, download=True, + ... transform=transforms.ToTensor()) + ... + ... # train/val split + ... mnist_train, mnist_val = random_split(mnist_train, [55000, 5000]) + ... + ... # assign to use in dataloaders + ... self.train_dataset = mnist_train + ... self.val_dataset = mnist_val + ... self.test_dataset = mnist_test + ... + ... def train_dataloader(self): + ... return DataLoader(self.train_dataset, batch_size=64) + ... + ... def val_dataloader(self): + ... return DataLoader(self.mnist_val, batch_size=64) + ... + ... def test_dataloader(self): + ... return DataLoader(self.mnist_test, batch_size=64) - def prepare_data(self): - # download - mnist_train = MNIST(os.getcwd(), train=True, download=True, - transform=transforms.ToTensor()) - mnist_test = MNIST(os.getcwd(), train=False, download=True, - transform=transforms.ToTensor()) +Note: + :meth:`~LightningModule.prepare_data` is called once. - # train/val split - mnist_train, mnist_val = random_split(mnist_train, [55000, 5000]) +Note: + Do anything with data that needs to happen ONLY once here, like download, tokenize, etc... - # assign to use in dataloaders - self.train_dataset = mnist_train - self.val_dataset = mnist_val - self.test_dataset = mnist_test - - def train_dataloader(self): - return DataLoader(self.train_dataset, batch_size=64) - - def val_dataloader(self): - return DataLoader(self.mnist_val, batch_size=64) - - def test_dataloader(self): - return DataLoader(self.mnist_test, batch_size=64) - -.. note:: ``prepare_data`` is called once. - -.. note:: Do anything with data that needs to happen ONLY once here, like download, tokenize, etc... Lifecycle --------- -The methods in the LightningModule are called in this order: +The methods in the :class:`~LightningModule` are called in this order: - 1. ```__init__``` - 2. ```prepare_data``` - 3. ```configure_optimizers``` - 4. ```train_dataloader``` +1. :meth:`~LightningModule.__init__` +2. :meth:`~LightningModule.prepare_data` +3. :meth:`~LightningModule.configure_optimizers` +4. :meth:`~LightningModule.train_dataloader` - If you define a validation loop then +If you define a validation loop then - 5. ```val_dataloader``` +5. :meth:`~LightningModule.val_dataloader` - And if you define a test loop: +And if you define a test loop: - 6. ```test_dataloader``` +6. :meth:`~LightningModule.test_dataloader` -.. note:: ``test_dataloader`` is only called with ``.test()`` +Note: + :meth:`~LightningModule.test_dataloader` is only called with ``.test()`` In every epoch, the loop methods are called in this frequency: -1. ```validation_step``` called every batch -2. ```validation_epoch_end``` called every epoch +1. :meth:`~LightningModule.validation_step` called every batch +2. :meth:`~LightningModule.validation_epoch_end` called every epoch Live demo --------- diff --git a/pytorch_lightning/core/hooks.py b/pytorch_lightning/core/hooks.py index 9c4b999d..67b67ec1 100644 --- a/pytorch_lightning/core/hooks.py +++ b/pytorch_lightning/core/hooks.py @@ -1,19 +1,3 @@ -""" -Model Hooks -=========== - -There are cases when you might want to do something different at different parts of the training/validation loop. - To enable a hook, simply override the method in your LightningModule and the trainer will call it at the correct time. - -**Contributing** If there's a hook you'd like to add, simply: - -1. Fork PyTorchLightning. - -2. Add the hook :py:mod:`pytorch_lightning.base_module.hooks.py`. - -3. Add the correct place in the :py:mod:`pytorch_lightning.models.trainer` where it should be called. - -""" from typing import Any import torch @@ -30,75 +14,93 @@ else: class ModelHooks(torch.nn.Module): + # TODO: remove in v0.9.0 def on_sanity_check_start(self): """ - Called before starting evaluate - .. warning:: will be deprecated. - :return: + Called before starting evaluation. + + Warning: + Deprecated. Will be removed in v0.9.0. """ def on_train_start(self) -> None: - """Called at the beginning of training before sanity check + """ + Called at the beginning of training before sanity check. """ # do something at the start of training def on_train_end(self) -> None: """ - Called at the end of training before logger experiment is closed + Called at the end of training before logger experiment is closed. """ # do something at the end of training def on_batch_start(self, batch: Any) -> None: - """Called in the training loop before anything happens for that batch. + """ + Called in the training loop before anything happens for that batch. If you return -1 here, you will skip training for the rest of the current epoch. - :param batch: + Args: + batch: The batched data as it is returned by the training DataLoader. """ # do something when the batch starts def on_batch_end(self) -> None: - """Called in the training loop after the batch.""" + """ + Called in the training loop after the batch. + """ # do something when the batch ends def on_epoch_start(self) -> None: - """Called in the training loop at the very beginning of the epoch.""" + """ + Called in the training loop at the very beginning of the epoch. + """ # do something when the epoch starts def on_epoch_end(self) -> None: - """Called in the training loop at the very end of the epoch.""" + """ + Called in the training loop at the very end of the epoch. + """ # do something when the epoch ends def on_pre_performance_check(self) -> None: - """Called at the very beginning of the validation loop.""" + """ + Called at the very beginning of the validation loop. + """ # do something before validation starts def on_post_performance_check(self) -> None: - """Called at the very end of the validation loop.""" + """ + Called at the very end of the validation loop. + """ # do something before validation end def on_before_zero_grad(self, optimizer: Optimizer) -> None: - """Called after optimizer.step() and before optimizer.zero_grad() + """ + Called after optimizer.step() and before optimizer.zero_grad(). Called in the training loop after taking an optimizer step and before zeroing grads. Good place to inspect weight information with weights updated. - for optimizer in optimizers:: + This is where it is called:: - optimizer.step() - model.on_before_zero_grad(optimizer) # < ---- called here - optimizer.zero_grad + for optimizer in optimizers: + optimizer.step() + model.on_before_zero_grad(optimizer) # < ---- called here + optimizer.zero_grad - :param optimizer: The optimizer for which grads should be zeroed. + Args: + optimizer: The optimizer for which grads should be zeroed. """ # do something with the optimizer or inspect it. def on_after_backward(self) -> None: - """Called in the training loop after loss.backward() and before optimizers do anything. + """ + Called in the training loop after loss.backward() and before optimizers do anything. + This is the ideal place to inspect or log gradient information. - This is the ideal place to inspect or log gradient information - - .. code-block:: python + Example:: def on_after_backward(self): # example to inspect gradient information in tensorboard @@ -113,19 +115,21 @@ class ModelHooks(torch.nn.Module): """ def backward(self, trainer, loss: Tensor, optimizer: Optimizer, optimizer_idx: int) -> None: - """Override backward with your own implementation if you need to + """ + Override backward with your own implementation if you need to. - :param trainer: Pointer to the trainer - :param loss: Loss is already scaled by accumulated grads - :param optimizer: Current optimizer being used - :param optimizer_idx: Index of the current optimizer being used + Args: + trainer: Pointer to the trainer + loss: Loss is already scaled by accumulated grads + optimizer: Current optimizer being used + optimizer_idx: Index of the current optimizer being used Called to perform backward step. Feel free to override as needed. The loss passed in has already been scaled for accumulated gradients if requested. - .. code-block:: python + Example:: def backward(self, use_amp, loss, optimizer): if use_amp: diff --git a/pytorch_lightning/core/memory.py b/pytorch_lightning/core/memory.py index da315733..8cbd281d 100644 --- a/pytorch_lightning/core/memory.py +++ b/pytorch_lightning/core/memory.py @@ -46,7 +46,7 @@ class ModelSummary(object): return list(mods) def get_variable_sizes(self) -> None: - """ Run sample input through each layer to get output sizes """ + """ Run sample input through each layer to get output sizes. """ mods = self.named_modules() in_sizes = [] out_sizes = [] @@ -116,7 +116,7 @@ class ModelSummary(object): self.layer_types = layer_types def get_parameter_sizes(self) -> None: - """ Get sizes of all parameters in `model` """ + """ Get sizes of all parameters in `model`. """ mods = self.named_modules() sizes = [] for _, m in mods: @@ -127,7 +127,7 @@ class ModelSummary(object): self.param_sizes = sizes def get_parameter_nums(self) -> None: - """ Get number of parameters in each layer """ + """ Get number of parameters in each layer. """ param_nums = [] for mod in self.param_sizes: all_params = 0 @@ -235,10 +235,19 @@ def count_mem_items() -> Tuple[int, int]: # pragma: no-cover def get_memory_profile(mode: str) -> Union[Dict[str, int], Dict[int, int]]: """ Get a profile of the current memory usage. - :param mode: There are two modes: - - 'all' means return memory for all gpus - - 'min_max' means return memory for max and min - :return: + Args: + mode: There are two modes: + + - 'all' means return memory for all gpus + - 'min_max' means return memory for max and min + + Return: + A dictionary in which the keys are device ids as integers and + values are memory usage as integers in MB. + If mode is 'min_max', the dictionary will also contain two additional keys: + + - 'min_gpu_mem': the minimum memory usage in MB + - 'max_gpu_mem': the maximum memory usage in MB """ memory_map = get_gpu_memory_map() @@ -280,15 +289,25 @@ def get_human_readable_count(number: int) -> str: billions and trillions, respectively. Examples: - 123 -> 123 - 1234 -> 1 K (one thousand) - 2e6 -> 2 M (two million) - 3e9 -> 3 B (three billion) - 4e12 -> 4 T (four trillion) - 5e15 -> 5,000 T + >>> get_human_readable_count(123) + '123 ' + >>> get_human_readable_count(1234) # (one thousand) + '1 K' + >>> get_human_readable_count(2e6) # (two million) + '2 M' + >>> get_human_readable_count(3e9) # (three billion) + '3 B' + >>> get_human_readable_count(4e12) # (four trillion) + '4 T' + >>> get_human_readable_count(5e15) # (more than trillion) + '5,000 T' + + Args: + number: a positive integer number + + Return: + A string formatted according to the pattern described above. - :param number: a positive integer number - :return: a string formatted according to the pattern described above. """ assert number >= 0 labels = [' ', 'K', 'M', 'B', 'T'] diff --git a/pytorch_lightning/core/saving.py b/pytorch_lightning/core/saving.py index 62fb92d3..d5e8d2a6 100644 --- a/pytorch_lightning/core/saving.py +++ b/pytorch_lightning/core/saving.py @@ -10,16 +10,21 @@ class ModelIO(object): def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None: """ - Do something with the checkpoint - Gives model a chance to load something before state_dict is restored - :param checkpoint: - :return: + Do something with the checkpoint. + Gives model a chance to load something before ``state_dict`` is restored. + + Args: + checkpoint: A dictionary with variables from the checkpoint. """ def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> None: """ Give the model a chance to add something to the checkpoint. - state_dict is already there + ``state_dict`` is already there. + + Args: + checkpoint: A dictionary in which you can save variables to save in a checkpoint. + Contents need to be pickleable. """ # ------------------------- @@ -27,12 +32,19 @@ class ModelIO(object): # ------------------------- def on_hpc_save(self, checkpoint: Dict[str, Any]) -> None: """ - Hook to do whatever you need right before Slurm manager saves the model + Hook to do whatever you need right before Slurm manager saves the model. + + Args: + checkpoint: A dictionary in which you can save variables to save in a checkpoint. + Contents need to be pickleable. """ def on_hpc_load(self, checkpoint: Dict[str, Any]) -> None: """ - Hook to do whatever you need right before Slurm manager loads the model + Hook to do whatever you need right before Slurm manager loads the model. + + Args: + checkpoint: A dictionary with variables from the checkpoint. """