From 045951db790838d2eb72b89a2dfab40d9c20d48d Mon Sep 17 00:00:00 2001 From: Mel Massadian Date: Sat, 26 Jul 2025 22:20:25 +0200 Subject: [PATCH] =?UTF-8?q?feat:=20=E2=9C=A8=20add=20LazyProxyTensor?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit utility extended from an idea by @AustinMroz --- utils.py | 201 +++++++++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 196 insertions(+), 5 deletions(-) diff --git a/utils.py b/utils.py index 739e97c..88ee45a 100644 --- a/utils.py +++ b/utils.py @@ -8,11 +8,14 @@ import shutil import socket import subprocess import sys +import textwrap import uuid +import warnings from collections.abc import Callable, Sequence from enum import Enum from functools import reduce from pathlib import Path +from types import EllipsisType from typing import TypeVar from urllib.parse import urlparse @@ -461,8 +464,6 @@ def _run_command(shell_cmd, ignored_lines_start): print("Command executed successfully!") - - # endregion @@ -526,6 +527,196 @@ PIL_FILTER_MAP = { # region TENSOR Utilities + + +class LazyProxyTensor: + """Memory-efficient proxy that wrap a tensor but presents itself as a different dtype (e.g., float32). + + It mimics a torch.Tensor's read-only attributes and methods. Data conversion + and normalization happen lazily on access (e.g., via slicing), avoiding + the high memory cost of a full conversion. + + Supported source dtypes: + - torch.uint8 (normalized from [0, 255]) + - torch.uint16 (normalized from [0, 65535]) + - All float types (passed through, assumed to be in [0, 1] range)" + """ + + _source_tensor: torch.Tensor + _target_dtype: torch.dtype + _target_element_size: int + _scale_divisor: float + _warned_inefficient_access: bool + + def __init__( + self, source_tensor, target_dtype=torch.float32, target_device=None + ): + if not isinstance(source_tensor, torch.Tensor): + raise ValueError("Input must be a torch.Tensor.") + + self._source_tensor = source_tensor + self._target_dtype = target_dtype + self._target_device = ( + target_device + if target_device is not None + else source_tensor.device + ) + + # Determine the normalization divisor based on source dtype + # fmt: off + if source_tensor.dtype == torch.uint8: self._scale_divisor = 255.0 + elif source_tensor.dtype == torch.uint16: self._scale_divisor = 65535.0 + elif torch.is_floating_point(source_tensor): self._scale_divisor = 1.0 + else: raise ValueError(f"Unsupported source dtype for LazyProxyTensor: {source_tensor.dtype}") + # fmt: on + + self._target_element_size = torch.empty( + (), dtype=self._target_dtype + ).element_size() + self._warned_inefficient_access = False + + def is_contiguous(self, *args, **kwargs): + return self._source_tensor.is_contiguous(*args, **kwargs) + + def stride(self, *args, **kwargs): + return self._source_tensor.stride(*args, **kwargs) + + @property + def shape(self): + return self._source_tensor.shape + + @property + def requires_grad(self): + return False + + def nelement(self): + """Return the total number of elements in the (pretend) tensor.""" + return self._source_tensor.nelement() + + def element_size(self): + """Return the size in bytes of an individual (pretend) float element.""" + return self._target_element_size + + @property + def dtype(self): + return self._target_dtype + + @property + def device(self): + return self._target_device + + def __len__(self): + return self._source_tensor.shape[0] + + def __getitem__(self, key): + if ( + self._source_tensor.device != self._target_device + and not self._warned_inefficient_access + ): + warnings.warn( + "Inefficient access pattern detected for LazyProxyTensor. " + "You are slicing a device-proxied tensor, which causes slow, " + "repeated data transfers. For performance, use the .iter_chunks() method." + ) + self._warned_inefficient_access = True + + subset = self._source_tensor[key] + + return ( + subset.to(self._target_device).to(self._target_dtype) + / self._scale_divisor + ) + + # def __iter__(self): + # for i in range(len(self)): + # yield self[i] + + def iter_chunks(self, chunk_size=16): + for i in range(0, len(self), chunk_size): + chunk = self._source_tensor[i : i + chunk_size] + yield ( + chunk.to(self._target_device, non_blocking=True).to( + self._target_dtype + ) + / self._scale_divisor + ) + + def squeeze(self, dim: str | EllipsisType | None = None): + squeezed = self._source_tensor.squeeze(dim) + return LazyProxyTensor(squeezed, self._target_dtype) + + def unsqueeze(self, dim: int = 0): + unsqueezed = self._source_tensor.unsqueeze(dim) + return LazyProxyTensor(unsqueezed, self._target_dtype) + + def repeat(self, *sizes): + repeated = self._source_tensor.repeat(*sizes) + return LazyProxyTensor(repeated, self._target_dtype) + + def _format_mem_size(self, mem_bytes): + if mem_bytes > 1e9: + return f"{mem_bytes / 1e9:.2f} GB" + if mem_bytes > 1e6: + return f"{mem_bytes / 1e6:.2f} MB" + if mem_bytes > 1e3: + return f"{mem_bytes / 1e3:.2f} KB" + return f"{mem_bytes} B" + + def __repr__(self): + + actual_info = get_torch_tensor_info(self._source_tensor, name="Source") + target_info = get_torch_tensor_info(self, name="Target") + + info = f""" + {target_info} + + {actual_info} + """ + return textwrap.dedent(info).strip() + + +def get_torch_tensor_info( + tensor: torch.Tensor | LazyProxyTensor | np.ndarray, + *, + name: str | None = None, +): + mem_str = "N/A" + + is_tensor = isinstance(tensor, torch.Tensor | LazyProxyTensor) + if is_tensor: + mem_bytes = tensor.element_size() * tensor.nelement() + else: + mem_bytes = tensor.itemsize * tensor.size + + if mem_bytes > 1e9: + mem_str = f"{mem_bytes / 1e9:.2f} GB" + elif mem_bytes > 1e6: + mem_str = f"{mem_bytes / 1e6:.2f} MB" + elif mem_bytes > 1e3: + mem_str = f"{mem_bytes / 1e3:.2f} KB" + else: + mem_str = f"{mem_bytes} B" + + device = "N/A" + grad = "False" + type_name = name or "Tensor" if is_tensor else "Numpy Array" + + if is_tensor: + device = tensor.device + grad = str(tensor.requires_grad) + + text = f""" + {type_name} + shape: {tensor.shape} + dtype: {str(tensor.dtype).replace("torch.", "")} + device: {device} + requires grad: {grad} + memory: {mem_str} + """ + + return textwrap.dedent(text).strip() + + def to_numpy(image: torch.Tensor) -> npt.NDArray[np.uint8]: """Converts a tensor to a ndarray with proper scaling and type conversion.""" np_array = np.clip(255.0 * image.cpu().numpy(), 0, 255).astype(np.uint8) @@ -536,12 +727,12 @@ def handle_batch( tensor: torch.Tensor, func: Callable[[torch.Tensor], Image.Image | npt.NDArray[np.uint8]], ) -> list[Image.Image] | list[npt.NDArray[np.uint8]]: - """Handles batch processing for a given tensor and conversion function.""" + """Handle batch processing for a given tensor and conversion function.""" return [func(tensor[i]) for i in range(tensor.shape[0])] def tensor2pil(tensor: torch.Tensor) -> list[Image.Image]: - """Converts a batch of tensors to a list of PIL Images.""" + """Convert a batch of tensors to a list of PIL Images.""" def single_tensor2pil(t: torch.Tensor) -> Image.Image: np_array = to_numpy(t) @@ -558,7 +749,7 @@ def tensor2pil(tensor: torch.Tensor) -> list[Image.Image]: def pil2tensor(images: Image.Image | list[Image.Image]) -> torch.Tensor: - """Converts a PIL Image or a list of PIL Images to a tensor.""" + """Convert a PIL Image or a list of PIL Images to a tensor.""" def single_pil2tensor(image: Image.Image) -> torch.Tensor: np_image = np.array(image).astype(np.float32) / 255.0