Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

clean AMP logic (#5994) 

* clean AMP logic

* cleaning

* ...

* ...

* Even apex
This commit is contained in:
Jirka Borovec 2021-02-16 20:06:47 +01:00 committed by GitHub
parent b40d414463
commit 936f42aa1c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 25 additions and 40 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

50
pytorch_lightning/accelerators/accelerator_connector.py
View File

@ -298,46 +298,42 @@ class BackendConnector(object):
if self.on_tpu:
return TPUHalfPrecisionPlugin()
if self.amp_type == "native":
if not _NATIVE_AMP_AVAILABLE:
rank_zero_warn(
"You have asked for native AMP but your PyTorch version does not support it."
" Consider upgrading with `pip install torch>=1.6`."
" We will attempt to use NVIDIA Apex for this session."
)
if not _APEX_AVAILABLE and self.on_cpu:
raise MisconfigurationException(
"You have asked for native AMP on CPU, but AMP is only available on GPU."
)
self.amp_type = "apex"
elif self.on_cpu:
self.amp_type = AMPType(self.amp_type)
if self.amp_type == AMPType.NATIVE:
if self.on_cpu:
raise MisconfigurationException(
"You have asked for native AMP on CPU, but AMP is only available on GPU."
)
elif not _NATIVE_AMP_AVAILABLE:
msg = "You have asked for native AMP but your PyTorch version does not support it." \
" Consider upgrading with `pip install torch>=1.6`."
if _APEX_AVAILABLE:
self.amp_type = AMPType.APEX
msg += " We will attempt to use NVIDIA Apex for this session."
rank_zero_warn(msg)
else:
raise MisconfigurationException(msg)
else:
log.info("Using native 16bit precision.")
if isinstance(self.training_type_plugin, (DDPShardedPlugin, DDPSpawnShardedPlugin)):
return ShardedNativeMixedPrecisionPlugin()
self.amp_type = AMPType.NATIVE
return NativeMixedPrecisionPlugin()
if self.amp_type == "apex":
if self.amp_type == AMPType.APEX:
if not _APEX_AVAILABLE:
rank_zero_warn(
raise MisconfigurationException(
"You have asked for Apex AMP but you have not installed it yet."
" Install apex first using this guide: https://github.com/NVIDIA/apex#linux"
)
else:
if isinstance(self.training_type_plugin, (DDPShardedPlugin, DDPSpawnShardedPlugin)):
raise MisconfigurationException(
"Sharded Plugin is not supported with Apex AMP, "
"please using native AMP for 16-bit precision."
)
log.info("Using APEX 16bit precision.")
self.amp_type = AMPType.APEX
return ApexMixedPrecisionPlugin(self.amp_level)
else:
raise NotImplementedError("We only support precisions 32 and 16!")
if isinstance(self.training_type_plugin, (DDPShardedPlugin, DDPSpawnShardedPlugin)):
raise MisconfigurationException(
"Sharded Plugin is not supported with Apex AMP,"
" please using native AMP for 16-bit precision."
)
log.info("Using APEX 16bit precision.")
return ApexMixedPrecisionPlugin(self.amp_level)
raise NotImplementedError("We only support precisions 32 and 16!")
def select_training_type_plugin(self) -> TrainingTypePlugin:
if self.use_ddp2:

15
tests/models/test_amp.py
View File

@ -18,7 +18,6 @@ import pytest
import torch
from torch import optim
import tests.helpers.pipelines as tpipes
import tests.helpers.utils as tutils
from pytorch_lightning import Trainer
from pytorch_lightning.plugins.environments import SLURMEnvironment
@ -155,21 +154,11 @@ def test_amp_gpu_ddp_slurm_managed(tmpdir):
assert generated == 'abc23'
@pytest.mark.skipif(torch.cuda.is_available(), reason="test is restricted only on CPU")
def test_cpu_model_with_amp(tmpdir):
"""Make sure model trains on CPU."""
trainer_options = dict(
default_root_dir=tmpdir,
progress_bar_refresh_rate=0,
max_epochs=1,
limit_train_batches=0.4,
limit_val_batches=0.4,
precision=16,
)
model = BoringModel()
with pytest.raises(MisconfigurationException, match="AMP is only available on GPU"):
tpipes.run_model_test(trainer_options, model, on_gpu=False)
Trainer(precision=16)
@mock.patch.dict(os.environ, {"PL_DEV_DEBUG": "1"})