unify torch.Tensor >> Tensor

Author: Borda

src/pytorch_lightning/callbacks/callback.py

@@ -18,7 +18,7 @@
 
 from typing import Any, Dict, List, Optional, Type
 
-import torch
+from torch import Tensor
 from torch.optim import Optimizer
 
 import pytorch_lightning as pl
@@ -342,7 +342,7 @@ def on_load_checkpoint(
             checkpoint dictionary instead of only the callback state from the checkpoint.
         """
 
-    def on_before_backward(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", loss: torch.Tensor) -> None:
+    def on_before_backward(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", loss: Tensor) -> None:
         """Called before ``loss.backward()``."""
 
     def on_after_backward(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:

src/pytorch_lightning/callbacks/early_stopping.py

@@ -23,6 +23,7 @@
 
 import numpy as np
 import torch
+from torch import Tensor
 
 import pytorch_lightning as pl
 from pytorch_lightning.callbacks.callback import Callback
@@ -203,7 +204,7 @@ def _run_early_stopping_check(self, trainer: "pl.Trainer") -> None:
         if reason and self.verbose:
             self._log_info(trainer, reason)
 
-    def _evaluate_stopping_criteria(self, current: torch.Tensor) -> Tuple[bool, Optional[str]]:
+    def _evaluate_stopping_criteria(self, current: Tensor) -> Tuple[bool, Optional[str]]:
         should_stop = False
         reason = None
         if self.check_finite and not torch.isfinite(current):
@@ -242,7 +243,7 @@ def _evaluate_stopping_criteria(self, current: torch.Tensor) -> Tuple[bool, Opti
 
         return should_stop, reason
 
-    def _improvement_message(self, current: torch.Tensor) -> str:
+    def _improvement_message(self, current: Tensor) -> str:
         """Formats a log message that informs the user about an improvement in the monitored score."""
         if torch.isfinite(self.best_score):
             msg = (

src/pytorch_lightning/callbacks/model_checkpoint.py

@@ -31,6 +31,7 @@
 import numpy as np
 import torch
 import yaml
+from torch import Tensor
 
 import pytorch_lightning as pl
 from pytorch_lightning.callbacks.callback import Callback
@@ -477,7 +478,7 @@ def __init_triggers(
     def every_n_epochs(self) -> Optional[int]:
         return self._every_n_epochs
 
-    def check_monitor_top_k(self, trainer: "pl.Trainer", current: Optional[torch.Tensor] = None) -> bool:
+    def check_monitor_top_k(self, trainer: "pl.Trainer", current: Optional[Tensor] = None) -> bool:
         if current is None:
             return False
 
@@ -628,11 +629,9 @@ def _monitor_candidates(self, trainer: "pl.Trainer") -> Dict[str, _METRIC]:
         # cast to int if necessary because `self.log("epoch", 123)` will convert it to float. if it's not a tensor
         # or does not exist we overwrite it as it's likely an error
         epoch = monitor_candidates.get("epoch")
-        monitor_candidates["epoch"] = (
-            epoch.int() if isinstance(epoch, torch.Tensor) else torch.tensor(trainer.current_epoch)
-        )
+        monitor_candidates["epoch"] = epoch.int() if isinstance(epoch, Tensor) else torch.tensor(trainer.current_epoch)
         step = monitor_candidates.get("step")
-        monitor_candidates["step"] = step.int() if isinstance(step, torch.Tensor) else torch.tensor(trainer.global_step)
+        monitor_candidates["step"] = step.int() if isinstance(step, Tensor) else torch.tensor(trainer.global_step)
         return monitor_candidates
 
     def _save_last_checkpoint(self, trainer: "pl.Trainer", monitor_candidates: Dict[str, _METRIC]) -> None:
@@ -670,7 +669,7 @@ def _save_none_monitor_checkpoint(self, trainer: "pl.Trainer", monitor_candidate
             trainer.strategy.remove_checkpoint(previous)
 
     def _update_best_and_save(
-        self, current: torch.Tensor, trainer: "pl.Trainer", monitor_candidates: Dict[str, _METRIC]
+        self, current: Tensor, trainer: "pl.Trainer", monitor_candidates: Dict[str, _METRIC]
     ) -> None:
         k = len(self.best_k_models) + 1 if self.save_top_k == -1 else self.save_top_k
 
@@ -680,7 +679,7 @@ def _update_best_and_save(
             self.best_k_models.pop(del_filepath)
 
         # do not save nan, replace with +/- inf
-        if isinstance(current, torch.Tensor) and torch.isnan(current):
+        if isinstance(current, Tensor) and torch.isnan(current):
             current = torch.tensor(float("inf" if self.mode == "min" else "-inf"), device=current.device)
 
         filepath = self._get_metric_interpolated_filepath_name(monitor_candidates, trainer, del_filepath)

src/pytorch_lightning/callbacks/pruning.py

@@ -21,9 +21,8 @@
 from functools import partial
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
 
-import torch
 import torch.nn.utils.prune as pytorch_prune
-from torch import nn
+from torch import nn, Tensor
 from typing_extensions import TypedDict
 
 import pytorch_lightning as pl
@@ -275,7 +274,7 @@ def make_pruning_permanent(self, module: nn.Module) -> None:
     def _copy_param(new: nn.Module, old: nn.Module, name: str) -> None:
         dst = getattr(new, name)
         src = getattr(old, name)
-        if dst is None or src is None or not isinstance(dst, torch.Tensor) or not isinstance(src, torch.Tensor):
+        if dst is None or src is None or not isinstance(dst, Tensor) or not isinstance(src, Tensor):
             return
         dst.data = src.data.to(dst.device)
 
@@ -418,11 +417,11 @@ def _make_pruning_permanent_on_state_dict(self, pl_module: LightningModule) -> D
             # make weights permanent
             state_dict[tensor_name] = mask.to(dtype=orig.dtype) * orig
 
-        def move_to_cpu(tensor: torch.Tensor) -> torch.Tensor:
+        def move_to_cpu(tensor: Tensor) -> Tensor:
             # each tensor and move them on cpu
             return tensor.cpu()
 
-        return apply_to_collection(state_dict, torch.Tensor, move_to_cpu)
+        return apply_to_collection(state_dict, Tensor, move_to_cpu)
 
     def on_save_checkpoint(
         self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", checkpoint: Dict[str, Any]

src/pytorch_lightning/callbacks/stochastic_weight_avg.py

@@ -19,7 +19,7 @@
 from typing import Callable, List, Optional, Union
 
 import torch
-from torch import nn
+from torch import FloatTensor, nn, Tensor
 from torch.optim.swa_utils import SWALR
 
 import pytorch_lightning as pl
@@ -28,7 +28,7 @@
 from pytorch_lightning.utilities.rank_zero import rank_zero_info, rank_zero_warn
 from pytorch_lightning.utilities.types import LRSchedulerConfig
 
-_AVG_FN = Callable[[torch.Tensor, torch.Tensor, torch.LongTensor], torch.FloatTensor]
+_AVG_FN = Callable[[Tensor, Tensor, torch.LongTensor], FloatTensor]
 
 
 class StochasticWeightAveraging(Callback):
@@ -269,7 +269,7 @@ def update_parameters(
 
     @staticmethod
     def avg_fn(
-        averaged_model_parameter: torch.Tensor, model_parameter: torch.Tensor, num_averaged: torch.LongTensor
-    ) -> torch.FloatTensor:
+        averaged_model_parameter: Tensor, model_parameter: Tensor, num_averaged: torch.LongTensor
+    ) -> FloatTensor:
         """Adapted from https://github.com/pytorch/pytorch/blob/v1.7.1/torch/optim/swa_utils.py#L95-L97."""
         return averaged_model_parameter + (model_parameter - averaged_model_parameter) / (num_averaged + 1)

src/pytorch_lightning/core/hooks.py

@@ -16,6 +16,7 @@
 from typing import Any, Dict, List, Optional
 
 import torch
+from torch import Tensor
 from torch.optim.optimizer import Optimizer
 
 from pytorch_lightning.utilities import move_data_to_device
@@ -254,7 +255,7 @@ def on_before_zero_grad(self, optimizer: Optimizer) -> None:
             optimizer: The optimizer for which grads should be zeroed.
         """
 
-    def on_before_backward(self, loss: torch.Tensor) -> None:
+    def on_before_backward(self, loss: Tensor) -> None:
         """Called before ``loss.backward()``.
 
         Args:

src/pytorch_lightning/core/mixins/device_dtype_mixin.py

@@ -89,8 +89,9 @@ def to(self, *args: Any, **kwargs: Any) -> Self:
             Module: self
 
         Example::
+            >>> from torch import Tensor
             >>> class ExampleModule(DeviceDtypeModuleMixin):
-            ...     def __init__(self, weight: torch.Tensor):
+            ...     def __init__(self, weight: Tensor):
             ...         super().__init__()
             ...         self.register_buffer('weight', weight)
             >>> _ = torch.manual_seed(0)

src/pytorch_lightning/core/module.py

@@ -528,7 +528,7 @@ def __check_not_nested(value: dict, name: str) -> dict:
     def __check_allowed(v: Any, name: str, value: Any) -> None:
         raise ValueError(f"`self.log({name}, {value})` was called, but `{type(v).__name__}` values cannot be logged")
 
-    def __to_tensor(self, value: numbers.Number) -> torch.Tensor:
+    def __to_tensor(self, value: numbers.Number) -> Tensor:
         return torch.tensor(value, device=self.device)
 
     def log_grad_norm(self, grad_norm_dict: Dict[str, float]) -> None:
@@ -547,9 +547,7 @@ def log_grad_norm(self, grad_norm_dict):
         """
         self.log_dict(grad_norm_dict, on_step=True, on_epoch=True, prog_bar=False, logger=True)
 
-    def all_gather(
-        self, data: Union[torch.Tensor, Dict, List, Tuple], group: Optional[Any] = None, sync_grads: bool = False
-    ):
+    def all_gather(self, data: Union[Tensor, Dict, List, Tuple], group: Optional[Any] = None, sync_grads: bool = False):
         r"""
         Allows users to call ``self.all_gather()`` from the LightningModule, thus making the ``all_gather`` operation
         accelerator agnostic. ``all_gather`` is a function provided by accelerators to gather a tensor from several
@@ -567,7 +565,7 @@ def all_gather(
         group = group if group is not None else torch.distributed.group.WORLD
         all_gather = self.trainer.strategy.all_gather
         data = convert_to_tensors(data, device=self.device)
-        return apply_to_collection(data, torch.Tensor, all_gather, group=group, sync_grads=sync_grads)
+        return apply_to_collection(data, Tensor, all_gather, group=group, sync_grads=sync_grads)
 
     def forward(self, *args, **kwargs) -> Any:
         r"""
@@ -1701,15 +1699,15 @@ def tbptt_split_batch(self, batch, split_size):
             if :paramref:`~pytorch_lightning.core.module.LightningModule.truncated_bptt_steps` > 0.
             Each returned batch split is passed separately to :meth:`training_step`.
         """
-        time_dims = [len(x[0]) for x in batch if isinstance(x, (torch.Tensor, collections.Sequence))]
+        time_dims = [len(x[0]) for x in batch if isinstance(x, (Tensor, collections.Sequence))]
         assert len(time_dims) >= 1, "Unable to determine batch time dimension"
         assert all(x == time_dims[0] for x in time_dims), "Batch time dimension length is ambiguous"
 
         splits = []
         for t in range(0, time_dims[0], split_size):
             batch_split = []
             for i, x in enumerate(batch):
-                if isinstance(x, torch.Tensor):
+                if isinstance(x, Tensor):
                     split_x = x[:, t : t + split_size]
                 elif isinstance(x, collections.Sequence):
                     split_x = [None] * len(x)

src/pytorch_lightning/lite/lite.py

@@ -342,8 +342,8 @@ def barrier(self, name: Optional[str] = None) -> None:
         self._strategy.barrier(name=name)
 
     def all_gather(
-        self, data: Union[torch.Tensor, Dict, List, Tuple], group: Optional[Any] = None, sync_grads: bool = False
-    ) -> Union[torch.Tensor, Dict, List, Tuple]:
+        self, data: Union[Tensor, Dict, List, Tuple], group: Optional[Any] = None, sync_grads: bool = False
+    ) -> Union[Tensor, Dict, List, Tuple]:
         r"""
         Gather tensors or collections of tensors from multiple processes.
 
@@ -358,7 +358,7 @@ def all_gather(
         """
         group = group if group is not None else torch.distributed.group.WORLD
         data = convert_to_tensors(data, device=self.device)
-        return apply_to_collection(data, torch.Tensor, self._strategy.all_gather, group=group, sync_grads=sync_grads)
+        return apply_to_collection(data, Tensor, self._strategy.all_gather, group=group, sync_grads=sync_grads)
 
     def broadcast(self, obj: TBroadcast, src: int = 0) -> TBroadcast:
         return self._strategy.broadcast(obj, src=src)

src/pytorch_lightning/loggers/comet.py

@@ -21,8 +21,7 @@
 from argparse import Namespace
 from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Union
 
-import torch
-from torch import is_tensor
+from torch import is_tensor, Tensor
 
 import pytorch_lightning as pl
 from pytorch_lightning.loggers.logger import Logger, rank_zero_experiment
@@ -241,7 +240,7 @@ def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
         self.experiment.log_parameters(params)
 
     @rank_zero_only
-    def log_metrics(self, metrics: Dict[str, Union[torch.Tensor, float]], step: Optional[int] = None) -> None:
+    def log_metrics(self, metrics: Dict[str, Union[Tensor, float]], step: Optional[int] = None) -> None:
         assert rank_zero_only.rank == 0, "experiment tried to log from global_rank != 0"
         # Comet.ml expects metrics to be a dictionary of detached tensors on CPU
         metrics_without_epoch = metrics.copy()