diff --git a/examples/cifar10/main.py b/examples/cifar10/main.py index b8dbce5d9601..aadc310382e3 100644 --- a/examples/cifar10/main.py +++ b/examples/cifar10/main.py @@ -8,7 +8,7 @@ import torch.optim as optim import utils from torch.amp import autocast -from torch.cuda.amp import GradScaler +from torch.amp import GradScaler import ignite import ignite.distributed as idist diff --git a/examples/cifar100_amp_benchmark/benchmark_torch_cuda_amp.py b/examples/cifar100_amp_benchmark/benchmark_torch_cuda_amp.py index 746d7eb54c49..4ca1551ad823 100644 --- a/examples/cifar100_amp_benchmark/benchmark_torch_cuda_amp.py +++ b/examples/cifar100_amp_benchmark/benchmark_torch_cuda_amp.py @@ -1,7 +1,7 @@ import fire import torch from torch.amp import autocast -from torch.cuda.amp import GradScaler +from torch.amp import GradScaler from torch.nn import CrossEntropyLoss from torch.optim import SGD from torchvision.models import wide_resnet50_2 diff --git a/examples/cifar10_qat/main.py b/examples/cifar10_qat/main.py index 7b8366a2a63f..3fb7d59d13ba 100644 --- a/examples/cifar10_qat/main.py +++ b/examples/cifar10_qat/main.py @@ -7,7 +7,7 @@ import torch.optim as optim import utils from torch.amp import autocast -from torch.cuda.amp import GradScaler +from torch.amp import GradScaler import ignite import ignite.distributed as idist diff --git a/examples/transformers/main.py b/examples/transformers/main.py index f8118eabf90e..8eeca9768ac5 100644 --- a/examples/transformers/main.py +++ b/examples/transformers/main.py @@ -8,7 +8,7 @@ import torch.optim as optim import utils from torch.amp import autocast -from torch.cuda.amp import GradScaler +from torch.amp import GradScaler import ignite import ignite.distributed as idist diff --git a/ignite/engine/__init__.py b/ignite/engine/__init__.py index 6e82bc2f6bc7..1298932b843d 100644 --- a/ignite/engine/__init__.py +++ b/ignite/engine/__init__.py @@ -2,7 +2,7 @@ from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union import torch - +from torch.amp import GradScaler import ignite.distributed as idist from ignite.engine.deterministic import DeterministicEngine from ignite.engine.engine import Engine @@ -133,7 +133,7 @@ def supervised_training_step_amp( prepare_batch: Callable = _prepare_batch, model_transform: Callable[[Any], Any] = lambda output: output, output_transform: Callable[[Any, Any, Any, torch.Tensor], Any] = lambda x, y, y_pred, loss: loss.item(), - scaler: Optional["torch.cuda.amp.GradScaler"] = None, + scaler: Optional["torch.amp.GradScaler"] = None, gradient_accumulation_steps: int = 1, model_fn: Callable[[torch.nn.Module, Any], Any] = lambda model, x: model(x), ) -> Callable: @@ -393,8 +393,8 @@ def update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[to def _check_arg( - on_tpu: bool, on_mps: bool, amp_mode: Optional[str], scaler: Optional[Union[bool, "torch.cuda.amp.GradScaler"]] -) -> Tuple[Optional[str], Optional["torch.cuda.amp.GradScaler"]]: + on_tpu: bool, on_mps: bool, amp_mode: Optional[str], scaler: Optional[Union[bool, "torch.amp.GradScaler"]] +) -> Tuple[Optional[str], Optional["torch.amp.GradScaler"]]: """Checking tpu, mps, amp and GradScaler instance combinations.""" if on_mps and amp_mode: raise ValueError("amp_mode cannot be used with mps device. Consider using amp_mode=None or device='cuda'.") @@ -410,7 +410,7 @@ def _check_arg( raise ValueError(f"scaler argument is {scaler}, but amp_mode is {amp_mode}. Consider using amp_mode='amp'.") elif amp_mode == "amp" and isinstance(scaler, bool): try: - from torch.cuda.amp import GradScaler + from torch.amp import GradScaler except ImportError: raise ImportError("Please install torch>=1.6.0 to use scaler argument.") scaler = GradScaler(enabled=True) @@ -434,7 +434,7 @@ def create_supervised_trainer( output_transform: Callable[[Any, Any, Any, torch.Tensor], Any] = lambda x, y, y_pred, loss: loss.item(), deterministic: bool = False, amp_mode: Optional[str] = None, - scaler: Union[bool, "torch.cuda.amp.GradScaler"] = False, + scaler: Union[bool, "torch.amp.GradScaler"] = False, gradient_accumulation_steps: int = 1, model_fn: Callable[[torch.nn.Module, Any], Any] = lambda model, x: model(x), ) -> Engine: diff --git a/tests/ignite/conftest.py b/tests/ignite/conftest.py index a9a35de69813..dd1d8e45e5c4 100644 --- a/tests/ignite/conftest.py +++ b/tests/ignite/conftest.py @@ -48,7 +48,7 @@ def pytest_configure(config): config.addinivalue_line("markers", "distributed: run distributed") config.addinivalue_line("markers", "multinode_distributed: distributed") config.addinivalue_line("markers", "tpu: run on tpu") - if config.option.treat_unrun_as_failed: + if getattr(config.option, "treat_unrun_as_failed", False): unrun_tracker = UnrunTracker() config.pluginmanager.register(unrun_tracker, "unrun_tracker_plugin") @@ -611,6 +611,6 @@ def pytest_sessionfinish(session, exitstatus): run finished, right before returning the exit status to the system. """ # If requested by the user, track all unrun tests and add them to the lastfailed cache - if session.config.option.treat_unrun_as_failed: + if getattr(session.config.option, "treat_unrun_as_failed", False): unrun_tracker = session.config.pluginmanager.get_plugin("unrun_tracker_plugin") unrun_tracker.record_unrun_as_failed(session, exitstatus) diff --git a/tests/ignite/engine/test_create_supervised.py b/tests/ignite/engine/test_create_supervised.py index ba42baddddae..6bd759f9c2b5 100644 --- a/tests/ignite/engine/test_create_supervised.py +++ b/tests/ignite/engine/test_create_supervised.py @@ -48,7 +48,7 @@ def _default_create_supervised_trainer( trainer_device: Optional[str] = None, trace: bool = False, amp_mode: str = None, - scaler: Union[bool, "torch.cuda.amp.GradScaler"] = False, + scaler: Union[bool, "torch.amp.GradScaler"] = False, with_model_transform: bool = False, with_model_fn: bool = False, ): @@ -104,7 +104,7 @@ def _test_create_supervised_trainer( trainer_device: Optional[str] = None, trace: bool = False, amp_mode: str = None, - scaler: Union[bool, "torch.cuda.amp.GradScaler"] = False, + scaler: Union[bool, "torch.amp.GradScaler"] = False, with_model_transform: bool = False, with_model_fn: bool = False, ): @@ -170,10 +170,10 @@ def _(): @pytest.mark.skipif(Version(torch.__version__) < Version("1.12.0"), reason="Skip if < 1.12.0") def test_create_supervised_training_scalar_assignment(): with mock.patch("ignite.engine._check_arg") as check_arg_mock: - check_arg_mock.return_value = None, torch.cuda.amp.GradScaler(enabled=False) + check_arg_mock.return_value = None, torch.amp.GradScaler(enabled=False) trainer, _ = _default_create_supervised_trainer(model_device="cpu", trainer_device="cpu", scaler=True) assert hasattr(trainer.state, "scaler") - assert isinstance(trainer.state.scaler, torch.cuda.amp.GradScaler) + assert isinstance(trainer.state.scaler, torch.amp.GradScaler) def _test_create_mocked_supervised_trainer( @@ -181,7 +181,7 @@ def _test_create_mocked_supervised_trainer( trainer_device: Optional[str] = None, trace: bool = False, amp_mode: str = None, - scaler: Union[bool, "torch.cuda.amp.GradScaler"] = False, + scaler: Union[bool, "torch.amp.GradScaler"] = False, ): with mock.patch("ignite.engine.supervised_training_step_amp") as training_step_amp_mock: with mock.patch("ignite.engine.supervised_training_step_apex") as training_step_apex_mock: @@ -462,7 +462,7 @@ def test_create_supervised_trainer_amp_error(mock_torch_cuda_amp_module): @pytest.mark.skipif(Version(torch.__version__) < Version("1.12.0"), reason="Skip if < 1.12.0") def test_create_supervised_trainer_scaler_not_amp(): - scaler = torch.cuda.amp.GradScaler(enabled=torch.cuda.is_available()) + scaler = torch.amp.GradScaler(enabled=torch.cuda.is_available()) with pytest.raises(ValueError, match=f"scaler argument is {scaler}, but amp_mode is None."): _test_create_supervised_trainer(amp_mode=None, scaler=scaler) @@ -540,7 +540,7 @@ def test_create_supervised_trainer_on_cuda_amp_scaler(): _test_create_mocked_supervised_trainer( model_device=model_device, trainer_device=trainer_device, amp_mode="amp", scaler=True ) - scaler = torch.cuda.amp.GradScaler(enabled=torch.cuda.is_available()) + scaler = torch.amp.GradScaler(enabled=torch.cuda.is_available()) _test_create_supervised_trainer( gradient_accumulation_steps=1, model_device=model_device,