141 lines
4.4 KiB
Python
141 lines
4.4 KiB
Python
import os
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from pytorch_lightning import Trainer
|
|
from pytorch_lightning.callbacks import Callback
|
|
from pytorch_lightning.plugins.native_amp import NativeAMPPlugin
|
|
from pytorch_lightning.utilities import NATIVE_AMP_AVAILABLE
|
|
from tests.base.boring_model import BoringModel
|
|
|
|
|
|
@pytest.mark.skipif(not NATIVE_AMP_AVAILABLE, reason="Minimal PT version is set to 1.6")
|
|
@mock.patch.dict(os.environ, {
|
|
"CUDA_VISIBLE_DEVICES": "0,1",
|
|
"SLURM_NTASKS": "2",
|
|
"SLURM_JOB_NAME": "SOME_NAME",
|
|
"SLURM_NODEID": "0",
|
|
"LOCAL_RANK": "0",
|
|
"SLURM_LOCALID": "0"
|
|
})
|
|
@mock.patch('torch.cuda.device_count', return_value=2)
|
|
@pytest.mark.parametrize(['ddp_backend', 'gpus', 'num_processes'],
|
|
[('ddp_cpu', None, None), ('ddp', 2, 0), ('ddp2', 2, 0), ('ddp_spawn', 2, 0)])
|
|
def test_amp_choice_default_ddp_cpu(tmpdir, ddp_backend, gpus, num_processes):
|
|
|
|
class CB(Callback):
|
|
def on_fit_start(self, trainer, pl_module):
|
|
assert isinstance(trainer.precision_connector.backend, NativeAMPPlugin)
|
|
raise SystemExit()
|
|
|
|
model = BoringModel()
|
|
trainer = Trainer(
|
|
fast_dev_run=True,
|
|
precision=16,
|
|
amp_backend='native',
|
|
gpus=gpus,
|
|
num_processes=num_processes,
|
|
distributed_backend=ddp_backend,
|
|
callbacks=[CB()]
|
|
)
|
|
|
|
with pytest.raises(SystemExit):
|
|
trainer.fit(model)
|
|
|
|
|
|
@pytest.mark.skipif(not NATIVE_AMP_AVAILABLE, reason="Minimal PT version is set to 1.6")
|
|
@mock.patch.dict(os.environ, {
|
|
"CUDA_VISIBLE_DEVICES": "0,1",
|
|
"SLURM_NTASKS": "2",
|
|
"SLURM_JOB_NAME": "SOME_NAME",
|
|
"SLURM_NODEID": "0",
|
|
"LOCAL_RANK": "0",
|
|
"SLURM_LOCALID": "0"
|
|
})
|
|
@mock.patch('torch.cuda.device_count', return_value=2)
|
|
@pytest.mark.parametrize(['ddp_backend', 'gpus', 'num_processes'],
|
|
[('ddp_cpu', None, None), ('ddp', 2, 0), ('ddp2', 2, 0), ('ddp_spawn', 2, 0)])
|
|
def test_amp_choice_custom_ddp_cpu(tmpdir, ddp_backend, gpus, num_processes):
|
|
class MyNativeAMP(NativeAMPPlugin):
|
|
pass
|
|
|
|
class CB(Callback):
|
|
def on_fit_start(self, trainer, pl_module):
|
|
assert isinstance(trainer.precision_connector.backend, MyNativeAMP)
|
|
raise SystemExit()
|
|
|
|
model = BoringModel()
|
|
trainer = Trainer(
|
|
fast_dev_run=True,
|
|
precision=16,
|
|
amp_backend='native',
|
|
gpus=gpus,
|
|
num_processes=num_processes,
|
|
distributed_backend=ddp_backend,
|
|
plugins=[MyNativeAMP()],
|
|
callbacks=[CB()]
|
|
)
|
|
|
|
with pytest.raises(SystemExit):
|
|
trainer.fit(model)
|
|
|
|
|
|
class GradientUnscaleBoringModel(BoringModel):
|
|
def on_after_backward(self):
|
|
norm = torch.nn.utils.clip_grad_norm_(self.parameters(), 2)
|
|
if not (torch.isinf(norm) or torch.isnan(norm)):
|
|
assert norm.item() < 15.
|
|
|
|
|
|
@pytest.mark.skipif(not NATIVE_AMP_AVAILABLE, reason="Minimal PT version is set to 1.6")
|
|
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
|
|
def test_amp_gradient_unscale(tmpdir):
|
|
model = GradientUnscaleBoringModel()
|
|
|
|
trainer = Trainer(
|
|
max_epochs=2,
|
|
default_root_dir=os.getcwd(),
|
|
limit_train_batches=2,
|
|
limit_test_batches=2,
|
|
limit_val_batches=2,
|
|
amp_backend='native',
|
|
distributed_backend='ddp_spawn',
|
|
gpus=2,
|
|
precision=16,
|
|
track_grad_norm=2,
|
|
log_every_n_steps=1
|
|
)
|
|
trainer.fit(model)
|
|
|
|
|
|
class UnscaleAccumulateGradBatchesBoringModel(BoringModel):
|
|
|
|
def on_after_backward(self):
|
|
norm = torch.nn.utils.clip_grad_norm_(self.parameters(), 2)
|
|
if not (torch.isinf(norm) or torch.isnan(norm)):
|
|
assert norm.item() < 15.
|
|
|
|
|
|
@pytest.mark.skipif(not NATIVE_AMP_AVAILABLE, reason="Minimal PT version is set to 1.6")
|
|
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
|
|
def test_amp_gradient_unscale_accumulate_grad_batches(tmpdir):
|
|
model = UnscaleAccumulateGradBatchesBoringModel()
|
|
|
|
trainer = Trainer(
|
|
max_epochs=2,
|
|
default_root_dir=os.getcwd(),
|
|
limit_train_batches=2,
|
|
limit_test_batches=2,
|
|
limit_val_batches=2,
|
|
amp_backend='native',
|
|
distributed_backend='ddp_spawn',
|
|
gpus=2,
|
|
precision=16,
|
|
track_grad_norm=2,
|
|
log_every_n_steps=1,
|
|
accumulate_grad_batches=2,
|
|
)
|
|
trainer.fit(model)
|