wassname/pytorch-lightning
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-lightning.git synced 2026-09-06 16:50:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Improved docs for pytorch_lightning.core (continued) (#1483)

Browse Source 
* 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
This commit is contained in:
Adrian Wälchli
2020-04-16 12:04:55 -04:00
committed by GitHub
parent 6e1d72d98a
commit 2ab2f7d08d
5 changed files with 320 additions and 270 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/source/hooks.rst
+57 -39
View File
@@ -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 <https://github.com/PyTorchLightning/pytorch-lightning>`_.
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
- ``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:
pytorch_lightning/core/__init__.py
+162 -165
View File
@@ -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
---------
pytorch_lightning/core/hooks.py
+48 -44
View File
@@ -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:
pytorch_lightning/core/memory.py
+34 -15
View File
@@ -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']
pytorch_lightning/core/saving.py
+19 -7
View File
@@ -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.
"""