remove device_type

Author: four4fish

pytorch_lightning/callbacks/gpu_stats_monitor.py

@@ -28,8 +28,9 @@
 import torch
 
 import pytorch_lightning as pl
+from pytorch_lightning.accelerators import GPUAccelerator
 from pytorch_lightning.callbacks.base import Callback
-from pytorch_lightning.utilities import _AcceleratorType, rank_zero_deprecation, rank_zero_only
+from pytorch_lightning.utilities import rank_zero_deprecation, rank_zero_only
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.parsing import AttributeDict
 from pytorch_lightning.utilities.types import STEP_OUTPUT
@@ -126,7 +127,7 @@ def setup(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", stage: O
         if not trainer.logger:
             raise MisconfigurationException("Cannot use GPUStatsMonitor callback with Trainer that has no logger.")
 
-        if trainer._device_type != _AcceleratorType.GPU:
+        if not isinstance(trainer.accelerator, GPUAccelerator):
             raise MisconfigurationException(
                 "You are using GPUStatsMonitor but are not running on GPU"
                 f" since gpus attribute in Trainer is set to {trainer.gpus}."

pytorch_lightning/callbacks/xla_stats_monitor.py

@@ -20,8 +20,9 @@
 """
 import time
 
+from pytorch_lightning.accelerators import TPUAccelerator
 from pytorch_lightning.callbacks.base import Callback
-from pytorch_lightning.utilities import _AcceleratorType, _TPU_AVAILABLE, rank_zero_deprecation, rank_zero_info
+from pytorch_lightning.utilities import _TPU_AVAILABLE, rank_zero_deprecation, rank_zero_info
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 
 if _TPU_AVAILABLE:
@@ -70,7 +71,7 @@ def on_train_start(self, trainer, pl_module) -> None:
         if not trainer.logger:
             raise MisconfigurationException("Cannot use XLAStatsMonitor callback with Trainer that has no logger.")
 
-        if trainer._device_type != _AcceleratorType.TPU:
+        if not isinstance(trainer.accelerator, TPUAccelerator):
             raise MisconfigurationException(
                 "You are using XLAStatsMonitor but are not running on TPU"
                 f" since `tpu_cores` attribute in Trainer is set to {trainer.tpu_cores}."

pytorch_lightning/loops/optimization/optimizer_loop.py

@@ -19,6 +19,7 @@
 from torch import Tensor
 from torch.optim import Optimizer
 
+from pytorch_lightning.accelerators import TPUAccelerator
 from pytorch_lightning.core.optimizer import LightningOptimizer
 from pytorch_lightning.loops import Loop
 from pytorch_lightning.loops.optimization.closure import AbstractClosure, OutputResult
@@ -30,7 +31,7 @@
 )
 from pytorch_lightning.profiler import BaseProfiler, PassThroughProfiler
 from pytorch_lightning.trainer.progress import OptimizationProgress
-from pytorch_lightning.utilities import _AcceleratorType, AMPType
+from pytorch_lightning.utilities import AMPType
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.finite_checks import detect_nan_parameters
 from pytorch_lightning.utilities.imports import _TPU_AVAILABLE
@@ -374,7 +375,7 @@ def _optimizer_step(
             optimizer,
             opt_idx,
             train_step_and_backward_closure,
-            on_tpu=(self.trainer._device_type == _AcceleratorType.TPU and _TPU_AVAILABLE),
+            on_tpu=(isinstance(self.trainer.accelerator == TPUAccelerator) and _TPU_AVAILABLE),
             using_native_amp=(self.trainer.amp_backend is not None and self.trainer.amp_backend == AMPType.NATIVE),
             using_lbfgs=is_lbfgs,
         )

pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py

@@ -17,11 +17,12 @@
 import torch
 
 import pytorch_lightning as pl
+from pytorch_lightning.accelerators import GPUAccelerator
 from pytorch_lightning.loggers import LightningLoggerBase, LoggerCollection, TensorBoardLogger
 from pytorch_lightning.plugins.environments.slurm_environment import SLURMEnvironment
 from pytorch_lightning.trainer.connectors.logger_connector.result import _METRICS, _OUT_DICT, _PBAR_DICT
 from pytorch_lightning.trainer.states import RunningStage, TrainerFn
-from pytorch_lightning.utilities import _AcceleratorType, memory
+from pytorch_lightning.utilities import memory
 from pytorch_lightning.utilities.apply_func import apply_to_collection, move_data_to_device
 from pytorch_lightning.utilities.metrics import metrics_to_scalars
 from pytorch_lightning.utilities.warnings import rank_zero_deprecation
@@ -330,7 +331,7 @@ def gpus_metrics(self) -> Dict[str, float]:
         .. deprecated:: v1.5
             Will be removed in v1.7.
         """
-        if self.trainer._device_type == _AcceleratorType.GPU and self.log_gpu_memory:
+        if isinstance(self.trainer.accelerator, GPUAccelerator) and self.log_gpu_memory:
             mem_map = memory.get_memory_profile(self.log_gpu_memory)
             self._gpus_metrics.update(mem_map)
         return self._gpus_metrics

pytorch_lightning/trainer/trainer.py

@@ -73,7 +73,6 @@
 from pytorch_lightning.tuner.lr_finder import _LRFinder
 from pytorch_lightning.tuner.tuning import Tuner
 from pytorch_lightning.utilities import (
-    _AcceleratorType,
     _IPU_AVAILABLE,
     _StrategyType,
     _TPU_AVAILABLE,
@@ -1706,9 +1705,9 @@ def should_rank_save_checkpoint(self) -> bool:
     def _distrib_type(self) -> _StrategyType:
         return self._accelerator_connector._distrib_type
 
-    @property
-    def _device_type(self) -> _AcceleratorType:
-        return self._accelerator_connector._device_type
+    # @property
+    # def _device_type(self) -> _AcceleratorType:
+    #     return self._accelerator_connector._device_type
 
     @property
     def num_nodes(self) -> int:

pytorch_lightning/utilities/model_summary.py

@@ -23,7 +23,8 @@
 from torch.utils.hooks import RemovableHandle
 
 import pytorch_lightning as pl
-from pytorch_lightning.utilities import _AcceleratorType, AMPType
+from pytorch_lightning.accelerators import TPUAccelerator
+from pytorch_lightning.utilities import AMPType
 from pytorch_lightning.utilities.imports import _TORCH_GREATER_EQUAL_1_8
 from pytorch_lightning.utilities.warnings import WarningCache
 
@@ -264,7 +265,7 @@ def _forward_example_input(self) -> None:
         if (
             trainer is not None
             and trainer.amp_backend == AMPType.NATIVE
-            and trainer._device_type != _AcceleratorType.TPU
+            and isinstance(trainer.accelerator, TPUAccelerator)
         ):
             model.forward = torch.cuda.amp.autocast()(model.forward)
 

tests/accelerators/test_ipu.py

@@ -24,7 +24,7 @@
 from pytorch_lightning.plugins import IPUPlugin, IPUPrecisionPlugin
 from pytorch_lightning.trainer.states import RunningStage, TrainerFn
 from pytorch_lightning.trainer.supporters import CombinedLoader
-from pytorch_lightning.utilities import _AcceleratorType, _IPU_AVAILABLE
+from pytorch_lightning.utilities import _IPU_AVAILABLE
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.helpers.boring_model import BoringModel
 from tests.helpers.datamodules import ClassifDataModule
@@ -500,7 +500,6 @@ def test_accelerator_ipu():
 
     trainer = Trainer(accelerator="ipu", ipus=1)
 
-    assert trainer._device_type == "ipu"
     assert isinstance(trainer.accelerator, IPUAccelerator)
 
     with pytest.raises(
@@ -510,7 +509,6 @@ def test_accelerator_ipu():
 
     trainer = Trainer(accelerator="auto", ipus=8)
 
-    assert trainer._device_type == "ipu"
     assert isinstance(trainer.accelerator, IPUAccelerator)
 
 
@@ -519,7 +517,6 @@ def test_accelerator_cpu_with_ipus_flag():
 
     trainer = Trainer(accelerator="cpu", ipus=1)
 
-    assert trainer._device_type == "cpu"
     assert isinstance(trainer.accelerator, CPUAccelerator)
 
 
@@ -538,7 +535,7 @@ def test_accelerator_auto_with_devices_ipu():
 
     trainer = Trainer(accelerator="auto", devices=8)
 
-    assert trainer._device_type == "ipu"
+    assert isinstance(trainer.accelerator, IPUAccelerator)
     assert trainer.ipus == 8
 
 
@@ -567,11 +564,10 @@ def test_strategy_choice_ipu_plugin(tmpdir):
 
 
 @RunIf(ipu=True)
-def test_device_type_when_training_plugin_ipu_passed(tmpdir):
+def test_accelerator_type_when_training_plugin_ipu_passed(tmpdir):
 
     trainer = Trainer(strategy=IPUPlugin(), ipus=8)
     assert isinstance(trainer.training_type_plugin, IPUPlugin)
-    assert trainer._device_type == _AcceleratorType.IPU
     assert isinstance(trainer.accelerator, IPUAccelerator)
 
 

tests/accelerators/test_tpu.py

@@ -85,7 +85,6 @@ def test_accelerator_tpu():
 
     trainer = Trainer(accelerator="tpu", tpu_cores=8)
 
-    assert trainer._device_type == "tpu"
     assert isinstance(trainer.accelerator, TPUAccelerator)
 
     with pytest.raises(
@@ -99,7 +98,6 @@ def test_accelerator_cpu_with_tpu_cores_flag():
 
     trainer = Trainer(accelerator="cpu", tpu_cores=8)
 
-    assert trainer._device_type == "cpu"
     assert isinstance(trainer.accelerator, CPUAccelerator)
 
 
@@ -108,7 +106,6 @@ def test_accelerator_tpu_with_auto():
 
     trainer = Trainer(accelerator="auto", tpu_cores=8)
 
-    assert trainer._device_type == "tpu"
     assert isinstance(trainer.accelerator, TPUAccelerator)
 
 
@@ -127,7 +124,7 @@ def test_accelerator_auto_with_devices_tpu():
 
     trainer = Trainer(accelerator="auto", devices=8)
 
-    assert trainer._device_type == "tpu"
+    assert isinstance(trainer.accelerator, TPUAccelerator)
     assert trainer.tpu_cores == 8
 
 

tests/models/test_tpu.py

@@ -26,7 +26,7 @@
 from pytorch_lightning.callbacks import EarlyStopping
 from pytorch_lightning.plugins import TPUSpawnPlugin
 from pytorch_lightning.trainer.connectors.logger_connector.result import _Sync
-from pytorch_lightning.utilities import _AcceleratorType, _TPU_AVAILABLE
+from pytorch_lightning.utilities import _TPU_AVAILABLE
 from pytorch_lightning.utilities.distributed import ReduceOp
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.helpers import BoringModel, RandomDataset
@@ -474,5 +474,4 @@ def test_device_type_when_training_plugin_tpu_passed(tmpdir):
 
     trainer = Trainer(strategy=TPUSpawnPlugin(), tpu_cores=8)
     assert isinstance(trainer.training_type_plugin, TPUSpawnPlugin)
-    assert trainer._device_type == _AcceleratorType.TPU
     assert isinstance(trainer.accelerator, TPUAccelerator)