Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,62 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from torch.nn.parameter import ( # usort: skip
|
||||
Buffer as Buffer,
|
||||
Parameter as Parameter,
|
||||
UninitializedBuffer as UninitializedBuffer,
|
||||
UninitializedParameter as UninitializedParameter,
|
||||
)
|
||||
from torch.nn.modules import * # usort: skip # noqa: F403
|
||||
from torch.nn import (
|
||||
attention as attention,
|
||||
functional as functional,
|
||||
init as init,
|
||||
modules as modules,
|
||||
parallel as parallel,
|
||||
parameter as parameter,
|
||||
utils as utils,
|
||||
)
|
||||
from torch.nn.parallel import DataParallel as DataParallel
|
||||
|
||||
|
||||
def factory_kwargs(kwargs):
|
||||
r"""Return a canonicalized dict of factory kwargs.
|
||||
|
||||
Given kwargs, returns a canonicalized dict of factory kwargs that can be directly passed
|
||||
to factory functions like torch.empty, or errors if unrecognized kwargs are present.
|
||||
|
||||
This function makes it simple to write code like this::
|
||||
|
||||
class MyModule(nn.Module):
|
||||
def __init__(self, **kwargs):
|
||||
factory_kwargs = torch.nn.factory_kwargs(kwargs)
|
||||
self.weight = Parameter(torch.empty(10, **factory_kwargs))
|
||||
|
||||
Why should you use this function instead of just passing `kwargs` along directly?
|
||||
|
||||
1. This function does error validation, so if there are unexpected kwargs we will
|
||||
immediately report an error, instead of deferring it to the factory call
|
||||
2. This function supports a special `factory_kwargs` argument, which can be used to
|
||||
explicitly specify a kwarg to be used for factory functions, in the event one of the
|
||||
factory kwargs conflicts with an already existing argument in the signature (e.g.
|
||||
in the signature ``def f(dtype, **kwargs)``, you can specify ``dtype`` for factory
|
||||
functions, as distinct from the dtype argument, by saying
|
||||
``f(dtype1, factory_kwargs={"dtype": dtype2})``)
|
||||
"""
|
||||
if kwargs is None:
|
||||
return {}
|
||||
simple_keys = {"device", "dtype", "memory_format"}
|
||||
expected_keys = simple_keys | {"factory_kwargs"}
|
||||
if not kwargs.keys() <= expected_keys:
|
||||
raise TypeError(f"unexpected kwargs {kwargs.keys() - expected_keys}")
|
||||
|
||||
# guarantee no input kwargs is untouched
|
||||
r = dict(kwargs.get("factory_kwargs", {}))
|
||||
for k in simple_keys:
|
||||
if k in kwargs:
|
||||
if k in r:
|
||||
raise TypeError(
|
||||
f"{k} specified twice, in **kwargs and in factory_kwargs"
|
||||
)
|
||||
r[k] = kwargs[k]
|
||||
|
||||
return r
|
||||
@@ -0,0 +1,60 @@
|
||||
import warnings
|
||||
|
||||
|
||||
# NB: Keep this file in sync with enums in aten/src/ATen/core/Reduction.h
|
||||
|
||||
|
||||
def get_enum(reduction: str) -> int:
|
||||
if reduction == "none":
|
||||
ret = 0
|
||||
elif reduction == "mean":
|
||||
ret = 1
|
||||
elif reduction == "elementwise_mean":
|
||||
warnings.warn(
|
||||
"reduction='elementwise_mean' is deprecated. "
|
||||
"Please use reduction='mean' instead.",
|
||||
stacklevel=2,
|
||||
)
|
||||
ret = 1
|
||||
elif reduction == "sum":
|
||||
ret = 2
|
||||
else:
|
||||
ret = -1 # TODO: remove once JIT exceptions support control flow
|
||||
raise ValueError(f"{reduction} is not a valid value for reduction")
|
||||
return ret
|
||||
|
||||
|
||||
# In order to support previous versions, accept boolean size_average and reduce
|
||||
# and convert them into the new constants for now
|
||||
|
||||
|
||||
# We use these functions in torch/legacy as well, in which case we'll silence the warning
|
||||
def legacy_get_string(
|
||||
size_average: bool | None,
|
||||
reduce: bool | None,
|
||||
emit_warning: bool = True,
|
||||
) -> str:
|
||||
warning = "size_average and reduce args will be deprecated, please use reduction='{}' instead."
|
||||
|
||||
if size_average is None:
|
||||
size_average = True
|
||||
if reduce is None:
|
||||
reduce = True
|
||||
|
||||
if size_average and reduce:
|
||||
ret = "mean"
|
||||
elif reduce:
|
||||
ret = "sum"
|
||||
else:
|
||||
ret = "none"
|
||||
if emit_warning:
|
||||
warnings.warn(warning.format(ret), stacklevel=2)
|
||||
return ret
|
||||
|
||||
|
||||
def legacy_get_enum(
|
||||
size_average: bool | None,
|
||||
reduce: bool | None,
|
||||
emit_warning: bool = True,
|
||||
) -> int:
|
||||
return get_enum(legacy_get_string(size_average, reduce, emit_warning))
|
||||
@@ -0,0 +1,197 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""This module contains functions and classes that alter the behavior of torch.nn.functional.scaled_dot_product_attention"""
|
||||
|
||||
import contextlib
|
||||
from collections.abc import Iterable
|
||||
from typing import Union
|
||||
from warnings import warn
|
||||
|
||||
import torch.backends.cuda
|
||||
from torch._C import _SDPBackend as SDPBackend
|
||||
from torch.backends.cuda import (
|
||||
can_use_efficient_attention,
|
||||
can_use_flash_attention,
|
||||
SDPAParams,
|
||||
)
|
||||
|
||||
|
||||
__all__: list[str] = [
|
||||
"SDPBackend",
|
||||
"sdpa_kernel",
|
||||
"WARN_FOR_UNFUSED_KERNELS",
|
||||
"register_flash_attention_impl",
|
||||
"activate_flash_attention_impl",
|
||||
"list_flash_attention_impls",
|
||||
"current_flash_attention_impl",
|
||||
"restore_flash_attention_impl",
|
||||
]
|
||||
|
||||
|
||||
# Note: [SDPA warnings]
|
||||
# TODO: Consider using this for sdpa regardless of subclasses
|
||||
# This only effects users of bias subclasses
|
||||
# If this is set to True, we will warn the user if they are not using the fused kernels
|
||||
# As well, it will raise warnings for all the reasons why the fused kernels can't be run.
|
||||
# To set this to True, run
|
||||
# torch.nn.attention.WARN_FOR_UNFUSED_KERNELS = True
|
||||
WARN_FOR_UNFUSED_KERNELS = False
|
||||
|
||||
|
||||
r"""An enum-like class that contains the different backends for scaled dot product attention.
|
||||
This backend class is designed to be used with the sdpa_kernel context manager.
|
||||
|
||||
The following Enums are available:
|
||||
- ERROR: An error occurred when trying to determine the backend.
|
||||
- MATH: The math backend for scaled dot product attention.
|
||||
- FLASH_ATTENTION: The flash attention backend for scaled dot product attention.
|
||||
- EFFICIENT_ATTENTION: The efficient attention backend for scaled dot product attention.
|
||||
- CUDNN_ATTENTION: The cuDNN backend for scaled dot product attention.
|
||||
- OVERRIDEABLE: The overridable backend for extension.
|
||||
|
||||
See :func:`torch.nn.attention.sdpa_kernel` for more details.
|
||||
|
||||
.. warning:: This class is in beta and subject to change.
|
||||
"""
|
||||
SDPBackend.__module__ = __name__
|
||||
SDPBackend.__name__ = "SDPBackend"
|
||||
|
||||
|
||||
def _raise_kernel_warnings(params: SDPAParams) -> None:
|
||||
"""
|
||||
If WARN_FOR_UNFUSED_KERNELS is set to True, this will raise warnings
|
||||
for all the reasons why the fused kernels can't be run. If using subclasses
|
||||
"""
|
||||
if WARN_FOR_UNFUSED_KERNELS:
|
||||
if not can_use_efficient_attention(params):
|
||||
warn("Efficient attention can't be used because:", stacklevel=2)
|
||||
can_use_efficient_attention(params, True)
|
||||
if not can_use_flash_attention(params):
|
||||
warn("Flash attention can't be used because:", stacklevel=2)
|
||||
can_use_flash_attention(params, True)
|
||||
|
||||
|
||||
_backend_names = {
|
||||
"cudnn": "CUDNN_ATTENTION",
|
||||
"flash": "FLASH_ATTENTION",
|
||||
"mem_efficient": "EFFICIENT_ATTENTION",
|
||||
"math": "MATH",
|
||||
"overrideable": "OVERRIDEABLE",
|
||||
}
|
||||
|
||||
|
||||
def _backend_from_string(name: str):
|
||||
return getattr(SDPBackend, name)
|
||||
|
||||
|
||||
def _cur_sdpa_kernel_backends(with_priority: bool = False):
|
||||
backends = []
|
||||
for name, val in _backend_names.items():
|
||||
if getattr(torch._C, f"_get_{name}_sdp_enabled")():
|
||||
backends.append(getattr(SDPBackend, val))
|
||||
if with_priority:
|
||||
curr_priority = torch._C._get_sdp_priority_order()
|
||||
backends = sorted(
|
||||
backends, key=lambda backend: curr_priority.index(int(backend))
|
||||
)
|
||||
return backends
|
||||
|
||||
|
||||
def _sdpa_kernel(backends: Iterable, set_priority: bool = False) -> None:
|
||||
for name, val in _backend_names.items():
|
||||
enabled = getattr(SDPBackend, val) in backends
|
||||
getattr(torch._C, f"_set_sdp_use_{name}")(enabled)
|
||||
if set_priority:
|
||||
# backends should be a unique list
|
||||
user_priority = [int(backend) for backend in backends]
|
||||
previous_priority = torch._C._get_sdp_priority_order()
|
||||
for backend in previous_priority:
|
||||
if backend not in user_priority:
|
||||
user_priority.append(int(backend))
|
||||
torch._C._set_sdp_priority_order(user_priority)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def sdpa_kernel(backends: list[SDPBackend] | SDPBackend, set_priority: bool = False):
|
||||
r"""
|
||||
Context manager to select which backend to use for scaled dot product attention.
|
||||
|
||||
.. warning:: This function is beta and subject to change.
|
||||
|
||||
Args:
|
||||
backends (Union[List[SDPBackend], SDPBackend]): A backend or list of backends for scaled dot product attention.
|
||||
set_priority (bool=False): Whether the ordering of the backends is interpreted as their priority order.
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from torch.nn.functional import scaled_dot_product_attention
|
||||
from torch.nn.attention import SDPBackend, sdpa_kernel
|
||||
|
||||
# Only enable flash attention backend
|
||||
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
|
||||
scaled_dot_product_attention(...)
|
||||
|
||||
# Enable the Math or Efficient attention backends
|
||||
with sdpa_kernel([SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION]):
|
||||
scaled_dot_product_attention(...)
|
||||
|
||||
# Enable the cuDNN or flash attention backends, and in that order
|
||||
with sdpa_kernel(
|
||||
[SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION], set_priority=True
|
||||
):
|
||||
scaled_dot_product_attention(...)
|
||||
|
||||
This context manager can be used to select which backend to use for scaled dot product attention.
|
||||
Upon exiting the context manager, the previous state of the flags will be restored, enabling all backends.
|
||||
"""
|
||||
if not isinstance(backends, (list, SDPBackend)):
|
||||
raise AssertionError(
|
||||
f"Backend must be an instance of SDPBackend or a list of SDPBackend instances, got {type(backends).__name__}"
|
||||
)
|
||||
|
||||
if isinstance(backends, SDPBackend):
|
||||
backends = [backends]
|
||||
|
||||
backends = list(dict.fromkeys(backends))
|
||||
|
||||
previous_backends = _cur_sdpa_kernel_backends(with_priority=set_priority)
|
||||
try:
|
||||
_sdpa_kernel(backends, set_priority)
|
||||
yield {}
|
||||
finally:
|
||||
_sdpa_kernel(previous_backends, set_priority)
|
||||
|
||||
|
||||
# variadic version of sdpa_kernel for dynamo to use while reconstructing
|
||||
@contextlib.contextmanager
|
||||
def _sdpa_kernel_variadic(*backends: SDPBackend):
|
||||
with sdpa_kernel(list(backends)):
|
||||
yield
|
||||
|
||||
|
||||
def _get_flash_version() -> str:
|
||||
"""This returns the closest matching tag for the flash attention backend"""
|
||||
return "2.5.7"
|
||||
|
||||
|
||||
from . import _registry
|
||||
|
||||
|
||||
# Re-export registry types and functions for public API
|
||||
_FlashAttentionImpl = _registry._FlashAttentionImpl
|
||||
_RegisterFn = _registry._RegisterFn
|
||||
register_flash_attention_impl = _registry.register_flash_attention_impl
|
||||
activate_flash_attention_impl = _registry.activate_flash_attention_impl
|
||||
list_flash_attention_impls = _registry.list_flash_attention_impls
|
||||
current_flash_attention_impl = _registry.current_flash_attention_impl
|
||||
restore_flash_attention_impl = _registry.restore_flash_attention_impl
|
||||
|
||||
register_flash_attention_impl.__module__ = __name__
|
||||
activate_flash_attention_impl.__module__ = __name__
|
||||
list_flash_attention_impls.__module__ = __name__
|
||||
current_flash_attention_impl.__module__ = __name__
|
||||
restore_flash_attention_impl.__module__ = __name__
|
||||
|
||||
# Import built-in implementations to trigger self-registration
|
||||
from . import _fa3, _fa4 # noqa: F401 # noqa: F401
|
||||
@@ -0,0 +1,763 @@
|
||||
"""
|
||||
PROTOTYPE!
|
||||
Flash Attention 3 implementation.
|
||||
For fp8: only supports forward pass right now.
|
||||
For fp16/bf16: supports forward and backward pass.
|
||||
"""
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import warnings
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
from dataclasses import dataclass
|
||||
from functools import cache
|
||||
from typing_extensions import TypeVarTuple, Unpack
|
||||
|
||||
import torch
|
||||
from torch.library import Library
|
||||
|
||||
from . import _registry
|
||||
|
||||
|
||||
__all__ = [
|
||||
"register_flash_attention_fa3",
|
||||
]
|
||||
|
||||
|
||||
_FA3_CUDA_FWD: Callable | None = None # Cache for torch.ops.flash_attn_3.fwd
|
||||
_FA3_CUDA_BWD: Callable | None = None # Cache for torch.ops.flash_attn_3.bwd
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FA3Handle:
|
||||
library: Library | None
|
||||
|
||||
def remove(self) -> None:
|
||||
self.library = None
|
||||
# Clear the C++ flag
|
||||
torch._C._set_sdp_use_fa3(False)
|
||||
|
||||
|
||||
@cache
|
||||
def _get_device_major(device: torch.device) -> int:
|
||||
major, _ = torch.cuda.get_device_capability(device)
|
||||
return major
|
||||
|
||||
|
||||
def register_flash_attention_fa3(
|
||||
module_path: str = "flash_attn_interface",
|
||||
) -> _FA3Handle:
|
||||
"""
|
||||
Register FA3 flash attention kernels with the PyTorch dispatcher.
|
||||
|
||||
Args:
|
||||
module_path: Python module path to the FA3 implementation.
|
||||
"""
|
||||
_fa3_import_module(module_path)
|
||||
|
||||
# Expose FA3 registration status to C++
|
||||
torch._C._set_sdp_use_fa3(True)
|
||||
|
||||
return _FA3Handle(_fa3_register_kernels())
|
||||
|
||||
|
||||
def _fa3_import_module(module_path: str) -> None:
|
||||
importlib.import_module(module_path)
|
||||
if not hasattr(torch.ops, "flash_attn_3"):
|
||||
raise RuntimeError(f"Module '{module_path}' does not expose FA3 kernels")
|
||||
if not hasattr(torch.ops.flash_attn_3, "fwd"):
|
||||
raise RuntimeError(
|
||||
f"Module '{module_path}' does not expose FA3 forward kernels"
|
||||
)
|
||||
if not hasattr(torch.ops.flash_attn_3, "bwd"):
|
||||
raise RuntimeError(
|
||||
f"Module '{module_path}' does not expose FA3 backward kernels"
|
||||
)
|
||||
global _FA3_CUDA_FWD, _FA3_CUDA_BWD
|
||||
_FA3_CUDA_FWD = torch.ops.flash_attn_3.fwd
|
||||
_FA3_CUDA_BWD = torch.ops.flash_attn_3.bwd
|
||||
|
||||
|
||||
def _fa3_register_kernels() -> Library:
|
||||
lib = Library("aten", "IMPL", "CUDA") # noqa: TOR901
|
||||
lib.impl(
|
||||
"_flash_attention_forward.quantized", _fa3_flash_attention_forward_impl, "CUDA"
|
||||
)
|
||||
lib.impl(
|
||||
"_scaled_dot_product_flash_attention.quantized",
|
||||
_fa3_scaled_dot_product_flash_attention_forward_impl,
|
||||
"CUDA",
|
||||
)
|
||||
lib.impl(
|
||||
"_flash_attention_forward", _fa3_flash_attention_forward_impl_default, "CUDA"
|
||||
)
|
||||
lib.impl(
|
||||
"_flash_attention_forward_no_dropout_inplace",
|
||||
_fa3_flash_attention_forward_no_dropout_inplace_impl,
|
||||
"CUDA",
|
||||
)
|
||||
lib.impl(
|
||||
"_scaled_dot_product_flash_attention",
|
||||
_fa3_scaled_dot_product_flash_attention_forward_impl_default,
|
||||
"CUDA",
|
||||
)
|
||||
|
||||
lib.impl("_flash_attention_backward", _fa3_flash_attention_backward_impl, "CUDA")
|
||||
lib.impl(
|
||||
"_scaled_dot_product_flash_attention_backward",
|
||||
_fa3_scaled_dot_product_flash_attention_backward_impl,
|
||||
"CUDA",
|
||||
)
|
||||
return lib
|
||||
|
||||
|
||||
def _fa3_common_support_error(
|
||||
query: torch.Tensor,
|
||||
tensors: tuple[torch.Tensor, ...],
|
||||
dropout_p: float,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
q_descale: torch.Tensor | None,
|
||||
k_descale: torch.Tensor | None,
|
||||
v_descale: torch.Tensor | None,
|
||||
) -> str | None:
|
||||
if dropout_p != 0.0:
|
||||
return "dropout_p must be 0"
|
||||
|
||||
if not all(t.is_cuda for t in tensors):
|
||||
return "inputs must be CUDA tensors"
|
||||
if len({t.device for t in tensors}) != 1:
|
||||
return "inputs must share device"
|
||||
if query.dtype == torch.float8_e4m3fn and (
|
||||
q_descale is None or k_descale is None or v_descale is None
|
||||
):
|
||||
warnings.warn(
|
||||
"When using SDPA with fp8, descale tensor should always be used"
|
||||
" for accurate dequantization. Please use "
|
||||
"_scaled_dot_product_attention_quantized and "
|
||||
"provide the descale tensors.",
|
||||
UserWarning,
|
||||
)
|
||||
if cum_seq_q is None and query.dim() != 4:
|
||||
return "dense query must be 4D"
|
||||
if cum_seq_q is not None and query.dim() != 3:
|
||||
return "ragged query must be 3D"
|
||||
if not torch.cuda.is_available():
|
||||
return "CUDA not available"
|
||||
if _get_device_major(query.device) != 9:
|
||||
return "FA3 requires compute capability 9.0"
|
||||
return None
|
||||
|
||||
|
||||
def _fa3_forward_support_error(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
dropout_p: float,
|
||||
return_debug_mask: bool,
|
||||
alibi_slopes: torch.Tensor | None,
|
||||
seqused_k: torch.Tensor | None,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
q_descale: torch.Tensor | None,
|
||||
k_descale: torch.Tensor | None,
|
||||
v_descale: torch.Tensor | None,
|
||||
) -> str | None:
|
||||
if return_debug_mask:
|
||||
return "return_debug_mask must be False"
|
||||
if alibi_slopes is not None:
|
||||
return "alibi_slopes not supported"
|
||||
if seqused_k is not None:
|
||||
if seqused_k.dtype != torch.int32:
|
||||
return "seqused_k must be int32"
|
||||
if not seqused_k.is_cuda:
|
||||
return "seqused_k must be CUDA"
|
||||
supported_dtypes = (torch.float8_e4m3fn, torch.float16, torch.bfloat16)
|
||||
if not all(t.dtype in supported_dtypes for t in {query, key, value}):
|
||||
return f"inputs must be one of {supported_dtypes}"
|
||||
if len({t.dtype for t in {query, key, value}}) != 1:
|
||||
return "all inputs must have the same dtype"
|
||||
error = _fa3_common_support_error(
|
||||
query,
|
||||
(query, key, value),
|
||||
dropout_p,
|
||||
cum_seq_q,
|
||||
q_descale,
|
||||
k_descale,
|
||||
v_descale,
|
||||
)
|
||||
if error is not None:
|
||||
if error == "inputs must share device":
|
||||
return "query, key, value must be on same device"
|
||||
return error
|
||||
return None
|
||||
|
||||
|
||||
def _fa3_backward_support_error(
|
||||
grad_out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
logsumexp: torch.Tensor,
|
||||
dropout_p: float,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
window_size_left: int | None,
|
||||
window_size_right: int | None,
|
||||
) -> str | None:
|
||||
# FA3 backward ONLY supports fp16/bf16, NOT fp8
|
||||
if query.dtype == torch.float8_e4m3fn:
|
||||
return (
|
||||
"FA3 backward does not support fp8 - use inference only (torch.no_grad())"
|
||||
)
|
||||
if logsumexp.dtype != torch.float32:
|
||||
return "logsumexp dtype must be float32"
|
||||
supported_dtypes = (torch.float16, torch.bfloat16)
|
||||
if not all(t.dtype in supported_dtypes for t in {grad_out, query, key, value, out}):
|
||||
return f"inputs must be one of {supported_dtypes}"
|
||||
if len({t.dtype for t in {grad_out, query, key, value, out}}) != 1:
|
||||
return "all inputs must have the same dtype"
|
||||
error = _fa3_common_support_error(
|
||||
query,
|
||||
(grad_out, query, key, value, out, logsumexp),
|
||||
dropout_p,
|
||||
cum_seq_q,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
if error is not None:
|
||||
return error
|
||||
return None
|
||||
|
||||
|
||||
Ts = TypeVarTuple("Ts")
|
||||
|
||||
|
||||
def _transpose_dense(*tensors: Unpack[Ts]) -> tuple[Unpack[Ts]]:
|
||||
return tuple(t.transpose(1, 2) for t in tensors) # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def _maybe_contiguous(x: torch.Tensor | None) -> torch.Tensor | None:
|
||||
"""Ensure tensor is contiguous in the last dimension."""
|
||||
return x.contiguous() if x is not None and x.stride(-1) != 1 else x
|
||||
|
||||
|
||||
def _fa3_run_forward(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor | None,
|
||||
cu_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
scale: float | None,
|
||||
is_causal: bool,
|
||||
window_size_left: int | None,
|
||||
window_size_right: int | None,
|
||||
seqused_k: torch.Tensor | None,
|
||||
out: torch.Tensor | None = None,
|
||||
q_descale: torch.Tensor | None = None,
|
||||
k_descale: torch.Tensor | None = None,
|
||||
v_descale: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Run the FA3 forward pass by calling the C++ kernel directly.
|
||||
"""
|
||||
if _FA3_CUDA_FWD is None:
|
||||
raise RuntimeError("FA3 not registered")
|
||||
# Ensure contiguous in the last dimension
|
||||
q = _maybe_contiguous(query)
|
||||
k = _maybe_contiguous(key)
|
||||
v = (
|
||||
value.contiguous()
|
||||
if value.dtype == torch.float8_e4m3fn
|
||||
and value.stride(-1) != 1
|
||||
and value.stride(-3) != 1
|
||||
else _maybe_contiguous(value)
|
||||
)
|
||||
|
||||
cu_seqlens_q = _maybe_contiguous(cu_seq_q)
|
||||
cu_seqlens_k = _maybe_contiguous(cu_seq_k)
|
||||
seqused_k = _maybe_contiguous(seqused_k)
|
||||
block_table = _maybe_contiguous(block_table)
|
||||
|
||||
out, softmax_lse, out_accum, softmax_lse_accum = _FA3_CUDA_FWD(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
None, # k_new
|
||||
None, # v_new
|
||||
None, # qv
|
||||
out, # out_ (pre-allocated output)
|
||||
cu_seqlens_q, # cu_seqlens_q
|
||||
cu_seqlens_k, # cu_seqlens_k
|
||||
None, # cu_seqlens_k_new
|
||||
None, # seqused_q
|
||||
seqused_k, # seqused_k
|
||||
max_q, # max_seqlen_q
|
||||
max_k, # max_seqlen_k
|
||||
block_table, # block_table,
|
||||
None, # kv_batch_idx,
|
||||
None, # leftpad_k,
|
||||
None, # rotary_cos,
|
||||
None, # rotary_sin,
|
||||
None, # seqlens_rotary,
|
||||
q_descale, # q_descale,
|
||||
k_descale, # k_descale,
|
||||
v_descale, # v_descale,
|
||||
scale, # softmax_scale,
|
||||
is_causal, # causal,
|
||||
window_size_left if window_size_left is not None else -1, # window_size_left
|
||||
window_size_right if window_size_right is not None else -1, # window_size_right
|
||||
0, # attention_chunk,
|
||||
0.0, # softcap,
|
||||
True, # rotary_interleaved,
|
||||
None, # scheduler_metadata,
|
||||
num_splits
|
||||
or (1 if torch.are_deterministic_algorithms_enabled() else 0), # num_splits,
|
||||
None, # pack_gqa,
|
||||
torch._C._get_sm_carveout_experimental() or 0, # sm_margin,
|
||||
)
|
||||
return out, softmax_lse.contiguous()
|
||||
|
||||
|
||||
def _fa3_run_backward(
|
||||
grad_out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
logsumexp: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor | None,
|
||||
cu_seq_k: torch.Tensor | None,
|
||||
max_seqlen_q: int | None,
|
||||
max_seqlen_k: int | None,
|
||||
scale: float | None,
|
||||
is_causal: bool,
|
||||
window_size_left: int,
|
||||
window_size_right: int,
|
||||
deterministic: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
if _FA3_CUDA_BWD is None:
|
||||
raise RuntimeError("FA3 not registered")
|
||||
|
||||
# Ensure contiguous
|
||||
dout = _maybe_contiguous(grad_out)
|
||||
q = query.contiguous() if query.stride(-1) != 1 else query
|
||||
k = key.contiguous() if key.stride(-1) != 1 else key
|
||||
v = value.contiguous() if value.stride(-1) != 1 else value
|
||||
o = _maybe_contiguous(out)
|
||||
lse = _maybe_contiguous(logsumexp)
|
||||
|
||||
# Pre-allocate gradient tensors
|
||||
dq = torch.empty_like(q)
|
||||
dk = torch.empty_like(k)
|
||||
dv = torch.empty_like(v)
|
||||
_FA3_CUDA_BWD(
|
||||
dout,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
o,
|
||||
lse,
|
||||
dq,
|
||||
dk,
|
||||
dv,
|
||||
cu_seq_q,
|
||||
cu_seq_k,
|
||||
None,
|
||||
None,
|
||||
max_seqlen_q,
|
||||
max_seqlen_k,
|
||||
scale,
|
||||
is_causal,
|
||||
window_size_left,
|
||||
window_size_right,
|
||||
0.0,
|
||||
deterministic,
|
||||
torch._C._get_sm_carveout_experimental() or 0,
|
||||
)
|
||||
return dq, dk, dv
|
||||
|
||||
|
||||
def _fa3_flash_attention_forward_impl(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
cum_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
dropout_p: float,
|
||||
is_causal: bool,
|
||||
return_debug_mask: bool,
|
||||
q_descale: torch.Tensor | None = None,
|
||||
k_descale: torch.Tensor | None = None,
|
||||
v_descale: torch.Tensor | None = None,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
window_size_left: int = -1,
|
||||
window_size_right: int = -1,
|
||||
seqused_k: torch.Tensor | None = None,
|
||||
alibi_slopes: torch.Tensor | None = None,
|
||||
out: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
compute_auxiliary: bool = True,
|
||||
num_splits: int | None = None,
|
||||
):
|
||||
error = _fa3_forward_support_error(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
dropout_p,
|
||||
return_debug_mask,
|
||||
alibi_slopes,
|
||||
seqused_k,
|
||||
cum_seq_q,
|
||||
q_descale,
|
||||
k_descale,
|
||||
v_descale,
|
||||
)
|
||||
if error is not None:
|
||||
raise RuntimeError(f"FA3 flash_attention forward unsupported: {error}")
|
||||
out, lse = _fa3_run_forward(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cum_seq_q,
|
||||
cum_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
scale,
|
||||
is_causal,
|
||||
window_size_left,
|
||||
window_size_right,
|
||||
seqused_k,
|
||||
out,
|
||||
q_descale,
|
||||
k_descale,
|
||||
v_descale,
|
||||
block_table,
|
||||
num_splits,
|
||||
)
|
||||
if compute_auxiliary:
|
||||
rng_state = torch.zeros((2,), dtype=torch.uint64, device=query.device)
|
||||
philox_offset = torch.zeros((), dtype=torch.uint64, device=query.device)
|
||||
debug_mask = torch.empty(0, dtype=query.dtype, device=query.device)
|
||||
else:
|
||||
rng_state = None
|
||||
philox_offset = None
|
||||
debug_mask = None
|
||||
return out, lse, rng_state, philox_offset, debug_mask
|
||||
|
||||
|
||||
def _fa3_flash_attention_forward_no_dropout_inplace_impl(
|
||||
out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
cum_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
dropout_p: float,
|
||||
is_causal: bool,
|
||||
return_debug_mask: bool,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
window_size_left: int = -1,
|
||||
window_size_right: int = -1,
|
||||
seqused_k: torch.Tensor | None = None,
|
||||
alibi_slopes: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
):
|
||||
_, lse, _, _, _ = _fa3_flash_attention_forward_impl(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cum_seq_q,
|
||||
cum_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
dropout_p,
|
||||
is_causal,
|
||||
return_debug_mask,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
scale=scale,
|
||||
window_size_left=window_size_left,
|
||||
window_size_right=window_size_right,
|
||||
seqused_k=seqused_k,
|
||||
alibi_slopes=alibi_slopes,
|
||||
out=out,
|
||||
block_table=block_table,
|
||||
compute_auxiliary=False,
|
||||
num_splits=num_splits,
|
||||
)
|
||||
return lse
|
||||
|
||||
|
||||
def _fa3_flash_attention_forward_impl_default(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
cum_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
dropout_p: float,
|
||||
is_causal: bool,
|
||||
return_debug_mask: bool,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
window_size_left: int = -1,
|
||||
window_size_right: int = -1,
|
||||
seqused_k: torch.Tensor | None = None,
|
||||
alibi_slopes: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
out: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
):
|
||||
return _fa3_flash_attention_forward_impl(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cum_seq_q,
|
||||
cum_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
dropout_p,
|
||||
is_causal,
|
||||
return_debug_mask,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
scale=scale,
|
||||
window_size_left=window_size_left,
|
||||
window_size_right=window_size_right,
|
||||
seqused_k=seqused_k,
|
||||
alibi_slopes=alibi_slopes,
|
||||
out=out,
|
||||
block_table=block_table,
|
||||
num_splits=num_splits,
|
||||
)
|
||||
|
||||
|
||||
def _fa3_flash_attention_backward_impl(
|
||||
grad_out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
logsumexp: torch.Tensor,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
cum_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
dropout_p: float,
|
||||
is_causal: bool,
|
||||
rng_state: torch.Tensor,
|
||||
unused: torch.Tensor,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
window_size_left: int | None = None,
|
||||
window_size_right: int | None = None,
|
||||
):
|
||||
"""FA3 implementation of _flash_attention_backward."""
|
||||
error = _fa3_backward_support_error(
|
||||
grad_out,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
out,
|
||||
logsumexp,
|
||||
dropout_p,
|
||||
cum_seq_q,
|
||||
window_size_left,
|
||||
window_size_right,
|
||||
)
|
||||
|
||||
if error is not None:
|
||||
raise RuntimeError(f"FA3 flash_attention backward unsupported: {error}")
|
||||
|
||||
deterministic = torch.are_deterministic_algorithms_enabled()
|
||||
|
||||
dq, dk, dv = _fa3_run_backward(
|
||||
grad_out,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
out,
|
||||
logsumexp,
|
||||
cum_seq_q,
|
||||
cum_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
scale,
|
||||
is_causal,
|
||||
window_size_left if window_size_left is not None else -1,
|
||||
window_size_right if window_size_right is not None else -1,
|
||||
deterministic,
|
||||
)
|
||||
return dq, dk, dv
|
||||
|
||||
|
||||
def _fa3_scaled_dot_product_flash_attention_forward_impl(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
q_descale: torch.Tensor | None = None,
|
||||
k_descale: torch.Tensor | None = None,
|
||||
v_descale: torch.Tensor | None = None,
|
||||
dropout_p: float = 0.0,
|
||||
is_causal: bool = False,
|
||||
return_debug_mask: bool = False,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
):
|
||||
error = _fa3_forward_support_error(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
dropout_p,
|
||||
return_debug_mask,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
q_descale,
|
||||
k_descale,
|
||||
v_descale,
|
||||
)
|
||||
if error is not None:
|
||||
raise RuntimeError(f"FA3 SDPA forward unsupported: {error}")
|
||||
q, k, v = _transpose_dense(query, key, value)
|
||||
|
||||
# Pre-allocate output with query's strides (BHSD layout), then create
|
||||
# a BSHD view for the kernel. This ensures the returned output has
|
||||
# the same memory layout as the input query.
|
||||
out_dtype = torch.bfloat16 if query.dtype == torch.float8_e4m3fn else query.dtype
|
||||
out_bhsd = torch.empty_like(query, dtype=out_dtype)
|
||||
out_bshd = out_bhsd.transpose(1, 2)
|
||||
|
||||
max_q_flash = q.size(1)
|
||||
max_k_flash = k.size(1)
|
||||
_, lse, rng_state, philox_offset, debug_mask = _fa3_flash_attention_forward_impl(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
None,
|
||||
None,
|
||||
max_q_flash,
|
||||
max_k_flash,
|
||||
dropout_p,
|
||||
is_causal,
|
||||
return_debug_mask,
|
||||
scale=scale,
|
||||
out=out_bshd,
|
||||
q_descale=q_descale,
|
||||
k_descale=k_descale,
|
||||
v_descale=v_descale,
|
||||
)
|
||||
max_q = query.size(2)
|
||||
max_k = key.size(2)
|
||||
return (
|
||||
out_bhsd,
|
||||
lse,
|
||||
None,
|
||||
None,
|
||||
max_q,
|
||||
max_k,
|
||||
rng_state,
|
||||
philox_offset,
|
||||
debug_mask,
|
||||
)
|
||||
|
||||
|
||||
def _fa3_scaled_dot_product_flash_attention_forward_impl_default(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
dropout_p: float = 0.0,
|
||||
is_causal: bool = False,
|
||||
return_debug_mask: bool = False,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
):
|
||||
return _fa3_scaled_dot_product_flash_attention_forward_impl(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
dropout_p,
|
||||
is_causal,
|
||||
return_debug_mask,
|
||||
scale=scale,
|
||||
)
|
||||
|
||||
|
||||
def _fa3_scaled_dot_product_flash_attention_backward_impl(
|
||||
grad_out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
logsumexp: torch.Tensor,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
cum_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
dropout_p: float,
|
||||
is_causal: bool,
|
||||
philox_seed: torch.Tensor,
|
||||
philox_offset: torch.Tensor,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
):
|
||||
"""FA3 implementation of _scaled_dot_product_flash_attention_backward."""
|
||||
error = _fa3_backward_support_error(
|
||||
grad_out, query, key, value, out, logsumexp, dropout_p, None, None, None
|
||||
)
|
||||
if error is not None:
|
||||
raise RuntimeError(f"FA3 SDPA backward unsupported: {error}")
|
||||
|
||||
# SDPA uses BHSD layout, FA3 uses BSHD - transpose
|
||||
grad_out_t, q_t, k_t, v_t, out_t = _transpose_dense(
|
||||
grad_out, query, key, value, out
|
||||
)
|
||||
|
||||
dq, dk, dv = _fa3_flash_attention_backward_impl(
|
||||
grad_out_t,
|
||||
q_t,
|
||||
k_t,
|
||||
v_t,
|
||||
out_t,
|
||||
logsumexp,
|
||||
None, # cum_seq_q (dense attention)
|
||||
None, # cum_seq_k
|
||||
max_q, # max_seqlen_q
|
||||
max_k, # max_seqlen_k
|
||||
dropout_p,
|
||||
is_causal,
|
||||
philox_seed,
|
||||
philox_offset,
|
||||
scale=scale,
|
||||
)
|
||||
|
||||
# Transpose gradients back to BHSD layout
|
||||
dq_out, dk_out, dv_out = _transpose_dense(dq, dk, dv)
|
||||
return dq_out, dk_out, dv_out
|
||||
|
||||
|
||||
_registry.register_flash_attention_impl("FA3", register_fn=register_flash_attention_fa3)
|
||||
@@ -0,0 +1,541 @@
|
||||
"""UBER PROTOTYPE!!!"""
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
from dataclasses import dataclass
|
||||
from functools import cache
|
||||
from typing import Any, TYPE_CHECKING
|
||||
from typing_extensions import TypeVarTuple, Unpack
|
||||
|
||||
from . import _registry
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from types import ModuleType
|
||||
|
||||
import torch
|
||||
from torch.library import Library
|
||||
|
||||
|
||||
__all__ = [
|
||||
"register_flash_attention_fa4",
|
||||
]
|
||||
|
||||
|
||||
_FA4_MODULE_PATH: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FA4Handle:
|
||||
library: Library | None
|
||||
|
||||
def remove(self) -> None:
|
||||
self.library = None
|
||||
|
||||
|
||||
@cache
|
||||
def _get_device_major(device: torch.device) -> int:
|
||||
major, _ = torch.cuda.get_device_capability(device)
|
||||
return major
|
||||
|
||||
|
||||
def register_flash_attention_fa4(
|
||||
module_path: str = "flash_attn.cute.interface",
|
||||
) -> _FA4Handle:
|
||||
"""
|
||||
Register FA4 flash attention kernels with the PyTorch dispatcher.
|
||||
|
||||
Args:
|
||||
module_path: Python module path to the FA4 implementation.
|
||||
"""
|
||||
global _FA4_MODULE_PATH
|
||||
_ = _fa4_import_module(module_path)
|
||||
_FA4_MODULE_PATH = module_path
|
||||
return _FA4Handle(_fa4_register_kernels())
|
||||
|
||||
|
||||
@cache
|
||||
def _fa4_import_module(module_path: str) -> ModuleType:
|
||||
module = importlib.import_module(module_path)
|
||||
if not hasattr(module, "_flash_attn_fwd") or not hasattr(module, "_flash_attn_bwd"):
|
||||
raise RuntimeError(f"Module '{module_path}' does not expose FA4 kernels")
|
||||
return module
|
||||
|
||||
|
||||
def _fa4_register_kernels() -> Library:
|
||||
lib = Library("aten", "IMPL", "CUDA") # noqa: TOR901
|
||||
lib.impl("_flash_attention_forward", _fa4_flash_attention_forward_impl, "CUDA")
|
||||
lib.impl(
|
||||
"_flash_attention_forward_no_dropout_inplace",
|
||||
_fa4_flash_attention_forward_no_dropout_inplace_impl,
|
||||
"CUDA",
|
||||
)
|
||||
lib.impl("_flash_attention_backward", _fa4_flash_attention_backward_impl, "CUDA")
|
||||
lib.impl(
|
||||
"_scaled_dot_product_flash_attention",
|
||||
_fa4_scaled_dot_product_flash_attention_forward_impl,
|
||||
"CUDA",
|
||||
)
|
||||
lib.impl(
|
||||
"_scaled_dot_product_flash_attention_backward",
|
||||
_fa4_scaled_dot_product_flash_attention_backward_impl,
|
||||
"CUDA",
|
||||
)
|
||||
return lib
|
||||
|
||||
|
||||
def _fa4_common_support_error(
|
||||
query: torch.Tensor,
|
||||
tensors: tuple[torch.Tensor, ...],
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
require_fp32: tuple[tuple[str, torch.Tensor], ...] = (),
|
||||
) -> str | None:
|
||||
if not all(t.is_cuda for t in tensors):
|
||||
return "inputs must be CUDA tensors"
|
||||
if len({t.device for t in tensors}) != 1:
|
||||
return "inputs must share device"
|
||||
if query.dtype not in (torch.float16, torch.bfloat16):
|
||||
return "query dtype must be float16 or bfloat16"
|
||||
for name, tensor in require_fp32:
|
||||
if tensor.dtype != torch.float32:
|
||||
return f"{name} dtype must be float32"
|
||||
if cum_seq_q is None and query.dim() != 4:
|
||||
return "dense query must be 4D"
|
||||
if cum_seq_q is not None and query.dim() != 3:
|
||||
return "ragged query must be 3D"
|
||||
if not torch.cuda.is_available():
|
||||
return "CUDA not available"
|
||||
if _get_device_major(query.device) not in (9, 10):
|
||||
return "FA4 requires compute capability 9.0 or 10.0"
|
||||
return None
|
||||
|
||||
|
||||
def _fa4_forward_support_error(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
dropout_p: float,
|
||||
return_debug_mask: bool,
|
||||
alibi_slopes: torch.Tensor | None,
|
||||
seqused_k: torch.Tensor | None,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
) -> str | None:
|
||||
if dropout_p != 0.0:
|
||||
return "dropout_p must be 0"
|
||||
if return_debug_mask:
|
||||
return "return_debug_mask must be False"
|
||||
if alibi_slopes is not None:
|
||||
return "alibi_slopes not supported"
|
||||
if seqused_k is not None:
|
||||
if seqused_k.dtype != torch.int32:
|
||||
return "seqused_k must be int32"
|
||||
if not seqused_k.is_cuda:
|
||||
return "seqused_k must be CUDA"
|
||||
major = _get_device_major(query.device)
|
||||
if block_table is not None and major != 10:
|
||||
return f"paged KV (block_table) not supported on SM {major}0"
|
||||
if num_splits is not None and num_splits > 1 and major != 10:
|
||||
return f"SplitKV (num_splits > 1) not supported on SM {major}0"
|
||||
error = _fa4_common_support_error(
|
||||
query,
|
||||
(query, key, value),
|
||||
cum_seq_q,
|
||||
)
|
||||
if error is not None:
|
||||
if error == "inputs must share device":
|
||||
return "query, key, value must be on same device"
|
||||
return error
|
||||
return None
|
||||
|
||||
|
||||
def _fa4_backward_support_error(
|
||||
grad_out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
logsumexp: torch.Tensor,
|
||||
dropout_p: float,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
) -> str | None:
|
||||
if dropout_p != 0.0:
|
||||
return "dropout_p must be 0"
|
||||
error = _fa4_common_support_error(
|
||||
query,
|
||||
(grad_out, query, key, value, out, logsumexp),
|
||||
cum_seq_q,
|
||||
require_fp32=(("logsumexp", logsumexp),),
|
||||
)
|
||||
if error is not None:
|
||||
return error
|
||||
return None
|
||||
|
||||
|
||||
def _aten_to_fa4_window_size(val: int | None) -> int | None:
|
||||
"""need to convert -1 to None for FA4"""
|
||||
return None if val == -1 else val
|
||||
|
||||
|
||||
Ts = TypeVarTuple("Ts")
|
||||
|
||||
|
||||
def _transpose_dense(*tensors: Unpack[Ts]) -> tuple[Unpack[Ts]]:
|
||||
return tuple(t.transpose(1, 2) for t in tensors) # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def _fa4_run_forward(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor | None,
|
||||
cu_seq_k: torch.Tensor | None,
|
||||
max_q: int | None,
|
||||
max_k: int | None,
|
||||
scale: float | None,
|
||||
is_causal: bool,
|
||||
window_size_left: int | None,
|
||||
window_size_right: int | None,
|
||||
seqused_k: torch.Tensor | None,
|
||||
out: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
if _FA4_MODULE_PATH is None:
|
||||
raise RuntimeError("FA4 not registered")
|
||||
module = _fa4_import_module(_FA4_MODULE_PATH)
|
||||
|
||||
kwargs: dict[str, Any] = {
|
||||
"softmax_scale": scale,
|
||||
"causal": is_causal,
|
||||
"window_size_left": _aten_to_fa4_window_size(window_size_left),
|
||||
"window_size_right": _aten_to_fa4_window_size(window_size_right),
|
||||
"return_lse": True,
|
||||
"cu_seqlens_q": cu_seq_q,
|
||||
"cu_seqlens_k": cu_seq_k,
|
||||
"max_seqlen_q": max_q,
|
||||
"max_seqlen_k": max_k,
|
||||
"seqused_k": seqused_k.contiguous() if seqused_k is not None else None,
|
||||
"page_table": block_table,
|
||||
"num_splits": num_splits or 1,
|
||||
"out": out,
|
||||
}
|
||||
out, lse = module._flash_attn_fwd(query, key, value, **kwargs)
|
||||
return out, lse.contiguous()
|
||||
|
||||
|
||||
def _fa4_run_backward(
|
||||
grad_out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
logsumexp: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor | None,
|
||||
cu_seq_k: torch.Tensor | None,
|
||||
scale: float | None,
|
||||
is_causal: bool,
|
||||
window_size_left: int | None,
|
||||
window_size_right: int | None,
|
||||
deterministic: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
if _FA4_MODULE_PATH is None:
|
||||
raise RuntimeError("FA4 not registered")
|
||||
module = _fa4_import_module(_FA4_MODULE_PATH)
|
||||
dq, dk, dv = module._flash_attn_bwd(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
out,
|
||||
grad_out,
|
||||
logsumexp.contiguous(),
|
||||
softmax_scale=scale,
|
||||
causal=is_causal,
|
||||
window_size_left=_aten_to_fa4_window_size(window_size_left),
|
||||
window_size_right=_aten_to_fa4_window_size(window_size_right),
|
||||
cu_seqlens_q=cu_seq_q,
|
||||
cu_seqlens_k=cu_seq_k,
|
||||
deterministic=deterministic,
|
||||
)
|
||||
return dq, dk, dv
|
||||
|
||||
|
||||
def _fa4_flash_attention_forward_impl(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
cum_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
dropout_p: float,
|
||||
is_causal: bool,
|
||||
return_debug_mask: bool,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
window_size_left: int | None = None,
|
||||
window_size_right: int | None = None,
|
||||
seqused_k: torch.Tensor | None = None,
|
||||
alibi_slopes: torch.Tensor | None = None,
|
||||
out: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
compute_auxiliary: bool = True,
|
||||
num_splits: int | None = None,
|
||||
):
|
||||
error = _fa4_forward_support_error(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
dropout_p,
|
||||
return_debug_mask,
|
||||
alibi_slopes,
|
||||
seqused_k,
|
||||
cum_seq_q,
|
||||
block_table,
|
||||
num_splits,
|
||||
)
|
||||
if error is not None:
|
||||
raise RuntimeError(f"FA4 flash_attention forward unsupported: {error}")
|
||||
out, lse = _fa4_run_forward(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cum_seq_q,
|
||||
cum_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
scale,
|
||||
is_causal,
|
||||
window_size_left,
|
||||
window_size_right,
|
||||
seqused_k,
|
||||
out,
|
||||
block_table,
|
||||
num_splits,
|
||||
)
|
||||
if compute_auxiliary:
|
||||
rng_state = torch.zeros((2,), dtype=torch.uint64, device=query.device)
|
||||
philox_offset = torch.zeros((), dtype=torch.uint64, device=query.device)
|
||||
debug_mask = torch.empty(0, dtype=query.dtype, device=query.device)
|
||||
else:
|
||||
rng_state = None
|
||||
philox_offset = None
|
||||
debug_mask = None
|
||||
return out, lse, rng_state, philox_offset, debug_mask
|
||||
|
||||
|
||||
def _fa4_flash_attention_forward_no_dropout_inplace_impl(
|
||||
out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
cum_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
dropout_p: float,
|
||||
is_causal: bool,
|
||||
return_debug_mask: bool,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
window_size_left: int | None = None,
|
||||
window_size_right: int | None = None,
|
||||
seqused_k: torch.Tensor | None = None,
|
||||
alibi_slopes: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
):
|
||||
_, lse, _, _, _ = _fa4_flash_attention_forward_impl(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cum_seq_q,
|
||||
cum_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
dropout_p,
|
||||
is_causal,
|
||||
return_debug_mask,
|
||||
scale=scale,
|
||||
window_size_left=window_size_left,
|
||||
window_size_right=window_size_right,
|
||||
seqused_k=seqused_k,
|
||||
alibi_slopes=alibi_slopes,
|
||||
out=out,
|
||||
block_table=block_table,
|
||||
compute_auxiliary=False,
|
||||
num_splits=num_splits,
|
||||
)
|
||||
return lse
|
||||
|
||||
|
||||
def _fa4_flash_attention_backward_impl(
|
||||
grad_out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
logsumexp: torch.Tensor,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
cum_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
dropout_p: float,
|
||||
is_causal: bool,
|
||||
rng_state: torch.Tensor,
|
||||
unused: torch.Tensor,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
window_size_left: int | None = None,
|
||||
window_size_right: int | None = None,
|
||||
):
|
||||
error = _fa4_backward_support_error(
|
||||
grad_out,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
out,
|
||||
logsumexp,
|
||||
dropout_p,
|
||||
cum_seq_q,
|
||||
)
|
||||
if error is not None:
|
||||
raise RuntimeError(f"FA4 flash_attention backward unsupported: {error}")
|
||||
deterministic = torch.are_deterministic_algorithms_enabled()
|
||||
dq, dk, dv = _fa4_run_backward(
|
||||
grad_out,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
out,
|
||||
logsumexp,
|
||||
cum_seq_q,
|
||||
cum_seq_k,
|
||||
scale,
|
||||
is_causal,
|
||||
window_size_left,
|
||||
window_size_right,
|
||||
deterministic,
|
||||
)
|
||||
return dq, dk, dv
|
||||
|
||||
|
||||
def _fa4_scaled_dot_product_flash_attention_forward_impl(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
dropout_p: float = 0.0,
|
||||
is_causal: bool = False,
|
||||
return_debug_mask: bool = False,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
):
|
||||
error = _fa4_forward_support_error(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
dropout_p,
|
||||
return_debug_mask,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
if error is not None:
|
||||
raise RuntimeError(f"FA4 SDPA forward unsupported: {error}")
|
||||
q, k, v = _transpose_dense(query, key, value)
|
||||
|
||||
# Pre-allocate output with query's strides (BHSD layout), then create
|
||||
# a BSHD view for the kernel. This ensures the returned output has
|
||||
# the same memory layout as the input query.
|
||||
out_bhsd = torch.empty_like(query)
|
||||
out_bshd = out_bhsd.transpose(1, 2)
|
||||
|
||||
max_q_flash = q.size(1)
|
||||
max_k_flash = k.size(1)
|
||||
_, lse, rng_state, philox_offset, debug_mask = _fa4_flash_attention_forward_impl(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
None,
|
||||
None,
|
||||
max_q_flash,
|
||||
max_k_flash,
|
||||
dropout_p,
|
||||
is_causal,
|
||||
return_debug_mask,
|
||||
scale=scale,
|
||||
out=out_bshd,
|
||||
)
|
||||
max_q = query.size(2)
|
||||
max_k = key.size(2)
|
||||
return (
|
||||
out_bhsd,
|
||||
lse,
|
||||
None,
|
||||
None,
|
||||
max_q,
|
||||
max_k,
|
||||
rng_state,
|
||||
philox_offset,
|
||||
debug_mask,
|
||||
)
|
||||
|
||||
|
||||
def _fa4_scaled_dot_product_flash_attention_backward_impl(
|
||||
grad_out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
logsumexp: torch.Tensor,
|
||||
cum_seq_q: torch.Tensor | None,
|
||||
cum_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
dropout_p: float,
|
||||
is_causal: bool,
|
||||
philox_seed: torch.Tensor,
|
||||
philox_offset: torch.Tensor,
|
||||
*,
|
||||
scale: float | None = None,
|
||||
):
|
||||
error = _fa4_backward_support_error(
|
||||
grad_out,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
out,
|
||||
logsumexp,
|
||||
dropout_p,
|
||||
None,
|
||||
)
|
||||
if error is not None:
|
||||
raise RuntimeError(f"FA4 SDPA backward unsupported: {error}")
|
||||
q, k, v, o, go = _transpose_dense(query, key, value, out, grad_out)
|
||||
max_q = query.size(2)
|
||||
max_k = key.size(2)
|
||||
dq, dk, dv = _fa4_flash_attention_backward_impl(
|
||||
go,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
o,
|
||||
logsumexp,
|
||||
None,
|
||||
None,
|
||||
max_q,
|
||||
max_k,
|
||||
dropout_p,
|
||||
is_causal,
|
||||
philox_seed,
|
||||
philox_offset,
|
||||
scale=scale,
|
||||
)
|
||||
dq, dk, dv = _transpose_dense(dq, dk, dv)
|
||||
return dq, dk, dv
|
||||
|
||||
|
||||
_registry.register_flash_attention_impl("FA4", register_fn=register_flash_attention_fa4)
|
||||
@@ -0,0 +1,137 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""Registry for flash attention implementations.
|
||||
|
||||
This module contains the registration system for flash attention implementations.
|
||||
It has no torch dependencies to avoid circular imports during initialization.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from typing import Literal, Protocol
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class FlashAttentionHandle(Protocol):
|
||||
def remove(self) -> None: ...
|
||||
|
||||
|
||||
_RegisterFn = Callable[..., FlashAttentionHandle | None]
|
||||
_FlashAttentionImpl = Literal["FA3", "FA4"]
|
||||
|
||||
_FLASH_ATTENTION_IMPLS: dict[str, _RegisterFn] = {}
|
||||
|
||||
_FLASH_ATTENTION_ACTIVE: tuple[str, FlashAttentionHandle] | None = None
|
||||
|
||||
|
||||
def register_flash_attention_impl(
|
||||
impl: str | _FlashAttentionImpl,
|
||||
*,
|
||||
register_fn: _RegisterFn,
|
||||
) -> None:
|
||||
"""
|
||||
Register the callable that activates a flash attention impl.
|
||||
|
||||
.. note::
|
||||
This function is intended for SDPA backend providers to register their
|
||||
implementations. End users should use :func:`activate_flash_attention_impl`
|
||||
to activate a registered implementation.
|
||||
|
||||
Args:
|
||||
impl: Implementation identifier (e.g., ``"FA4"``).
|
||||
register_fn: Callable that performs the actual dispatcher registration.
|
||||
This function will be invoked by :func:`activate_flash_attention_impl`
|
||||
and should register custom kernels with the PyTorch dispatcher.
|
||||
It may optionally return a handle implementing
|
||||
:class:`FlashAttentionHandle` to keep any necessary state alive.
|
||||
|
||||
Example:
|
||||
>>> def my_impl_register(module_path: str = "my_flash_impl"):
|
||||
... # Register custom kernels with torch dispatcher
|
||||
... pass # doctest: +SKIP
|
||||
>>> register_flash_attention_impl(
|
||||
... "MyImpl", register_fn=my_impl_register
|
||||
... ) # doctest: +SKIP
|
||||
"""
|
||||
global _FLASH_ATTENTION_IMPLS
|
||||
_FLASH_ATTENTION_IMPLS[impl] = register_fn
|
||||
|
||||
|
||||
def activate_flash_attention_impl(
|
||||
impl: str | _FlashAttentionImpl,
|
||||
) -> None:
|
||||
"""
|
||||
Activate into the dispatcher a previously registered flash attention impl.
|
||||
|
||||
.. note::
|
||||
Backend providers should NOT automatically activate their implementation
|
||||
on import. Users should explicitly opt-in by calling this function or via
|
||||
environment variables to ensure multiple provider libraries can coexist.
|
||||
|
||||
Args:
|
||||
impl: Implementation identifier to activate. See
|
||||
:func:`~torch.nn.attention.list_flash_attention_impls` for available
|
||||
implementations.
|
||||
If the backend's :func:`register_flash_attention_impl` callable
|
||||
returns a :class:`FlashAttentionHandle`, the registry keeps that
|
||||
handle alive for the lifetime of the process (until explicit
|
||||
uninstall support exists).
|
||||
|
||||
Example:
|
||||
>>> activate_flash_attention_impl("FA4") # doctest: +SKIP
|
||||
"""
|
||||
global _FLASH_ATTENTION_ACTIVE, _FLASH_ATTENTION_IMPLS
|
||||
|
||||
restore_flash_attention_impl(
|
||||
_raise_warn=False
|
||||
) # first restore any prev overrides (if any) to default
|
||||
|
||||
register_fn = _FLASH_ATTENTION_IMPLS.get(impl)
|
||||
if register_fn is None:
|
||||
raise ValueError(
|
||||
f"Unknown flash attention impl '{impl}'. "
|
||||
f"Available implementations: {list_flash_attention_impls()}"
|
||||
)
|
||||
|
||||
handle = register_fn()
|
||||
if handle is not None:
|
||||
_FLASH_ATTENTION_ACTIVE = (impl, handle)
|
||||
|
||||
|
||||
def list_flash_attention_impls() -> list[str]:
|
||||
"""Return the names of all available flash attention implementations."""
|
||||
return sorted(_FLASH_ATTENTION_IMPLS.keys())
|
||||
|
||||
|
||||
def current_flash_attention_impl() -> str | None:
|
||||
"""
|
||||
Return the currently activated flash attention impl name, if any.
|
||||
|
||||
``None`` indicates that no custom impl has been activated.
|
||||
"""
|
||||
return (
|
||||
_FLASH_ATTENTION_ACTIVE[0]
|
||||
if _FLASH_ATTENTION_ACTIVE is not None
|
||||
else _FLASH_ATTENTION_ACTIVE
|
||||
)
|
||||
|
||||
|
||||
def restore_flash_attention_impl(_raise_warn: bool = True) -> None:
|
||||
"""
|
||||
Restore the default FA2 implementation
|
||||
"""
|
||||
global _FLASH_ATTENTION_ACTIVE
|
||||
|
||||
handle = None
|
||||
if _FLASH_ATTENTION_ACTIVE is not None:
|
||||
handle = _FLASH_ATTENTION_ACTIVE[1]
|
||||
|
||||
if handle is not None:
|
||||
handle.remove()
|
||||
elif _raise_warn:
|
||||
logger.warning(
|
||||
"Trying to restore default FA2 impl when no custom impl was activated"
|
||||
)
|
||||
|
||||
_FLASH_ATTENTION_ACTIVE = None # default
|
||||
@@ -0,0 +1,61 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""Defines utilities for interacting with scaled_dot_product_attention"""
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
__all__: list[str] = []
|
||||
|
||||
|
||||
def _input_requires_grad(*tensors: torch.Tensor) -> bool:
|
||||
"""Returns True if any of the tensors requires grad"""
|
||||
return any(t.requires_grad for t in tensors)
|
||||
|
||||
|
||||
def _postprocess_flash_output(inpt_tensor: torch.Tensor, og_size: int) -> torch.Tensor:
|
||||
"""Handles the unpad of the last dimension"""
|
||||
if inpt_tensor.size(-1) != og_size:
|
||||
return inpt_tensor[..., :og_size]
|
||||
return inpt_tensor
|
||||
|
||||
|
||||
def _calculate_scale(head_dim_size: int, scale: float | None) -> float:
|
||||
"""
|
||||
For FlashAttention we pad the head dimension to be a multiple of 8 so we need to scale the output
|
||||
by the original head size and not the padded.
|
||||
"""
|
||||
if scale is not None:
|
||||
return scale
|
||||
return 1.0 / math.sqrt(head_dim_size)
|
||||
|
||||
|
||||
def _validate_sdpa_input(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_mask: torch.Tensor | None = None,
|
||||
dropout_p=0.0,
|
||||
is_causal=False,
|
||||
scale=None,
|
||||
allow_lowp_kv=False,
|
||||
) -> None:
|
||||
if not allow_lowp_kv:
|
||||
if query.dtype != key.dtype or query.dtype != value.dtype:
|
||||
raise ValueError(
|
||||
f"Expected query, key, and value to have the same dtype, "
|
||||
f"but got query.dtype: {query.dtype}, key.dtype: {key.dtype}, "
|
||||
f"and value.dtype: {value.dtype} instead."
|
||||
)
|
||||
if query.device != key.device or query.device != value.device:
|
||||
raise ValueError(
|
||||
f"Expected query, key, and value to have the same device type, "
|
||||
f"but got query.device: {query.device}, key.device: {key.device}, "
|
||||
f"and value.device: {value.device} instead."
|
||||
)
|
||||
if query.dim() < 2 or key.dim() < 2 or value.dim() < 2:
|
||||
raise ValueError(
|
||||
f"Expected query, key, and value to all be at least 2 dimensional, but got query.dim: "
|
||||
f"{query.dim()}, key.dim: {key.dim()} and value.dim: {value.dim()} instead."
|
||||
)
|
||||
@@ -0,0 +1,376 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""Defines bias subclasses that work with scaled_dot_product_attention"""
|
||||
|
||||
from enum import auto, IntEnum
|
||||
from warnings import warn
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.backends.cuda import (
|
||||
can_use_efficient_attention,
|
||||
can_use_flash_attention,
|
||||
is_flash_attention_available,
|
||||
SDPAParams,
|
||||
)
|
||||
from torch.nn.attention import _raise_kernel_warnings
|
||||
from torch.nn.attention._utils import (
|
||||
_calculate_scale,
|
||||
_input_requires_grad,
|
||||
_postprocess_flash_output,
|
||||
_validate_sdpa_input,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["causal_upper_left", "causal_lower_right", "CausalVariant", "CausalBias"]
|
||||
|
||||
|
||||
torch._dynamo.allow_in_graph(is_flash_attention_available)
|
||||
torch._dynamo.allow_in_graph(can_use_flash_attention)
|
||||
torch._dynamo.allow_in_graph(can_use_efficient_attention)
|
||||
torch._dynamo.allow_in_graph(SDPAParams)
|
||||
|
||||
|
||||
class CausalVariant(IntEnum):
|
||||
r"""
|
||||
Enum for causal variants used in attention mechanisms.
|
||||
|
||||
Defines two types of causal biases:
|
||||
|
||||
``UPPER_LEFT``: Represents upper-left triangular bias for standard causal attention.
|
||||
The equivalent pytorch code for constructing this bias is:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
torch.tril(torch.ones(size, dtype=torch.bool))
|
||||
|
||||
For instance, with ``shape=(3,4)``, the materialized bias tensor will be:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
[[1, 0, 0, 0],
|
||||
[1, 1, 0, 0],
|
||||
[1, 1, 1, 0]]
|
||||
|
||||
|
||||
``LOWER_RIGHT``: Represents lower-right triangular bias, the include values are aligned to the lower
|
||||
right corner of the matrix.
|
||||
|
||||
The equivalent pytorch code for constructing this bias is:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
diagonal_offset = size[1] - size[0]
|
||||
torch.tril(
|
||||
torch.ones(size, dtype=torch.bool),
|
||||
diagonal=diagonal_offset,
|
||||
)
|
||||
|
||||
For instance, with ``shape=(3,4)``, the materialized bias tensor will be:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
[[1, 1, 0, 0],
|
||||
[1, 1, 1, 0],
|
||||
[1, 1, 1, 1]]
|
||||
|
||||
Note that these variants are equivalent to each other when the sequence lengths of the query and key/value
|
||||
tensors are equal since the triangular matrix is square.
|
||||
|
||||
.. warning:: This enum is a prototype and subject to change.
|
||||
"""
|
||||
|
||||
UPPER_LEFT = auto()
|
||||
LOWER_RIGHT = auto()
|
||||
|
||||
|
||||
class CausalBias(torch.Tensor):
|
||||
"""
|
||||
A bias representing causal attention patterns. For an overview of the bias structure, see the :class:`CausalVariant` enum.
|
||||
|
||||
This class is used for defining causal (triangular) attention biases. For construing the bias, there exist
|
||||
two factory functions: :func:`causal_upper_left` and :func:`causal_lower_right`.
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from torch.nn.attention.bias import causal_lower_right
|
||||
|
||||
bsz, num_heads, seqlen_q, seqlen_kv, head_dim = 32, 8, 4, 12, 8
|
||||
|
||||
# Create a lower-right causal bias
|
||||
attn_bias = causal_lower_right(seqlen_q, seqlen_kv)
|
||||
|
||||
q = torch.randn(
|
||||
bsz, num_heads, seqlen_q, head_dim, device="cuda", dtype=torch.float16
|
||||
)
|
||||
k = torch.randn(
|
||||
bsz, num_heads, seqlen_kv, head_dim, device="cuda", dtype=torch.float16
|
||||
)
|
||||
v = torch.randn(
|
||||
bsz, num_heads, seqlen_kv, head_dim, device="cuda", dtype=torch.float16
|
||||
)
|
||||
|
||||
out = F.scaled_dot_product_attention(q, k, v, attn_bias)
|
||||
|
||||
.. warning:: This class is a prototype and subject to change.
|
||||
"""
|
||||
|
||||
def __init__(self, variant: CausalVariant, seq_len_q: int, seq_len_kv: int) -> None:
|
||||
"""
|
||||
Initializes the CausalBias instance with a specified variant and sequence lengths.
|
||||
|
||||
Args:
|
||||
variant (CausalVariant): The type of causal bias to use (either UPPER_LEFT or LOWER_RIGHT).
|
||||
seq_len_q (int): The sequence length of the query tensor.
|
||||
seq_len_kv (int): The sequence length of the key/value tensor.
|
||||
|
||||
Raises a warning if the LOWER_RIGHT variant is used with seq_len_q > seq_len_kv, as it may produce NaNs.
|
||||
"""
|
||||
if not isinstance(variant, CausalVariant):
|
||||
raise AssertionError(
|
||||
f"variant must be a CausalVariant, got {type(variant).__name__}"
|
||||
)
|
||||
super().__init__()
|
||||
self.variant = variant
|
||||
self.seq_len_q = seq_len_q
|
||||
self.seq_len_kv = seq_len_kv
|
||||
if seq_len_q > seq_len_kv and variant == CausalVariant.LOWER_RIGHT:
|
||||
warn(
|
||||
"Lower right causal bias will produce NaNs in the output when seq_len_q > seq_len_kv!",
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
def _upper_left(self, device: torch.device) -> torch.Tensor:
|
||||
"""Upper left causal bias"""
|
||||
return torch.tril(
|
||||
torch.ones(self.seq_len_q, self.seq_len_kv, device=device, dtype=torch.bool)
|
||||
)
|
||||
|
||||
def _lower_right(self, device: torch.device) -> torch.Tensor:
|
||||
"""Lower right causal bias"""
|
||||
diagonal_offset = self.seq_len_kv - self.seq_len_q
|
||||
return torch.tril(
|
||||
torch.ones(
|
||||
self.seq_len_q, self.seq_len_kv, device=device, dtype=torch.bool
|
||||
),
|
||||
diagonal=diagonal_offset,
|
||||
)
|
||||
|
||||
# pyrefly: ignore [bad-return]
|
||||
def _materialize(self, device: torch.device | None = None) -> torch.Tensor:
|
||||
"""
|
||||
Materializes the causal bias into a tensor form.
|
||||
|
||||
Depending on the variant, this method generates either an upper-left or lower-right
|
||||
triangular matrix to represent the causal bias.
|
||||
|
||||
Args:
|
||||
device (Optional[torch.device]): The device on which to create the tensor. Defaults to CPU.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The materialized bias tensor.
|
||||
"""
|
||||
if device is None:
|
||||
device = torch.device("cpu")
|
||||
if self.variant == CausalVariant.UPPER_LEFT:
|
||||
return self._upper_left(device)
|
||||
elif self.variant == CausalVariant.LOWER_RIGHT:
|
||||
return self._lower_right(device)
|
||||
|
||||
@staticmethod
|
||||
def _dispatch(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attn_mask: "CausalBias",
|
||||
dropout_p: float = 0.0,
|
||||
is_causal: bool = False,
|
||||
scale: float | None = None,
|
||||
enable_gqa: bool = False,
|
||||
) -> torch.Tensor:
|
||||
r"""
|
||||
Handles the logic for computing attention with the specified causal bias.
|
||||
|
||||
Args:
|
||||
query (Tensor): Query tensor; shape :math:`(N, ..., L, E)`.
|
||||
key (Tensor): Key tensor; shape :math:`(N, ..., S, E)`.
|
||||
value (Tensor): Value tensor; shape :math:`(N, ..., S, Ev)`.
|
||||
attn_mask (CausalBias): The type of causal attention to apply.
|
||||
A boolean mask where a value of True indicates that the element *should* take part in attention.
|
||||
A float mask of the same type as query, key, value that is added to the attention score.
|
||||
dropout_p (float): Dropout probability; if greater than 0.0, dropout is applied
|
||||
is_causal (bool): If true, assumes upper left causal attention masking and errors if both attn_mask and is_causal
|
||||
are set.
|
||||
scale (optional float): Scaling factor applied prior to softmax. If None, the default value is set
|
||||
to :math:`\frac{1}{\sqrt{E}}`.
|
||||
enable_gqa (optional bool): If set to True, Grouped Query Attention (GQA) is enabled, by default it is set to False.
|
||||
|
||||
Returns:
|
||||
output (Tensor): Attention output; shape :math:`(N, ..., L, Ev)`.
|
||||
|
||||
Raises:
|
||||
ValueError: If the causal bias variant is not a CausalVariant type.
|
||||
|
||||
"""
|
||||
if is_causal:
|
||||
raise ValueError("CausalBias should not be used with causal=True")
|
||||
|
||||
if (
|
||||
attn_mask.seq_len_q == attn_mask.seq_len_kv
|
||||
or attn_mask.variant == CausalVariant.UPPER_LEFT
|
||||
):
|
||||
return F.scaled_dot_product_attention(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
attn_mask=None,
|
||||
dropout_p=dropout_p,
|
||||
is_causal=True,
|
||||
scale=scale,
|
||||
enable_gqa=enable_gqa,
|
||||
)
|
||||
elif attn_mask.variant == CausalVariant.LOWER_RIGHT:
|
||||
_validate_sdpa_input(query, key, value, None, dropout_p, is_causal, scale)
|
||||
sdpa_params = SDPAParams(
|
||||
query, key, value, None, dropout_p, is_causal, enable_gqa
|
||||
)
|
||||
if can_use_flash_attention(sdpa_params):
|
||||
alignment = 64 if query.device.type == "xpu" else 8
|
||||
og_head_size = query.size(-1)
|
||||
og_scale = _calculate_scale(og_head_size, scale)
|
||||
needs_padding = og_head_size % alignment != 0
|
||||
if needs_padding:
|
||||
pad_len = alignment - (og_head_size % alignment)
|
||||
query = torch.nn.functional.pad(query, (0, pad_len))
|
||||
key = torch.nn.functional.pad(key, (0, pad_len))
|
||||
value = torch.nn.functional.pad(value, (0, pad_len))
|
||||
out = torch.ops.aten._scaled_dot_product_flash_attention(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
dropout_p,
|
||||
is_causal=True, # TODO: Flash accepts causal = True and for this particular op it means lower right
|
||||
return_debug_mask=False,
|
||||
scale=og_scale,
|
||||
)[0]
|
||||
return _postprocess_flash_output(out, og_head_size)
|
||||
if can_use_efficient_attention(sdpa_params):
|
||||
compute_log_sumexp = False
|
||||
if _input_requires_grad(query, key, value):
|
||||
compute_log_sumexp = True
|
||||
return torch.ops.aten._efficient_attention_forward(
|
||||
query.transpose(1, 2),
|
||||
key.transpose(1, 2),
|
||||
value.transpose(1, 2),
|
||||
bias=None,
|
||||
cu_seqlens_q=None,
|
||||
cu_seqlens_k=None,
|
||||
max_seqlen_q=None,
|
||||
max_seqlen_k=None,
|
||||
dropout_p=dropout_p,
|
||||
custom_mask_type=int(attn_mask.variant),
|
||||
compute_log_sumexp=compute_log_sumexp,
|
||||
scale=scale,
|
||||
seqlen_k=None,
|
||||
)[0].transpose(1, 2)
|
||||
else:
|
||||
_raise_kernel_warnings(sdpa_params)
|
||||
# We can't use efficient attention the only support for lower right is via materialization
|
||||
return F.scaled_dot_product_attention(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
attn_mask=attn_mask._materialize(query.device),
|
||||
dropout_p=dropout_p,
|
||||
is_causal=False,
|
||||
scale=scale,
|
||||
enable_gqa=enable_gqa,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"CausalBias.variant must be a CausalVariant type, but found: {attn_mask.variant}"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def __torch_function__(cls, func, types, args=(), kwargs=None):
|
||||
"""Defines the behavior of torch.nn.functional.scaled_dot_product_attention when the attn_bias is an AttnBias"""
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
if func is torch.nn.functional.scaled_dot_product_attention:
|
||||
return cls._dispatch(*args, **kwargs)
|
||||
return super().__torch_function__(func, types, args, kwargs)
|
||||
|
||||
def __repr__(self) -> str: # type:ignore[override]
|
||||
return self._materialize().__repr__()
|
||||
|
||||
|
||||
def causal_upper_left(*size) -> CausalBias:
|
||||
"""
|
||||
Creates an upper-left triangular causal bias.
|
||||
|
||||
This function generates a upper-left triangular matrix to represent causal attention bias with a
|
||||
diagonal offset set so that the inclusive values are aligned to the upper left corner of the matrix.
|
||||
This equivalent to the `is_causal=True` argument in `scaled_dot_product_attention`.
|
||||
|
||||
The equivalent pytorch code for constructing this bias is:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
torch.tril(torch.ones(size, dtype=torch.bool))
|
||||
|
||||
For instance, with `shape=(3,4)`, the materialized bias tensor will be:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
[[1, 0, 0, 0],
|
||||
[1, 1, 0, 0],
|
||||
[1, 1, 1, 0]]
|
||||
|
||||
Args:
|
||||
size: The size of the bias matrix.
|
||||
|
||||
Returns:
|
||||
CausalBias: The UPPER_LEFT triangular causal bias variant.
|
||||
"""
|
||||
if len(size) != 2:
|
||||
raise AssertionError("causal_upper_left only supports 2D tensors")
|
||||
seq_len_q, seq_len_kv = size
|
||||
return CausalBias(CausalVariant.UPPER_LEFT, seq_len_q, seq_len_kv)
|
||||
|
||||
|
||||
def causal_lower_right(*size) -> CausalBias:
|
||||
"""
|
||||
Creates a lower-right triangular causal bias.
|
||||
|
||||
This function generates a lower-right triangular matrix to represent causal attention bias with a
|
||||
diagonal offset set so that the inclusive values are aligned to the lower right corner of the matrix.
|
||||
|
||||
The equivalent pytorch code for constructing this bias is:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
diagonal_offset = size[1] - size[0]
|
||||
torch.tril(
|
||||
torch.ones(size, dtype=torch.bool),
|
||||
diagonal=diagonal_offset,
|
||||
)
|
||||
|
||||
For instance, with `shape=(3,4)`, the materialized bias tensor will be:
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
[[1, 1, 0, 0],
|
||||
[1, 1, 1, 0],
|
||||
[1, 1, 1, 1]]
|
||||
|
||||
Args:
|
||||
size: The size of the bias matrix.
|
||||
|
||||
Returns:
|
||||
CausalBias: The LOWER_RIGHT triangular causal bias variant.
|
||||
"""
|
||||
if len(size) != 2:
|
||||
raise AssertionError("causal_lower_right only supports 2D tensors")
|
||||
seq_len_q, seq_len_kv = size
|
||||
return CausalBias(CausalVariant.LOWER_RIGHT, seq_len_q, seq_len_kv)
|
||||
@@ -0,0 +1,2 @@
|
||||
# Experimental features are not mature yet and are subject to change.
|
||||
# We do not provide any BC/FC guarantees
|
||||
+358
@@ -0,0 +1,358 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""
|
||||
This module implements Paged Attention on top of flex_attention.
|
||||
This module is experimental and subject to change.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch.nn.attention.flex_attention import (
|
||||
_identity,
|
||||
_mask_mod_signature,
|
||||
_score_mod_signature,
|
||||
BlockMask,
|
||||
noop_mask,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["PagedAttention"]
|
||||
|
||||
|
||||
def _cdiv(x: int | float | torch.Tensor, multiple: int | float | torch.Tensor):
|
||||
return (x + multiple - 1) // multiple
|
||||
|
||||
|
||||
class PagedAttention:
|
||||
"""
|
||||
PagedAttention supports flex attention inference with a large batch size.
|
||||
With PagedAttention, a batch of key/value tensors with varying kv length
|
||||
is split into tensor blocks of fixed length and cached in a compact way.
|
||||
Thus we can avoid redundant memory consumption due to varying kv length and
|
||||
support a larger batch size.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
n_pages: int,
|
||||
page_size: int,
|
||||
max_batch_size: int,
|
||||
device: str = "cuda",
|
||||
) -> None:
|
||||
# number of pages
|
||||
self.n_pages = n_pages
|
||||
|
||||
# number of tokens per page
|
||||
self.page_size = page_size
|
||||
|
||||
# page table: [batch, logical_block_idx] -> physical_page_idx
|
||||
self.page_table = -torch.ones(
|
||||
(max_batch_size, self.n_pages), dtype=torch.int64, device=device
|
||||
)
|
||||
|
||||
# capacity: batch_idx -> allocated sequence length
|
||||
self.capacity = torch.zeros(max_batch_size, dtype=torch.int64, device=device)
|
||||
|
||||
# index of empty pages that is available for allocation
|
||||
self.empty_pages = list(range(n_pages - 1, -1, -1))
|
||||
|
||||
# mapping from physical page index to logical page index
|
||||
self.physical_to_logical = -torch.ones(
|
||||
(max_batch_size, n_pages), dtype=torch.int64, device=device
|
||||
)
|
||||
|
||||
def reserve(self, batch_idx: torch.Tensor, seq_len: torch.Tensor) -> None:
|
||||
"""
|
||||
Requests the capacity of a given batch to be at least enough to
|
||||
hold `seq_len` elements.
|
||||
|
||||
Args:
|
||||
batch_idx (Tensor): batch index to be reserved; shape :math:`(1)`.
|
||||
seq_len (Tensor): minimum capacity for the given batch; shape :math:`(1)`.
|
||||
"""
|
||||
|
||||
if seq_len <= self.capacity[batch_idx]:
|
||||
return
|
||||
|
||||
num_pages_to_allocate = _cdiv(
|
||||
seq_len - self.capacity[batch_idx], self.page_size
|
||||
)
|
||||
|
||||
if len(self.empty_pages) < num_pages_to_allocate:
|
||||
raise AssertionError(
|
||||
f"requested {num_pages_to_allocate.item()} pages "
|
||||
f"but there are only {len(self.empty_pages)} empty pages"
|
||||
)
|
||||
|
||||
start_page_idx = self.capacity[batch_idx] // self.page_size
|
||||
end_page_idx = start_page_idx + num_pages_to_allocate
|
||||
|
||||
# find empty physical pages
|
||||
allocated_pages = torch.tensor(
|
||||
self.empty_pages[-num_pages_to_allocate:],
|
||||
device=num_pages_to_allocate.device,
|
||||
)
|
||||
self.empty_pages = self.empty_pages[:-num_pages_to_allocate]
|
||||
|
||||
# update page table
|
||||
self.page_table[
|
||||
batch_idx,
|
||||
start_page_idx:end_page_idx,
|
||||
] = allocated_pages
|
||||
|
||||
# update metadata
|
||||
self.physical_to_logical[batch_idx, allocated_pages] = torch.arange(
|
||||
start_page_idx.item(),
|
||||
end_page_idx.item(),
|
||||
device=num_pages_to_allocate.device,
|
||||
)
|
||||
self.capacity[batch_idx] += num_pages_to_allocate * self.page_size
|
||||
|
||||
def erase(self, batch_idx: torch.Tensor) -> None:
|
||||
"""
|
||||
Removes a single batch from paged attention.
|
||||
|
||||
Args:
|
||||
batch_idx (Tensor): batch index to be removed; shape :math:`(1)`.
|
||||
"""
|
||||
|
||||
# find allocated pages
|
||||
allocated_page_idx = self.page_table[batch_idx] != -1
|
||||
allocated_pages = self.page_table[batch_idx][allocated_page_idx]
|
||||
|
||||
# clean metadata
|
||||
self.capacity[batch_idx] = 0
|
||||
self.empty_pages += allocated_pages.tolist()
|
||||
self.physical_to_logical[batch_idx][:, allocated_pages] = -1
|
||||
self.page_table[batch_idx] = -1
|
||||
|
||||
def assign(
|
||||
self,
|
||||
batch_idx: torch.Tensor,
|
||||
input_pos: torch.Tensor,
|
||||
k_val: torch.Tensor,
|
||||
v_val: torch.Tensor,
|
||||
k_cache: torch.Tensor,
|
||||
v_cache: torch.Tensor,
|
||||
) -> None:
|
||||
"""
|
||||
Assigns new contents `val` to the storage `cache` at the location
|
||||
`batch_idx` and `input_pos`.
|
||||
|
||||
Args:
|
||||
batch_idx (Tensor): batch index; shape :math:`(B)`.
|
||||
input_pos (Tensor): input positions to be assigned for the given batch; shape :math:`(B, S)`.
|
||||
val (Tensor): value to be assigned; shape :math:`(B, H, S, D)`
|
||||
cache (Tensor): the cache to store the values; shape:`(1, H, MAX_S, D)`
|
||||
"""
|
||||
if k_val.requires_grad:
|
||||
raise RuntimeError("val must not require gradient")
|
||||
|
||||
B, H, S, K_D = k_val.shape
|
||||
V_D = v_val.shape[3]
|
||||
if B != batch_idx.shape[0]:
|
||||
raise RuntimeError(
|
||||
f"Expect val and batch_idx have the same batch size "
|
||||
f"but got B={B} and B={batch_idx.shape[0]}."
|
||||
)
|
||||
if H != k_cache.shape[1]:
|
||||
raise RuntimeError(
|
||||
f"Expect val and cache has the same number of heads "
|
||||
f"but got H={H} and H={k_cache.shape[1]}."
|
||||
)
|
||||
if S != input_pos.shape[1]:
|
||||
raise RuntimeError(
|
||||
f"Expect val and input_pos has the same length "
|
||||
f"but got S={S} and S={input_pos.shape[0]}."
|
||||
)
|
||||
if K_D != k_cache.shape[3]:
|
||||
raise RuntimeError(
|
||||
f"Expect k_val and k_cache has the same hidden dim "
|
||||
f"but got D={K_D} and D={k_cache.shape[3]}."
|
||||
)
|
||||
if V_D != v_cache.shape[3]:
|
||||
raise RuntimeError(
|
||||
f"Expect v_val and v_cache has the same hidden dim "
|
||||
f"but got D={V_D} and D={v_cache.shape[3]}."
|
||||
)
|
||||
|
||||
# find address
|
||||
logical_block_idx = input_pos // self.page_size # [B, S]
|
||||
logical_block_offset = input_pos % self.page_size # [B, S]
|
||||
physical_block_idx = torch.gather(
|
||||
self.page_table[batch_idx], 1, logical_block_idx.to(torch.int64)
|
||||
).to(torch.int32) # [B, S]
|
||||
|
||||
addr = (physical_block_idx * self.page_size + logical_block_offset).view(
|
||||
-1
|
||||
) # [B*S]
|
||||
|
||||
k_val = k_val.permute(1, 0, 2, 3).contiguous().view(1, H, B * S, K_D)
|
||||
v_val = v_val.permute(1, 0, 2, 3).contiguous().view(1, H, B * S, V_D)
|
||||
|
||||
k_cache[:, :, addr, :] = k_val
|
||||
v_cache[:, :, addr, :] = v_val
|
||||
|
||||
def convert_logical_block_mask(
|
||||
self,
|
||||
block_mask: BlockMask,
|
||||
batch_idx: torch.Tensor | None = None,
|
||||
kv_len: torch.Tensor | None = None,
|
||||
) -> BlockMask:
|
||||
"""
|
||||
Converts a logical block mask by mapping its logical kv indices to the corresponding
|
||||
physical kv indices.
|
||||
|
||||
Args:
|
||||
block_mask (BlockMask): logical block mask;
|
||||
kv_indices shape :math:`(B, H, ROWS, MAX_BLOCKS_IN_COL)`.
|
||||
batch_idx (Tensor): batch index corresponding to the block_mask
|
||||
batch dimension. This provides flexibility to convert a
|
||||
block mask with smaller batch size than the page table;
|
||||
shape :math:`(B)`.
|
||||
kv_len (Optional[Tensor]): actual KV sequence length for upper bound check;
|
||||
shape :math:`(B,)` to handle multiple batches.
|
||||
"""
|
||||
B, H, ROWS, MAX_BLOCKS_IN_COL = block_mask.kv_indices.shape
|
||||
|
||||
if block_mask.BLOCK_SIZE[1] != self.page_size:
|
||||
raise RuntimeError(
|
||||
f"Expect block_mask has the same column block size as page_size"
|
||||
f"but got size={block_mask.BLOCK_SIZE[1]} and size={self.page_size}"
|
||||
)
|
||||
|
||||
# Increase the num columns of converted block mask from logical block mask's
|
||||
# num columns to n_pages, since a) the converted block mask
|
||||
# may have larger indices values; and b) `_ordered_to_dense` realizes
|
||||
# a dense tensor with these converted indices. There would be an IndexError
|
||||
# if using the logical block mask's num columns.
|
||||
|
||||
device = block_mask.kv_num_blocks.device
|
||||
|
||||
if batch_idx is None:
|
||||
batch_idx = torch.arange(B, device=device)
|
||||
page_table = self.page_table[batch_idx]
|
||||
|
||||
new_kv_num_blocks = block_mask.kv_num_blocks.clone()
|
||||
|
||||
new_kv_indices = torch.zeros(
|
||||
(B, H, ROWS, self.n_pages), dtype=torch.int32, device=device
|
||||
)
|
||||
new_kv_indices[:, :, :, :MAX_BLOCKS_IN_COL] = (
|
||||
torch.gather(
|
||||
page_table, 1, block_mask.kv_indices.view(B, -1).to(torch.int64)
|
||||
)
|
||||
.view(block_mask.kv_indices.shape)
|
||||
.to(torch.int32)
|
||||
)
|
||||
|
||||
new_full_kv_indices, new_full_kv_num_blocks = None, None
|
||||
if block_mask.full_kv_num_blocks is not None:
|
||||
if block_mask.full_kv_indices is None:
|
||||
raise AssertionError(
|
||||
"block_mask.full_kv_indices must not be None when full_kv_num_blocks is not None"
|
||||
)
|
||||
new_full_kv_num_blocks = block_mask.full_kv_num_blocks.clone()
|
||||
new_full_kv_indices = torch.zeros(
|
||||
(B, H, ROWS, self.n_pages), dtype=torch.int32, device=device
|
||||
)
|
||||
new_full_kv_indices[:, :, :, :MAX_BLOCKS_IN_COL] = (
|
||||
torch.gather(
|
||||
page_table,
|
||||
1,
|
||||
block_mask.full_kv_indices.view(B, -1).to(torch.int64),
|
||||
)
|
||||
.view(block_mask.full_kv_indices.shape)
|
||||
.to(torch.int32)
|
||||
)
|
||||
|
||||
new_mask_mod = self.get_mask_mod(block_mask.mask_mod, kv_len)
|
||||
|
||||
seq_lengths = (block_mask.seq_lengths[0], self.n_pages * self.page_size)
|
||||
return BlockMask.from_kv_blocks(
|
||||
new_kv_num_blocks,
|
||||
new_kv_indices,
|
||||
new_full_kv_num_blocks,
|
||||
new_full_kv_indices,
|
||||
block_mask.BLOCK_SIZE,
|
||||
new_mask_mod,
|
||||
seq_lengths=seq_lengths,
|
||||
)
|
||||
|
||||
def get_mask_mod(
|
||||
self,
|
||||
mask_mod: _mask_mod_signature | None,
|
||||
kv_len: torch.Tensor | None = None,
|
||||
) -> _mask_mod_signature:
|
||||
"""
|
||||
Converts a mask_mod based on mapping from the physical block index to the logical
|
||||
block index.
|
||||
|
||||
Args:
|
||||
mask_mod (_mask_mod_signature): mask_mod based on the logical block index.
|
||||
kv_len (Optional[torch.Tensor]): actual KV sequence length for upper bound check.
|
||||
"""
|
||||
if mask_mod is None:
|
||||
mask_mod = noop_mask
|
||||
|
||||
def new_mask_mod(
|
||||
b: torch.Tensor,
|
||||
h: torch.Tensor,
|
||||
q_idx: torch.Tensor,
|
||||
physical_kv_idx: torch.Tensor,
|
||||
):
|
||||
physical_kv_block = physical_kv_idx // self.page_size
|
||||
physical_kv_offset = physical_kv_idx % self.page_size
|
||||
logical_block_idx = self.physical_to_logical[b, physical_kv_block]
|
||||
logical_kv_idx = logical_block_idx * self.page_size + physical_kv_offset
|
||||
live_block = logical_block_idx >= 0
|
||||
within_upper_bound = (
|
||||
logical_kv_idx < kv_len[b] if kv_len is not None else True
|
||||
)
|
||||
within_lower_bound = logical_kv_idx >= 0
|
||||
is_valid = live_block & within_upper_bound & within_lower_bound
|
||||
|
||||
return torch.where(is_valid, mask_mod(b, h, q_idx, logical_kv_idx), False)
|
||||
|
||||
return new_mask_mod
|
||||
|
||||
def get_score_mod(
|
||||
self,
|
||||
score_mod: _score_mod_signature | None,
|
||||
kv_len: torch.Tensor | None = None,
|
||||
) -> _score_mod_signature:
|
||||
"""
|
||||
Converts a score_mod based on mapping from the physical block index to the logical
|
||||
block index.
|
||||
|
||||
Args:
|
||||
score_mod (_score_mod_signature): score_mod based on the logical block index.
|
||||
`kv_len (Optional[torch.Tensor]): actual KV sequence length for upper bound check.
|
||||
|
||||
"""
|
||||
if score_mod is None:
|
||||
score_mod = _identity
|
||||
|
||||
def new_score_mod(
|
||||
score: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
h: torch.Tensor,
|
||||
q_idx: torch.Tensor,
|
||||
physical_kv_idx: torch.Tensor,
|
||||
):
|
||||
physical_kv_block = physical_kv_idx // self.page_size
|
||||
physical_kv_offset = physical_kv_idx % self.page_size
|
||||
logical_block_idx = self.physical_to_logical[b, physical_kv_block]
|
||||
logical_kv_idx = logical_block_idx * self.page_size + physical_kv_offset
|
||||
live_block = logical_block_idx >= 0
|
||||
within_upper_bound = (
|
||||
logical_kv_idx < kv_len[b] if kv_len is not None else True
|
||||
)
|
||||
within_lower_bound = logical_kv_idx >= 0
|
||||
is_valid = live_block & within_upper_bound & within_lower_bound
|
||||
|
||||
return torch.where(
|
||||
is_valid,
|
||||
score_mod(score, b, h, q_idx, logical_kv_idx),
|
||||
float("-inf"),
|
||||
)
|
||||
|
||||
return new_score_mod
|
||||
+154
@@ -0,0 +1,154 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""
|
||||
This operator implements FP8 scaled dot product attention using Flash Attention 3.
|
||||
This operator is experimental and subject to change.
|
||||
"""
|
||||
|
||||
import warnings
|
||||
from enum import IntEnum
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
class DescaleType(IntEnum):
|
||||
"""Describes the scaling granularity for FP8 descale tensors.
|
||||
|
||||
Used with _scaled_dot_product_attention_quantized to explicitly specify
|
||||
how the descale factors are applied to the quantized inputs.
|
||||
|
||||
.. warning::
|
||||
This enum is experimental and subject to change.
|
||||
"""
|
||||
|
||||
PER_HEAD = 0
|
||||
"""Per-head descaling. Descale tensor shape: (batch_size, num_kv_heads)."""
|
||||
|
||||
|
||||
def _validate_descale(
|
||||
descale: Tensor | None,
|
||||
name: str,
|
||||
query: Tensor,
|
||||
key: Tensor,
|
||||
descale_type: DescaleType,
|
||||
) -> None:
|
||||
"""Validate descale tensor for the specified scaling type.
|
||||
|
||||
Args:
|
||||
descale: The descale tensor to validate (may be None)
|
||||
name: Name of the descale tensor ("q", "k", or "v") for error messages
|
||||
query: Query tensor to get batch size
|
||||
key: Key tensor to get num_kv_heads
|
||||
descale_type: The scaling granularity being used
|
||||
|
||||
Raises:
|
||||
ValueError: If the descale tensor has invalid dtype, device, or shape
|
||||
|
||||
Note:
|
||||
All descale tensors (q, k, v) use num_kv_heads for the head dimension.
|
||||
For GQA/MQA where num_query_heads > num_kv_heads, q_descale is broadcast
|
||||
from (B, H_kv) to match the query heads internally.
|
||||
"""
|
||||
if descale is None:
|
||||
return
|
||||
|
||||
# Check dtype
|
||||
if descale.dtype != torch.float32:
|
||||
raise ValueError(f"{name}_descale must have dtype float32, got {descale.dtype}")
|
||||
|
||||
# Check device
|
||||
if not descale.is_cuda:
|
||||
raise ValueError(f"{name}_descale must be a CUDA tensor")
|
||||
|
||||
# Check shape based on descale type
|
||||
if descale_type == DescaleType.PER_HEAD:
|
||||
batch_size = query.size(0)
|
||||
# All descale tensors use num_kv_heads, even q_descale (broadcast internally)
|
||||
# For BHSD layout, num_kv_heads is at dim 1 of key
|
||||
num_kv_heads = key.size(1)
|
||||
|
||||
if descale.dim() != 2:
|
||||
raise ValueError(
|
||||
f"{name}_descale must be a 2D tensor with shape (batch_size, num_kv_heads) "
|
||||
f"for PER_HEAD descaling, got {descale.dim()}D tensor"
|
||||
)
|
||||
|
||||
if descale.size(0) != batch_size:
|
||||
raise ValueError(
|
||||
f"{name}_descale batch dimension must match query batch size, "
|
||||
f"expected {batch_size}, got {descale.size(0)}"
|
||||
)
|
||||
|
||||
if descale.size(1) != num_kv_heads:
|
||||
raise ValueError(
|
||||
f"{name}_descale head dimension must match num_kv_heads, "
|
||||
f"expected {num_kv_heads}, got {descale.size(1)}"
|
||||
)
|
||||
|
||||
|
||||
def _scaled_dot_product_attention_quantized(
|
||||
query: Tensor,
|
||||
key: Tensor,
|
||||
value: Tensor,
|
||||
is_causal: bool = False,
|
||||
scale: float | None = None,
|
||||
q_descale: Tensor | None = None,
|
||||
k_descale: Tensor | None = None,
|
||||
v_descale: Tensor | None = None,
|
||||
q_descale_type: DescaleType = DescaleType.PER_HEAD,
|
||||
k_descale_type: DescaleType = DescaleType.PER_HEAD,
|
||||
v_descale_type: DescaleType = DescaleType.PER_HEAD,
|
||||
) -> Tensor:
|
||||
r"""Scaled dot product attention for FP8 inputs.
|
||||
|
||||
This is a specialized version of scaled_dot_product_attention that supports
|
||||
FP8 quantized inputs (float8_e4m3fn) with per-head descaling. Requires the
|
||||
Flash Attention 3 backend to be activated.
|
||||
|
||||
.. warning::
|
||||
This function is experimental and only supports forward pass.
|
||||
|
||||
Args:
|
||||
query (Tensor): Query tensor; shape :math:`(N, H_q, L, E)` dtype float8_e4m3fn
|
||||
key (Tensor): Key tensor; shape :math:`(N, H, S, E)` dtype float8_e4m3fn
|
||||
value (Tensor): Value tensor; shape :math:`(N, H, S, E_v)` dtype float8_e4m3fn
|
||||
is_causal (bool): Apply causal attention mask
|
||||
scale (float, optional): Scaling factor for attention weights
|
||||
q_descale (Tensor, optional): Query descale tensor; shape :math:`(N, H)` for PER_HEAD
|
||||
k_descale (Tensor, optional): Key descale tensor; shape :math:`(N, H)` for PER_HEAD
|
||||
v_descale (Tensor, optional): Value descale tensor; shape :math:`(N, H)` for PER_HEAD
|
||||
q_descale_type (DescaleType): Specifies the descaling granularity for query. Default: PER_HEAD
|
||||
k_descale_type (DescaleType): Specifies the descaling granularity for key. Default: PER_HEAD
|
||||
v_descale_type (DescaleType): Specifies the descaling granularity for value. Default: PER_HEAD
|
||||
|
||||
Returns:
|
||||
Tensor: Attention output; shape :math:`(N, H_q, L, E_v)` dtype bfloat16
|
||||
"""
|
||||
# Validate descale tensors
|
||||
_validate_descale(q_descale, "q", query, key, q_descale_type)
|
||||
_validate_descale(k_descale, "k", query, key, k_descale_type)
|
||||
_validate_descale(v_descale, "v", query, key, v_descale_type)
|
||||
|
||||
if torch.is_grad_enabled() and (
|
||||
query.requires_grad or key.requires_grad or value.requires_grad
|
||||
):
|
||||
warnings.warn(
|
||||
"_scaled_dot_product_attention_quantized does not support backward pass. "
|
||||
"Gradients will not be computed for query, key, or value.",
|
||||
UserWarning,
|
||||
)
|
||||
# Directly call the internal flash attention operator which has descale support
|
||||
# NOTE: This should be torch._scaled_dot_product_flash_attention, but it does not work with torch.compile
|
||||
result = torch.ops.aten._scaled_dot_product_flash_attention.quantized(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
q_descale,
|
||||
k_descale,
|
||||
v_descale,
|
||||
0.0,
|
||||
is_causal,
|
||||
False,
|
||||
scale=scale,
|
||||
)
|
||||
return result[0] # Return the output tensor, mirroring scaled_dot_product_attention
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,671 @@
|
||||
"""
|
||||
Variable-length attention implementation using Flash Attention.
|
||||
|
||||
This module provides a high-level Python interface for variable-length attention
|
||||
that calls into the optimized Flash Attention kernels.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from functools import lru_cache
|
||||
from typing import Any, NamedTuple
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
__all__ = ["varlen_attn", "varlen_attn_out", "AuxRequest"]
|
||||
|
||||
|
||||
def _normalize_window_size(window_size: list[int] | None) -> list[int]:
|
||||
if window_size is None:
|
||||
window_size = [-1, -1]
|
||||
|
||||
if len(window_size) != 2:
|
||||
raise ValueError(f"window_size must have length 2, got {len(window_size)}")
|
||||
return window_size
|
||||
|
||||
|
||||
@lru_cache(maxsize=8)
|
||||
def _should_use_cudnn(device_index: int) -> bool:
|
||||
"""Cache device capability check to avoid repeated CUDA calls."""
|
||||
return False
|
||||
|
||||
|
||||
class AuxRequest(NamedTuple):
|
||||
"""
|
||||
Request which auxiliary outputs to compute from varlen_attn.
|
||||
|
||||
Each field is a boolean indicating whether that auxiliary output should be computed.
|
||||
"""
|
||||
|
||||
lse: bool = False
|
||||
|
||||
|
||||
@torch.library.custom_op("torch_attn::_varlen_attn", mutates_args={})
|
||||
def _varlen_attn(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor,
|
||||
cu_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
is_causal: bool = False,
|
||||
scale: float | None = None,
|
||||
window_size: list[int] | None = None,
|
||||
enable_gqa: bool = False,
|
||||
seqused_k: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Private custom op for variable-length attention.
|
||||
|
||||
This is the internal implementation. Users should use the public varlen_attn function instead.
|
||||
"""
|
||||
window_size = _normalize_window_size(window_size)
|
||||
|
||||
use_cudnn = query.is_cuda and _should_use_cudnn(query.device.index)
|
||||
|
||||
if use_cudnn:
|
||||
log.info("Using cuDNN backend for varlen_attn")
|
||||
|
||||
if enable_gqa:
|
||||
# TODO: check this
|
||||
raise RuntimeError("GQA is not supported with the cuDNN backend.")
|
||||
if num_splits is not None:
|
||||
# TODO: check this
|
||||
raise RuntimeError("num_splits is not supported with the cuDNN backend.")
|
||||
if window_size[0] != -1 or window_size[1] != -1:
|
||||
raise RuntimeError(
|
||||
"cuDNN backend does not support window attention. Please use Flash Attention backend."
|
||||
)
|
||||
if seqused_k is not None or block_table is not None:
|
||||
# TODO: cuDNN supports per-sequence KV lengths via SEQ_LEN_KV + padding_mask,
|
||||
# but _cudnn_attention_forward doesn't expose it yet.
|
||||
raise RuntimeError(
|
||||
"seqused_k/block_table is not yet supported with the cuDNN backend."
|
||||
)
|
||||
|
||||
result = torch.ops.aten._cudnn_attention_forward(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
None, # attn_bias
|
||||
cu_seq_q,
|
||||
cu_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
True, # compute_log_sumexp
|
||||
0.0, # dropout_p hardcoded to 0.0
|
||||
is_causal,
|
||||
False, # return_debug_mask
|
||||
scale=scale,
|
||||
)
|
||||
# cuDNN returns: (output, logsumexp, cum_seq_q, cum_seq_k, max_q, max_k, philox_seed, philox_offset, debug_attn_mask)
|
||||
output, softmax_lse, rng_state = result[0], result[1], result[6]
|
||||
else:
|
||||
log.info("Using Flash Attention backend for varlen_attn")
|
||||
output, softmax_lse, rng_state, _, _ = torch.ops.aten._flash_attention_forward(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cu_seq_q,
|
||||
cu_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
0.0, # dropout_p hardcoded to 0.0
|
||||
is_causal,
|
||||
return_debug_mask=False,
|
||||
scale=scale,
|
||||
window_size_left=window_size[0],
|
||||
window_size_right=window_size[1],
|
||||
seqused_k=seqused_k,
|
||||
block_table=block_table,
|
||||
num_splits=num_splits,
|
||||
)
|
||||
|
||||
rng_state_ = torch.zeros(
|
||||
(2,), dtype=torch.uint64, device=query.device
|
||||
) # hardcoded since dropout is hardcoded to 0
|
||||
return output, softmax_lse, rng_state_
|
||||
|
||||
|
||||
@_varlen_attn.register_fake
|
||||
def _varlen_attn_fake(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor,
|
||||
cu_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
is_causal: bool = False,
|
||||
scale: float | None = None,
|
||||
window_size: list[int] | None = None,
|
||||
enable_gqa: bool = False,
|
||||
seqused_k: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Fake implementation for meta tensor computation and tracing.
|
||||
|
||||
Based on the 3D varlen path from meta__flash_attention_forward:
|
||||
- query shape: (total, num_heads, head_dim)
|
||||
- logsumexp shape: (num_heads, total_q)
|
||||
"""
|
||||
window_size = _normalize_window_size(window_size)
|
||||
|
||||
# Output has same shape as query
|
||||
output = torch.empty_like(query)
|
||||
|
||||
# For varlen path: logsumexp shape is (num_heads, total_q)
|
||||
total_q = query.size(0)
|
||||
num_heads = query.size(1)
|
||||
logsumexp = torch.empty(
|
||||
(num_heads, total_q), dtype=torch.float, device=query.device
|
||||
)
|
||||
|
||||
if torch.version.hip:
|
||||
preferred = torch._C._get_rocm_fa_preferred_backend()
|
||||
if preferred == torch._C._ROCmFABackend.AOTriton:
|
||||
# AOTriton ROCm path uses batched 3D
|
||||
batch_size = cu_seq_q.size(0) - 1
|
||||
logsumexp = torch.empty(
|
||||
(batch_size, num_heads, max_q), dtype=torch.float, device=query.device
|
||||
)
|
||||
|
||||
rng_state = torch.empty((2,), dtype=torch.uint64, device=query.device)
|
||||
|
||||
return output, logsumexp, rng_state
|
||||
|
||||
|
||||
def varlen_attn(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor,
|
||||
cu_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
*,
|
||||
return_aux: AuxRequest | None = None,
|
||||
scale: float | None = None,
|
||||
window_size: tuple[int, int] = (-1, -1),
|
||||
enable_gqa: bool = False,
|
||||
seqused_k: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
|
||||
r"""Compute variable-length attention using Flash Attention.
|
||||
|
||||
This function is similar to scaled_dot_product_attention but optimized for
|
||||
variable-length sequences using cumulative sequence position tensors.
|
||||
|
||||
Args:
|
||||
query (Tensor): Query tensor; shape :math:`(T_q, H_q, D)`
|
||||
key (Tensor): Key tensor; shape :math:`(T_k, H_{kv}, D)`, or
|
||||
:math:`(\text{total\_pages}, \text{page\_size}, H_{kv}, D)` when ``block_table`` is provided.
|
||||
value (Tensor): Value tensor; shape :math:`(T_k, H_{kv}, D)`, or
|
||||
:math:`(\text{total\_pages}, \text{page\_size}, H_{kv}, D)` when ``block_table`` is provided.
|
||||
cu_seq_q (Tensor): Cumulative sequence positions for queries; shape :math:`(N+1,)`
|
||||
cu_seq_k (Tensor): Cumulative sequence positions for keys/values; shape :math:`(N+1,)`
|
||||
max_q (int): Maximum query sequence length in the batch.
|
||||
max_k (int): Maximum key/value sequence length in the batch.
|
||||
return_aux (Optional[AuxRequest]): If not None and ``return_aux.lse`` is True, also returns the logsumexp tensor.
|
||||
scale (float, optional): Scaling factor for attention scores
|
||||
window_size (tuple[int, int], optional): Window size for sliding window attention as (left, right).
|
||||
Use (-1, -1) for full attention (default), (-1, 0) for causal attention,
|
||||
or (W, 0) for causal attention with sliding window of size W.
|
||||
enable_gqa (bool): If set to True, enables Grouped Query Attention (GQA)
|
||||
and allows key/value to have fewer heads than query.
|
||||
Each KV head is shared by a group of :math:`H_q / H_{kv}` query heads,
|
||||
so :math:`H_q` must be divisible by :math:`H_{kv}`.
|
||||
Default is False.
|
||||
seqused_k (Tensor, optional): Number of valid KV tokens per batch element; shape :math:`(N,)`.
|
||||
When set, only the first ``seqused_k[i]`` tokens in the key/value sequence for batch
|
||||
element *i* participate in attention. Useful for KV-cache decoding where the cache slot
|
||||
is larger than the actual sequence. Inference-only (not supported in backward).
|
||||
block_table (Tensor, optional): Block table for paged KV cache; shape
|
||||
:math:`(N, \text{max\_pages\_per\_seq})`, dtype ``int32``.
|
||||
Requires ``seqused_k``. Inference-only (not supported in backward).
|
||||
|
||||
When ``block_table`` is provided, ``key`` and ``value`` are a "pool" of
|
||||
pages of tokens of KV data and the pages belong to any sequence/order.
|
||||
The ``block_table`` is what maps each sequence's logical chunks
|
||||
back to physical pages in this pool.
|
||||
|
||||
``seqused_k[i]`` tells the kernel how many tokens in sequence *i* are
|
||||
actually valid, since the last page is typically only partially filled.
|
||||
num_splits (int, optional): Number of splits for split-KV. Set to ``1``
|
||||
to disable split-KV which enables batch invariance. Split-KV
|
||||
parallelizes the key/value sequence dimension across multiple thread
|
||||
blocks and combines partial results. The split decision depends
|
||||
on ``max_k`` (the longest sequence in the batch), so different batch
|
||||
compositions can change the reduction order and produce different
|
||||
floating-point results for the same sequence. When this is disabled,
|
||||
bitwise identical outputs are guaranteed for a given sequence
|
||||
regardless of what other sequences are in the batch, at the
|
||||
cost of lower GPU utilization when there are few queries. When
|
||||
``None`` (default), the kernel chooses automatically.
|
||||
|
||||
Returns:
|
||||
output (Tensor): Output tensor from attention computation; shape :math:`(T_q, H_q, D)`.
|
||||
|
||||
If ``return_aux`` is not None and ``return_aux.lse`` is True:
|
||||
lse (Tensor): Log-sum-exp of attention scores; shape :math:`(T_q, H_q)`.
|
||||
|
||||
Shape legend:
|
||||
- :math:`N`: Batch size
|
||||
- :math:`T_q`: Total number of query tokens in the batch (sum of all query sequence lengths)
|
||||
- :math:`T_k`: Total number of key/value tokens in the batch (sum of all key/value sequence lengths)
|
||||
- :math:`H_q`: Number of query attention heads
|
||||
- :math:`H_{kv}`: Number of key/value attention heads (equal to :math:`H_q` unless GQA is enabled)
|
||||
- :math:`D`: Head dimension
|
||||
|
||||
Example::
|
||||
|
||||
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
|
||||
>>> batch_size, max_seq_len, embed_dim, num_heads = 2, 512, 1024, 16
|
||||
>>> head_dim = embed_dim // num_heads
|
||||
>>> seq_lengths = []
|
||||
>>> for _ in range(batch_size):
|
||||
... length = torch.randint(1, max_seq_len // 64 + 1, (1,)).item() * 64
|
||||
... seq_lengths.append(min(length, max_seq_len))
|
||||
>>> seq_lengths = torch.tensor(seq_lengths, device="cuda")
|
||||
>>> total_tokens = seq_lengths.sum().item()
|
||||
>>>
|
||||
>>> # Create packed query, key, value tensors
|
||||
>>> query = torch.randn(
|
||||
... total_tokens, num_heads, head_dim, dtype=torch.float16, device="cuda"
|
||||
... )
|
||||
>>> key = torch.randn(
|
||||
... total_tokens, num_heads, head_dim, dtype=torch.float16, device="cuda"
|
||||
... )
|
||||
>>> value = torch.randn(
|
||||
... total_tokens, num_heads, head_dim, dtype=torch.float16, device="cuda"
|
||||
... )
|
||||
>>>
|
||||
>>> # Build cumulative sequence tensor
|
||||
>>> cu_seq = torch.zeros(batch_size + 1, device="cuda", dtype=torch.int32)
|
||||
>>> cu_seq[1:] = seq_lengths.cumsum(0)
|
||||
>>> max_len = seq_lengths.max().item()
|
||||
>>>
|
||||
>>> # Call varlen_attn
|
||||
>>> output = varlen_attn(
|
||||
... query, key, value, cu_seq, cu_seq, max_len, max_len
|
||||
... )
|
||||
"""
|
||||
|
||||
num_heads_q = query.size(1)
|
||||
num_heads_k = key.size(2) if block_table is not None else key.size(1)
|
||||
if not enable_gqa and num_heads_q != num_heads_k:
|
||||
raise ValueError(
|
||||
f"Expect query and key/value to have the same number of heads "
|
||||
f"but got Hq={num_heads_q} and Hkv={num_heads_k}. "
|
||||
f"Try setting enable_gqa=True for GQA."
|
||||
)
|
||||
if enable_gqa and num_heads_q % num_heads_k != 0:
|
||||
raise ValueError(
|
||||
f"Expect number of query heads to be a multiple of kv heads for GQA "
|
||||
f"but got Hq={num_heads_q} and Hkv={num_heads_k}."
|
||||
)
|
||||
|
||||
is_causal = window_size == (-1, 0)
|
||||
out, lse, _ = torch.ops.torch_attn._varlen_attn(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cu_seq_q,
|
||||
cu_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
is_causal,
|
||||
scale,
|
||||
list(window_size),
|
||||
enable_gqa,
|
||||
seqused_k,
|
||||
block_table,
|
||||
num_splits,
|
||||
)
|
||||
if return_aux is not None and return_aux.lse:
|
||||
return out, lse
|
||||
return out
|
||||
|
||||
|
||||
@torch.library.custom_op("torch_attn::_varlen_attn_out", mutates_args={"out"})
|
||||
def _varlen_attn_out(
|
||||
out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor,
|
||||
cu_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
is_causal: bool = False,
|
||||
scale: float | None = None,
|
||||
window_size: list[int] | None = None,
|
||||
enable_gqa: bool = False,
|
||||
seqused_k: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Private custom op for variable-length attention with pre-allocated output.
|
||||
Same as _varlen_attn but writes the attention output into the provided out tensor.
|
||||
"""
|
||||
window_size = _normalize_window_size(window_size)
|
||||
|
||||
use_cudnn = query.is_cuda and _should_use_cudnn(query.device.index)
|
||||
|
||||
if use_cudnn:
|
||||
# TODO: look into this
|
||||
raise RuntimeError("cuDNN backend does not support out variant.")
|
||||
|
||||
log.info("Using Flash Attention backend for varlen_attn_out")
|
||||
softmax_lse = torch.ops.aten._flash_attention_forward_no_dropout_inplace(
|
||||
out,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cu_seq_q,
|
||||
cu_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
0.0, # dropout_p hardcoded to 0.0
|
||||
is_causal,
|
||||
False, # return_debug_mask
|
||||
scale=scale,
|
||||
window_size_left=window_size[0],
|
||||
window_size_right=window_size[1],
|
||||
seqused_k=seqused_k,
|
||||
block_table=block_table,
|
||||
num_splits=num_splits,
|
||||
)
|
||||
|
||||
return softmax_lse
|
||||
|
||||
|
||||
@_varlen_attn_out.register_fake
|
||||
def _varlen_attn_out_fake(
|
||||
out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor,
|
||||
cu_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
is_causal: bool = False,
|
||||
scale: float | None = None,
|
||||
window_size: list[int] | None = None,
|
||||
enable_gqa: bool = False,
|
||||
seqused_k: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Fake implementation for meta tensor computation and tracing.
|
||||
"""
|
||||
total_q = query.size(0)
|
||||
num_heads = query.size(1)
|
||||
logsumexp = torch.empty(
|
||||
(num_heads, total_q), dtype=torch.float, device=query.device
|
||||
)
|
||||
|
||||
if torch.version.hip:
|
||||
preferred = torch._C._get_rocm_fa_preferred_backend()
|
||||
if preferred == torch._C._ROCmFABackend.AOTriton:
|
||||
batch_size = cu_seq_q.size(0) - 1
|
||||
logsumexp = torch.empty(
|
||||
(batch_size, num_heads, max_q), dtype=torch.float, device=query.device
|
||||
)
|
||||
|
||||
return logsumexp
|
||||
|
||||
|
||||
def varlen_attn_out(
|
||||
out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor,
|
||||
cu_seq_k: torch.Tensor | None,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
*,
|
||||
return_aux: AuxRequest | None = None,
|
||||
scale: float | None = None,
|
||||
window_size: tuple[int, int] = (-1, -1),
|
||||
enable_gqa: bool = False,
|
||||
seqused_k: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
num_splits: int | None = None,
|
||||
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
|
||||
r"""Compute variable-length attention using Flash Attention with a pre-allocated output tensor.
|
||||
|
||||
Same as :func:`varlen_attn` but writes the attention output into the provided ``out`` tensor
|
||||
instead of allocating a new one.
|
||||
|
||||
"""
|
||||
num_heads_q = query.size(1)
|
||||
num_heads_k = key.size(2) if block_table is not None else key.size(1)
|
||||
if not enable_gqa and num_heads_q != num_heads_k:
|
||||
raise ValueError(
|
||||
f"Expect query and key/value to have the same number of heads "
|
||||
f"but got Hq={num_heads_q} and Hkv={num_heads_k}. "
|
||||
f"Try setting enable_gqa=True for GQA."
|
||||
)
|
||||
if enable_gqa and num_heads_q % num_heads_k != 0:
|
||||
raise ValueError(
|
||||
f"Expect number of query heads to be a multiple of kv heads for GQA "
|
||||
f"but got Hq={num_heads_q} and Hkv={num_heads_k}."
|
||||
)
|
||||
|
||||
is_causal = window_size == (-1, 0)
|
||||
lse = torch.ops.torch_attn._varlen_attn_out(
|
||||
out,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cu_seq_q,
|
||||
cu_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
is_causal,
|
||||
scale,
|
||||
list(window_size),
|
||||
enable_gqa,
|
||||
seqused_k,
|
||||
block_table,
|
||||
num_splits,
|
||||
)
|
||||
if return_aux is not None and return_aux.lse:
|
||||
return out, lse
|
||||
return out
|
||||
|
||||
|
||||
def _setup_context(ctx: Any, inputs: tuple[Any, ...], output: Any) -> None:
|
||||
(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
cu_seq_q,
|
||||
cu_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
is_causal,
|
||||
scale,
|
||||
window_size,
|
||||
enable_gqa,
|
||||
seqused_k,
|
||||
block_table,
|
||||
num_splits,
|
||||
) = inputs
|
||||
out, lse, rng_state = output
|
||||
|
||||
if seqused_k is not None:
|
||||
raise RuntimeError("seqused_k is an inference-only parameter.")
|
||||
if block_table is not None:
|
||||
raise RuntimeError("block_table is an inference-only parameter.")
|
||||
|
||||
ctx.save_for_backward(query, key, value, cu_seq_q, cu_seq_k, out, lse, rng_state)
|
||||
|
||||
ctx.max_q = max_q
|
||||
ctx.max_k = max_k
|
||||
ctx.is_causal = is_causal
|
||||
ctx.scale = scale
|
||||
ctx.window_size = window_size
|
||||
|
||||
|
||||
@torch.library.custom_op("torch_attn::_varlen_attn_backward", mutates_args={})
|
||||
def _varlen_attn_backward(
|
||||
grad_out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
lse: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor,
|
||||
cu_seq_k: torch.Tensor,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
is_causal: bool,
|
||||
rng_state: torch.Tensor,
|
||||
scale: float | None = None,
|
||||
window_size: list[int] | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
window_size = _normalize_window_size(window_size)
|
||||
|
||||
unused = torch.empty(0, device=query.device)
|
||||
|
||||
use_cudnn = query.is_cuda and _should_use_cudnn(query.device.index)
|
||||
if use_cudnn:
|
||||
log.info("Using cuDNN backend for varlen_attn")
|
||||
if window_size[0] != -1 or window_size[1] != -1:
|
||||
raise RuntimeError(
|
||||
"cuDNN backend does not support window attention. Please use Flash Attention backend."
|
||||
)
|
||||
dq, dk, dv = torch.ops.aten._cudnn_attention_backward(
|
||||
grad_out,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
out,
|
||||
lse,
|
||||
cu_seq_q,
|
||||
cu_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
0.0,
|
||||
is_causal,
|
||||
rng_state,
|
||||
unused,
|
||||
scale=scale,
|
||||
)
|
||||
else:
|
||||
log.info("Using Flash Attention backend for varlen_attn")
|
||||
dq, dk, dv = torch.ops.aten._flash_attention_backward(
|
||||
grad_out,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
out,
|
||||
lse,
|
||||
cu_seq_q,
|
||||
cu_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
0.0,
|
||||
is_causal,
|
||||
rng_state,
|
||||
unused,
|
||||
scale=scale,
|
||||
window_size_left=window_size[0],
|
||||
window_size_right=window_size[1],
|
||||
)
|
||||
return dq, dk, dv
|
||||
|
||||
|
||||
@_varlen_attn_backward.register_fake
|
||||
def _varlen_attn_backward_fake(
|
||||
grad_out: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
lse: torch.Tensor,
|
||||
cu_seq_q: torch.Tensor,
|
||||
cu_seq_k: torch.Tensor,
|
||||
max_q: int,
|
||||
max_k: int,
|
||||
is_causal: bool,
|
||||
rng_state: torch.Tensor,
|
||||
scale: float | None = None,
|
||||
window_size: list[int] | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Fake implementation for meta tensor computation and tracing.
|
||||
"""
|
||||
window_size = _normalize_window_size(window_size)
|
||||
|
||||
grad_query = torch.empty_like(query)
|
||||
grad_key = torch.empty_like(key)
|
||||
grad_value = torch.empty_like(value)
|
||||
|
||||
return grad_query, grad_key, grad_value
|
||||
|
||||
|
||||
def _backward(
|
||||
ctx: Any, grad_out: torch.Tensor, grad_lse: torch.Tensor, grad_rng: torch.Tensor
|
||||
) -> tuple[torch.Tensor | None, ...]:
|
||||
query, key, value, cu_seq_q, cu_seq_k, out, lse, rng_state = ctx.saved_tensors
|
||||
|
||||
max_q = ctx.max_q
|
||||
max_k = ctx.max_k
|
||||
is_causal = ctx.is_causal
|
||||
scale = ctx.scale
|
||||
window_size = ctx.window_size
|
||||
|
||||
dq, dk, dv = torch.ops.torch_attn._varlen_attn_backward(
|
||||
grad_out,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
out,
|
||||
lse,
|
||||
cu_seq_q,
|
||||
cu_seq_k,
|
||||
max_q,
|
||||
max_k,
|
||||
is_causal,
|
||||
rng_state,
|
||||
scale,
|
||||
window_size,
|
||||
)
|
||||
# cu_seq_q, cu_seq_k, max_q, max_k, is_causal, scale, window_size, \
|
||||
# enable_gqa, seqused_k, block_table, num_splits
|
||||
num_params = 11
|
||||
return (dq, dk, dv, *((None,) * num_params))
|
||||
|
||||
|
||||
_varlen_attn.register_autograd(_backward, setup_context=_setup_context)
|
||||
|
||||
torch._dynamo.disallow_in_graph(
|
||||
torch.ops.aten._flash_attention_forward_no_dropout_inplace
|
||||
)
|
||||
|
||||
from torch.utils.flop_counter import (
|
||||
_varlen_attn_backward_flop,
|
||||
_varlen_attn_forward_flop,
|
||||
_varlen_attn_out_flop,
|
||||
flop_registry,
|
||||
)
|
||||
|
||||
|
||||
flop_registry[torch.ops.torch_attn._varlen_attn] = _varlen_attn_forward_flop
|
||||
flop_registry[torch.ops.torch_attn._varlen_attn_out] = _varlen_attn_out_flop
|
||||
flop_registry[torch.ops.torch_attn._varlen_attn_backward] = _varlen_attn_backward_flop
|
||||
@@ -0,0 +1,6 @@
|
||||
# mypy: allow-untyped-defs
|
||||
# this is for historical pickle deserialization, it is not used otherwise
|
||||
|
||||
|
||||
def _get_thnn_function_backend() -> None:
|
||||
pass
|
||||
@@ -0,0 +1,46 @@
|
||||
from typing import TypeAlias as _TypeAlias, TypeVar
|
||||
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
# ruff: noqa: PYI042,PYI047
|
||||
|
||||
# Create some useful type aliases
|
||||
|
||||
# Template for arguments which can be supplied as a tuple, or which can be a scalar which PyTorch will internally
|
||||
# broadcast to a tuple.
|
||||
# Comes in several variants: A tuple of unknown size, and a fixed-size tuple for 1d, 2d, or 3d operations.
|
||||
T = TypeVar("T")
|
||||
_scalar_or_tuple_any_t: _TypeAlias = T | tuple[T, ...]
|
||||
_scalar_or_tuple_1_t: _TypeAlias = T | tuple[T]
|
||||
_scalar_or_tuple_2_t: _TypeAlias = T | tuple[T, T]
|
||||
_scalar_or_tuple_3_t: _TypeAlias = T | tuple[T, T, T]
|
||||
_scalar_or_tuple_4_t: _TypeAlias = T | tuple[T, T, T, T]
|
||||
_scalar_or_tuple_5_t: _TypeAlias = T | tuple[T, T, T, T, T]
|
||||
_scalar_or_tuple_6_t: _TypeAlias = T | tuple[T, T, T, T, T, T]
|
||||
|
||||
# For arguments which represent size parameters (eg, kernel size, padding)
|
||||
_size_any_t: _TypeAlias = _scalar_or_tuple_any_t[int]
|
||||
_size_1_t: _TypeAlias = _scalar_or_tuple_1_t[int]
|
||||
_size_2_t: _TypeAlias = _scalar_or_tuple_2_t[int]
|
||||
_size_3_t: _TypeAlias = _scalar_or_tuple_3_t[int]
|
||||
_size_4_t: _TypeAlias = _scalar_or_tuple_4_t[int]
|
||||
_size_5_t: _TypeAlias = _scalar_or_tuple_5_t[int]
|
||||
_size_6_t: _TypeAlias = _scalar_or_tuple_6_t[int]
|
||||
|
||||
# For arguments which represent optional size parameters (eg, adaptive pool parameters)
|
||||
_size_any_opt_t: _TypeAlias = _scalar_or_tuple_any_t[int | None]
|
||||
_size_2_opt_t: _TypeAlias = _scalar_or_tuple_2_t[int | None]
|
||||
_size_3_opt_t: _TypeAlias = _scalar_or_tuple_3_t[int | None]
|
||||
|
||||
# For arguments that represent a ratio to adjust each dimension of an input with (eg, upsampling parameters)
|
||||
_ratio_2_t: _TypeAlias = _scalar_or_tuple_2_t[float]
|
||||
_ratio_3_t: _TypeAlias = _scalar_or_tuple_3_t[float]
|
||||
_ratio_any_t: _TypeAlias = _scalar_or_tuple_any_t[float]
|
||||
|
||||
_tensor_list_t: _TypeAlias = _scalar_or_tuple_any_t[Tensor]
|
||||
|
||||
# For the return value of max pooling operations that may or may not return indices.
|
||||
# With the proposed 'Literal' feature to Python typing, it might be possible to
|
||||
# eventually eliminate this.
|
||||
_maybe_indices_t: _TypeAlias = _scalar_or_tuple_2_t[Tensor]
|
||||
@@ -0,0 +1,90 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""Functionality for Python <-> C++ frontend inter-op."""
|
||||
|
||||
from torch import nn
|
||||
|
||||
|
||||
class OrderedDictWrapper:
|
||||
"""A wrapper around a C++ OrderedDict.
|
||||
|
||||
It dynamically evaluates the OrderedDict getter on a bound C++ module, such
|
||||
that new changes on the C++ side are picked up. Otherwise accessing e.g.
|
||||
``cpp_module._parameters`` just once would get a frozen copy of the parameters
|
||||
at the time of access. ``torch.nn.Module`` accesses ``_parameters`` et al. via ``self.__dict__``
|
||||
so using properties does not work.
|
||||
"""
|
||||
|
||||
def __init__(self, cpp_module, attr) -> None:
|
||||
self.cpp_module = cpp_module
|
||||
self.attr = attr
|
||||
|
||||
@property
|
||||
def cpp_dict(self):
|
||||
return getattr(self.cpp_module, self.attr)
|
||||
|
||||
# Magic methods cannot be assigned dynamically and bypass ``getattr``, so we
|
||||
# must manually override them.
|
||||
|
||||
def items(self):
|
||||
return self.cpp_dict.items()
|
||||
|
||||
def keys(self):
|
||||
return self.cpp_dict.keys()
|
||||
|
||||
def values(self):
|
||||
return self.cpp_dict.values()
|
||||
|
||||
def __iter__(self):
|
||||
return self.cpp_dict.__iter__()
|
||||
|
||||
def __len__(self) -> int:
|
||||
return self.cpp_dict.__len__()
|
||||
|
||||
def __contains__(self, key) -> bool:
|
||||
return self.cpp_dict.__contains__(key)
|
||||
|
||||
def __getitem__(self, key):
|
||||
return self.cpp_dict.__getitem__(key)
|
||||
|
||||
|
||||
class ModuleWrapper(nn.Module):
|
||||
"""A subclass of ``torch.nn.Module`` that wraps a C++ frontend module and delegates all access."""
|
||||
|
||||
def __init__(self, cpp_module) -> None:
|
||||
# Assign before the super class constructor so ``self.training`` can be
|
||||
# assigned to in the super class constructor.
|
||||
self.cpp_module = cpp_module
|
||||
super().__init__()
|
||||
self._parameters = OrderedDictWrapper(cpp_module, "_parameters") # type: ignore[assignment]
|
||||
self._buffers: OrderedDictWrapper = OrderedDictWrapper(cpp_module, "_buffers") # type: ignore[assignment]
|
||||
self._modules: OrderedDictWrapper = OrderedDictWrapper(cpp_module, "_modules") # type: ignore[assignment]
|
||||
for attr in dir(cpp_module):
|
||||
# Skip magic methods and the three attributes above.
|
||||
if not attr.startswith("_"):
|
||||
setattr(self, attr, getattr(self.cpp_module, attr))
|
||||
|
||||
def _apply(self, fn, recurse=True):
|
||||
for param in self.parameters():
|
||||
# Tensors stored in modules are graph leaves, and we don't
|
||||
# want to create copy nodes, so we have to unpack the data.
|
||||
param.data = fn(param.data)
|
||||
if param._grad is not None:
|
||||
param._grad.data = fn(param._grad.data)
|
||||
|
||||
for buf in self.buffers():
|
||||
buf.data = fn(buf.data)
|
||||
|
||||
return self
|
||||
|
||||
# nn.Module defines training as a boolean
|
||||
@property # type: ignore[override]
|
||||
# pyrefly: ignore [bad-override]
|
||||
def training(self):
|
||||
return self.cpp_module.training
|
||||
|
||||
@training.setter
|
||||
def training(self, mode) -> None:
|
||||
self.cpp_module.train(mode)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return self.cpp_module.__repr__()
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,298 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""Gradient interface."""
|
||||
|
||||
import torch
|
||||
from torch.nn.modules.utils import _pair, _single, _triple
|
||||
|
||||
|
||||
def conv1d_input(
|
||||
input_size,
|
||||
weight,
|
||||
grad_output,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
):
|
||||
r"""Compute the gradient of conv1d with respect to the input of the convolution.
|
||||
|
||||
This is same as the 1D transposed convolution operator under the hood but requires
|
||||
the shape of the gradient w.r.t. input to be specified explicitly.
|
||||
|
||||
Args:
|
||||
input_size : Shape of the input gradient tensor
|
||||
weight: weight tensor (out_channels x in_channels/groups x kW)
|
||||
grad_output : output gradient tensor (minibatch x out_channels x oW)
|
||||
stride (int or tuple, optional): Stride of the convolution. Default: 1
|
||||
padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0
|
||||
dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
|
||||
groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
|
||||
|
||||
Examples::
|
||||
|
||||
>>> input = torch.randn(1, 1, 3, requires_grad=True)
|
||||
>>> weight = torch.randn(1, 1, 1, requires_grad=True)
|
||||
>>> output = F.conv1d(input, weight)
|
||||
>>> grad_output = torch.randn(output.shape)
|
||||
>>> grad_input = torch.autograd.grad(output, input, grad_output)
|
||||
>>> F.grad.conv1d_input(input.shape, weight, grad_output)
|
||||
|
||||
"""
|
||||
input = grad_output.new_empty(1).expand(input_size)
|
||||
|
||||
return torch.ops.aten.convolution_backward(
|
||||
grad_output,
|
||||
input,
|
||||
weight,
|
||||
None,
|
||||
_single(stride),
|
||||
_single(padding),
|
||||
_single(dilation),
|
||||
False,
|
||||
[0],
|
||||
groups,
|
||||
(True, False, False),
|
||||
)[0]
|
||||
|
||||
|
||||
def conv1d_weight(
|
||||
input,
|
||||
weight_size,
|
||||
grad_output,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
):
|
||||
r"""Compute the gradient of conv1d with respect to the weight of the convolution.
|
||||
|
||||
Args:
|
||||
input: input tensor of shape (minibatch x in_channels x iW)
|
||||
weight_size : Shape of the weight gradient tensor
|
||||
grad_output : output gradient tensor (minibatch x out_channels x oW)
|
||||
stride (int or tuple, optional): Stride of the convolution. Default: 1
|
||||
padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0
|
||||
dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
|
||||
groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
|
||||
|
||||
Examples::
|
||||
|
||||
>>> input = torch.randn(1, 1, 3, requires_grad=True)
|
||||
>>> weight = torch.randn(1, 1, 1, requires_grad=True)
|
||||
>>> output = F.conv1d(input, weight)
|
||||
>>> grad_output = torch.randn(output.shape)
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> grad_weight = torch.autograd.grad(output, filter, grad_output)
|
||||
>>> F.grad.conv1d_weight(input, weight.shape, grad_output)
|
||||
|
||||
"""
|
||||
weight = grad_output.new_empty(1).expand(weight_size)
|
||||
|
||||
return torch.ops.aten.convolution_backward(
|
||||
grad_output,
|
||||
input,
|
||||
weight,
|
||||
None,
|
||||
_single(stride),
|
||||
_single(padding),
|
||||
_single(dilation),
|
||||
False,
|
||||
[0],
|
||||
groups,
|
||||
(False, True, False),
|
||||
)[1]
|
||||
|
||||
|
||||
def conv2d_input(
|
||||
input_size,
|
||||
weight,
|
||||
grad_output,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
):
|
||||
r"""Compute the gradient of conv2d with respect to the input of the convolution.
|
||||
|
||||
This is same as the 2D transposed convolution operator under the hood but requires
|
||||
the shape of the gradient w.r.t. input to be specified explicitly.
|
||||
|
||||
Args:
|
||||
input_size : Shape of the input gradient tensor
|
||||
weight: weight tensor (out_channels x in_channels/groups x kH x kW)
|
||||
grad_output : output gradient tensor (minibatch x out_channels x oH x oW)
|
||||
stride (int or tuple, optional): Stride of the convolution. Default: 1
|
||||
padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0
|
||||
dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
|
||||
groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
|
||||
|
||||
Examples::
|
||||
|
||||
>>> input = torch.randn(1, 1, 3, 3, requires_grad=True)
|
||||
>>> weight = torch.randn(1, 1, 1, 2, requires_grad=True)
|
||||
>>> output = F.conv2d(input, weight)
|
||||
>>> grad_output = torch.randn(output.shape)
|
||||
>>> grad_input = torch.autograd.grad(output, input, grad_output)
|
||||
>>> F.grad.conv2d_input(input.shape, weight, grad_output)
|
||||
|
||||
"""
|
||||
input = grad_output.new_empty(1).expand(input_size)
|
||||
|
||||
return torch.ops.aten.convolution_backward(
|
||||
grad_output,
|
||||
input,
|
||||
weight,
|
||||
None,
|
||||
_pair(stride),
|
||||
_pair(padding),
|
||||
_pair(dilation),
|
||||
False,
|
||||
[0],
|
||||
groups,
|
||||
(True, False, False),
|
||||
)[0]
|
||||
|
||||
|
||||
def conv2d_weight(
|
||||
input,
|
||||
weight_size,
|
||||
grad_output,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
):
|
||||
r"""Compute the gradient of conv2d with respect to the weight of the convolution.
|
||||
|
||||
Args:
|
||||
input: input tensor of shape (minibatch x in_channels x iH x iW)
|
||||
weight_size : Shape of the weight gradient tensor
|
||||
grad_output : output gradient tensor (minibatch x out_channels x oH x oW)
|
||||
stride (int or tuple, optional): Stride of the convolution. Default: 1
|
||||
padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0
|
||||
dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
|
||||
groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
|
||||
|
||||
Examples::
|
||||
|
||||
>>> input = torch.randn(1, 1, 3, 3, requires_grad=True)
|
||||
>>> weight = torch.randn(1, 1, 1, 2, requires_grad=True)
|
||||
>>> output = F.conv2d(input, weight)
|
||||
>>> grad_output = torch.randn(output.shape)
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> grad_weight = torch.autograd.grad(output, filter, grad_output)
|
||||
>>> F.grad.conv2d_weight(input, weight.shape, grad_output)
|
||||
|
||||
"""
|
||||
weight = grad_output.new_empty(1).expand(weight_size)
|
||||
|
||||
return torch.ops.aten.convolution_backward(
|
||||
grad_output,
|
||||
input,
|
||||
weight,
|
||||
None,
|
||||
_pair(stride),
|
||||
_pair(padding),
|
||||
_pair(dilation),
|
||||
False,
|
||||
[0],
|
||||
groups,
|
||||
(False, True, False),
|
||||
)[1]
|
||||
|
||||
|
||||
def conv3d_input(
|
||||
input_size,
|
||||
weight,
|
||||
grad_output,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
):
|
||||
r"""Compute the gradient of conv3d with respect to the input of the convolution.
|
||||
|
||||
This is same as the 3D transposed convolution operator under the hood but requires
|
||||
the shape of the gradient w.r.t. input to be specified explicitly.
|
||||
|
||||
Args:
|
||||
input_size : Shape of the input gradient tensor
|
||||
weight: weights tensor (out_channels x in_channels/groups x kT x kH x kW)
|
||||
grad_output : output gradient tensor (minibatch x out_channels x oT x oH x oW)
|
||||
stride (int or tuple, optional): Stride of the convolution. Default: 1
|
||||
padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0
|
||||
dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
|
||||
groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
|
||||
|
||||
Examples::
|
||||
|
||||
>>> input = torch.randn(2, 8, 10, 10, 20, requires_grad=True)
|
||||
>>> weight = torch.randn(4, 8, 2, 3, 3, requires_grad=True)
|
||||
>>> output = F.conv3d(input, weight)
|
||||
>>> grad_output = torch.randn(output.shape)
|
||||
>>> grad_input = torch.autograd.grad(output, input, grad_output)
|
||||
>>> F.grad.conv3d_input(input.shape, weight, grad_output)
|
||||
|
||||
"""
|
||||
input = grad_output.new_empty(1).expand(input_size)
|
||||
|
||||
return torch.ops.aten.convolution_backward(
|
||||
grad_output,
|
||||
input,
|
||||
weight,
|
||||
None,
|
||||
_triple(stride),
|
||||
_triple(padding),
|
||||
_triple(dilation),
|
||||
False,
|
||||
[0],
|
||||
groups,
|
||||
(True, False, False),
|
||||
)[0]
|
||||
|
||||
|
||||
def conv3d_weight(
|
||||
input,
|
||||
weight_size,
|
||||
grad_output,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
):
|
||||
r"""Compute the gradient of conv3d with respect to the weight of the convolution.
|
||||
|
||||
Args:
|
||||
input: input tensor of shape (minibatch x in_channels x iT x iH x iW)
|
||||
weight_size : Shape of the weight gradient tensor
|
||||
grad_output : output gradient tensor (minibatch x out_channels x oT x oH x oW)
|
||||
stride (int or tuple, optional): Stride of the convolution. Default: 1
|
||||
padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0
|
||||
dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
|
||||
groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
|
||||
|
||||
Examples::
|
||||
|
||||
>>> input = torch.randn(2, 8, 10, 10, 20, requires_grad=True)
|
||||
>>> weight = torch.randn(4, 8, 2, 3, 3, requires_grad=True)
|
||||
>>> output = F.conv3d(input, weight)
|
||||
>>> grad_output = torch.randn(output.shape)
|
||||
>>> grad_weight = torch.autograd.grad(output, weight, grad_output)
|
||||
>>> F.grad.conv3d_weight(input, weight.shape, grad_output)
|
||||
|
||||
"""
|
||||
weight = grad_output.new_empty(1).expand(weight_size)
|
||||
|
||||
return torch.ops.aten.convolution_backward(
|
||||
grad_output,
|
||||
input,
|
||||
weight,
|
||||
None,
|
||||
_triple(stride),
|
||||
_triple(padding),
|
||||
_triple(dilation),
|
||||
False,
|
||||
[0],
|
||||
groups,
|
||||
(False, True, False),
|
||||
)[1]
|
||||
@@ -0,0 +1,798 @@
|
||||
"""This file contains utilities for initializing neural network parameters."""
|
||||
|
||||
import math
|
||||
import warnings
|
||||
from collections.abc import Callable
|
||||
from typing import Literal, TypeVar
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
__all__ = [
|
||||
"calculate_gain",
|
||||
"uniform_",
|
||||
"normal_",
|
||||
"trunc_normal_",
|
||||
"constant_",
|
||||
"ones_",
|
||||
"zeros_",
|
||||
"eye_",
|
||||
"dirac_",
|
||||
"xavier_uniform_",
|
||||
"xavier_normal_",
|
||||
"kaiming_uniform_",
|
||||
"kaiming_normal_",
|
||||
"orthogonal_",
|
||||
"sparse_",
|
||||
# Deprecated aliases (for backward compatibility)
|
||||
"uniform",
|
||||
"normal",
|
||||
"constant",
|
||||
"eye",
|
||||
"dirac",
|
||||
"xavier_uniform",
|
||||
"xavier_normal",
|
||||
"kaiming_uniform",
|
||||
"kaiming_normal",
|
||||
"orthogonal",
|
||||
"sparse",
|
||||
]
|
||||
|
||||
|
||||
_R = TypeVar("_R")
|
||||
_P = ParamSpec("_P")
|
||||
|
||||
_NonlinearityType = Literal[
|
||||
"linear",
|
||||
"conv1d",
|
||||
"conv2d",
|
||||
"conv3d",
|
||||
"conv_transpose1d",
|
||||
"conv_transpose2d",
|
||||
"conv_transpose3d",
|
||||
"sigmoid",
|
||||
"tanh",
|
||||
"relu",
|
||||
"leaky_relu",
|
||||
"selu",
|
||||
]
|
||||
|
||||
_FanMode = Literal["fan_in", "fan_out"]
|
||||
|
||||
|
||||
# These no_grad_* functions are necessary as wrappers around the parts of these
|
||||
# functions that use `with torch.no_grad()`. The JIT doesn't support context
|
||||
# managers, so these need to be implemented as builtins. Using these wrappers
|
||||
# lets us keep those builtins small and reusable.
|
||||
def _no_grad_uniform_(
|
||||
tensor: Tensor, a: float, b: float, generator: torch.Generator | None = None
|
||||
) -> Tensor:
|
||||
with torch.no_grad():
|
||||
return tensor.uniform_(a, b, generator=generator)
|
||||
|
||||
|
||||
def _no_grad_normal_(
|
||||
tensor: Tensor,
|
||||
mean: float,
|
||||
std: float,
|
||||
generator: torch.Generator | None = None,
|
||||
) -> Tensor:
|
||||
with torch.no_grad():
|
||||
return tensor.normal_(mean, std, generator=generator)
|
||||
|
||||
|
||||
def _no_grad_trunc_normal_(
|
||||
tensor: Tensor,
|
||||
mean: float,
|
||||
std: float,
|
||||
a: float,
|
||||
b: float,
|
||||
generator: torch.Generator | None = None,
|
||||
) -> Tensor:
|
||||
# Meta tensors have no storage, so sampling is a no-op.
|
||||
if tensor.is_meta:
|
||||
return tensor
|
||||
|
||||
def norm_cdf(x: float) -> float:
|
||||
return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0
|
||||
|
||||
if (mean < a - 2 * std) or (mean > b + 2 * std):
|
||||
warnings.warn(
|
||||
"mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
|
||||
"The distribution of values may be incorrect.",
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
p = norm_cdf((b - mean) / std) - norm_cdf((a - mean) / std)
|
||||
|
||||
if p > 0.3:
|
||||
# Cast bounds to tensor dtype so the rejection mask is consistent
|
||||
# with the tensor's representable values.
|
||||
lo = tensor.new_tensor(a, device="cpu").item()
|
||||
hi = tensor.new_tensor(b, device="cpu").item()
|
||||
result = tensor.normal_(mean, std, generator=generator)
|
||||
while True:
|
||||
mask = (result < lo) | (result > hi)
|
||||
if not mask.any():
|
||||
break
|
||||
result = torch.where(
|
||||
mask,
|
||||
torch.empty_like(result).normal_(mean, std, generator=generator),
|
||||
result,
|
||||
)
|
||||
if tensor is not result:
|
||||
tensor.copy_(result)
|
||||
else:
|
||||
mode = max(a, min(mean, b))
|
||||
log_peak = -0.5 * ((mode - mean) / std) ** 2
|
||||
|
||||
candidates = torch.empty_like(tensor)
|
||||
accept_buf = torch.empty_like(tensor)
|
||||
|
||||
# First iteration: sample directly into tensor to avoid
|
||||
# a where() + copy_() if all samples are accepted.
|
||||
tensor.uniform_(a, b, generator=generator)
|
||||
candidates.copy_(tensor)
|
||||
# log_pdf = -0.5 * ((candidates - mean) / std) ** 2
|
||||
candidates.sub_(mean).div_(std).pow_(2).mul_(-0.5).sub_(log_peak)
|
||||
pending = accept_buf.uniform_(generator=generator).log_().gt(candidates)
|
||||
if not pending.any():
|
||||
pass
|
||||
else:
|
||||
result = tensor
|
||||
while True:
|
||||
candidates.uniform_(a, b, generator=generator)
|
||||
result = torch.where(pending, candidates, result)
|
||||
# log_pdf = -0.5 * ((candidates - mean) / std) ** 2
|
||||
candidates.sub_(mean).div_(std).pow_(2).mul_(-0.5).sub_(log_peak)
|
||||
pending = torch.where(
|
||||
pending,
|
||||
accept_buf.uniform_(generator=generator).log_().gt(candidates),
|
||||
pending,
|
||||
)
|
||||
if not pending.any():
|
||||
break
|
||||
tensor.copy_(result)
|
||||
|
||||
return tensor
|
||||
|
||||
|
||||
def _no_grad_fill_(tensor: Tensor, val: float) -> Tensor:
|
||||
with torch.no_grad():
|
||||
return tensor.fill_(val)
|
||||
|
||||
|
||||
def _no_grad_zero_(tensor: Tensor) -> Tensor:
|
||||
with torch.no_grad():
|
||||
return tensor.zero_()
|
||||
|
||||
|
||||
def calculate_gain(
|
||||
nonlinearity: _NonlinearityType, param: int | float | None = None
|
||||
) -> float:
|
||||
r"""Return the recommended gain value for the given nonlinearity function.
|
||||
|
||||
The values are as follows:
|
||||
|
||||
================= ====================================================
|
||||
nonlinearity gain
|
||||
================= ====================================================
|
||||
Linear / Identity :math:`1`
|
||||
Conv{1,2,3}D :math:`1`
|
||||
Sigmoid :math:`1`
|
||||
Tanh :math:`\frac{5}{3}`
|
||||
ReLU :math:`\sqrt{2}`
|
||||
Leaky Relu :math:`\sqrt{\frac{2}{1 + \text{negative\_slope}^2}}`
|
||||
SELU :math:`\frac{3}{4}`
|
||||
================= ====================================================
|
||||
|
||||
.. warning::
|
||||
In order to implement `Self-Normalizing Neural Networks`_ ,
|
||||
you should use ``nonlinearity='linear'`` instead of ``nonlinearity='selu'``.
|
||||
This gives the initial weights a variance of ``1 / N``,
|
||||
which is necessary to induce a stable fixed point in the forward pass.
|
||||
In contrast, the default gain for ``SELU`` sacrifices the normalization
|
||||
effect for more stable gradient flow in rectangular layers.
|
||||
|
||||
Args:
|
||||
nonlinearity: the non-linear function (`nn.functional` name)
|
||||
param: optional parameter for the non-linear function
|
||||
|
||||
Examples:
|
||||
>>> gain = nn.init.calculate_gain(
|
||||
... "leaky_relu", 0.2
|
||||
... ) # leaky_relu with negative_slope=0.2
|
||||
|
||||
.. _Self-Normalizing Neural Networks: https://papers.nips.cc/paper/2017/hash/5d44ee6f2c3f71b73125876103c8f6c4-Abstract.html
|
||||
"""
|
||||
linear_fns = [
|
||||
"linear",
|
||||
"conv1d",
|
||||
"conv2d",
|
||||
"conv3d",
|
||||
"conv_transpose1d",
|
||||
"conv_transpose2d",
|
||||
"conv_transpose3d",
|
||||
]
|
||||
if nonlinearity in linear_fns or nonlinearity == "sigmoid":
|
||||
return 1
|
||||
elif nonlinearity == "tanh":
|
||||
return 5.0 / 3
|
||||
elif nonlinearity == "relu":
|
||||
return math.sqrt(2.0)
|
||||
elif nonlinearity == "leaky_relu":
|
||||
if param is None:
|
||||
negative_slope = 0.01
|
||||
elif (
|
||||
not isinstance(param, bool)
|
||||
and isinstance(param, int)
|
||||
or isinstance(param, float)
|
||||
):
|
||||
# True/False are instances of int, hence check above
|
||||
negative_slope = param
|
||||
else:
|
||||
raise ValueError(f"negative_slope {param} not a valid number")
|
||||
return math.sqrt(2.0 / (1 + negative_slope**2))
|
||||
elif nonlinearity == "selu":
|
||||
return (
|
||||
3.0 / 4
|
||||
) # Value found empirically (https://github.com/pytorch/pytorch/pull/50664)
|
||||
else:
|
||||
raise ValueError(f"Unsupported nonlinearity {nonlinearity}")
|
||||
|
||||
|
||||
def uniform_(
|
||||
tensor: Tensor,
|
||||
a: float = 0.0,
|
||||
b: float = 1.0,
|
||||
generator: torch.Generator | None = None,
|
||||
) -> Tensor:
|
||||
r"""Fill the input Tensor with values drawn from the uniform distribution.
|
||||
|
||||
:math:`\mathcal{U}(a, b)`.
|
||||
|
||||
Args:
|
||||
tensor: an n-dimensional `torch.Tensor`
|
||||
a: the lower bound of the uniform distribution
|
||||
b: the upper bound of the uniform distribution
|
||||
generator: the torch Generator to sample from (default: None)
|
||||
|
||||
Examples:
|
||||
>>> w = torch.empty(3, 5)
|
||||
>>> nn.init.uniform_(w)
|
||||
"""
|
||||
if torch.overrides.has_torch_function_variadic(tensor):
|
||||
return torch.overrides.handle_torch_function(
|
||||
uniform_, (tensor,), tensor=tensor, a=a, b=b, generator=generator
|
||||
)
|
||||
return _no_grad_uniform_(tensor, a, b, generator)
|
||||
|
||||
|
||||
def normal_(
|
||||
tensor: Tensor,
|
||||
mean: float = 0.0,
|
||||
std: float = 1.0,
|
||||
generator: torch.Generator | None = None,
|
||||
) -> Tensor:
|
||||
r"""Fill the input Tensor with values drawn from the normal distribution.
|
||||
|
||||
:math:`\mathcal{N}(\text{mean}, \text{std}^2)`.
|
||||
|
||||
Args:
|
||||
tensor: an n-dimensional `torch.Tensor`
|
||||
mean: the mean of the normal distribution
|
||||
std: the standard deviation of the normal distribution
|
||||
generator: the torch Generator to sample from (default: None)
|
||||
|
||||
Examples:
|
||||
>>> w = torch.empty(3, 5)
|
||||
>>> nn.init.normal_(w)
|
||||
"""
|
||||
if torch.overrides.has_torch_function_variadic(tensor):
|
||||
return torch.overrides.handle_torch_function(
|
||||
normal_, (tensor,), tensor=tensor, mean=mean, std=std, generator=generator
|
||||
)
|
||||
return _no_grad_normal_(tensor, mean, std, generator)
|
||||
|
||||
|
||||
def trunc_normal_(
|
||||
tensor: Tensor,
|
||||
mean: float = 0.0,
|
||||
std: float = 1.0,
|
||||
a: float = -2.0,
|
||||
b: float = 2.0,
|
||||
generator: torch.Generator | None = None,
|
||||
) -> Tensor:
|
||||
r"""Fill the input Tensor with values drawn from a truncated normal distribution.
|
||||
|
||||
The values are effectively drawn from the
|
||||
normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
|
||||
with values outside :math:`[a, b]` redrawn until they are within
|
||||
the bounds. The method used for generating the random values works
|
||||
best when :math:`a \leq \text{mean} \leq b`.
|
||||
|
||||
For reduced-precision types (``torch.float16`` and ``torch.bfloat16``),
|
||||
sampling quality depends on the underlying ``normal_()`` and ``uniform_()``
|
||||
implementations which operate at higher internal precision to avoid
|
||||
quantization artifacts.
|
||||
|
||||
Args:
|
||||
tensor: an n-dimensional `torch.Tensor`
|
||||
mean: the mean of the normal distribution
|
||||
std: the standard deviation of the normal distribution
|
||||
a: the minimum cutoff value
|
||||
b: the maximum cutoff value
|
||||
generator: the torch Generator to sample from (default: None)
|
||||
|
||||
Examples:
|
||||
>>> w = torch.empty(3, 5)
|
||||
>>> nn.init.trunc_normal_(w)
|
||||
"""
|
||||
return _no_grad_trunc_normal_(tensor, mean, std, a, b, generator=generator)
|
||||
|
||||
|
||||
def constant_(tensor: Tensor, val: float) -> Tensor:
|
||||
r"""Fill the input Tensor with the value :math:`\text{val}`.
|
||||
|
||||
Args:
|
||||
tensor: an n-dimensional `torch.Tensor`
|
||||
val: the value to fill the tensor with
|
||||
|
||||
Examples:
|
||||
>>> w = torch.empty(3, 5)
|
||||
>>> nn.init.constant_(w, 0.3)
|
||||
"""
|
||||
if torch.overrides.has_torch_function_variadic(tensor):
|
||||
return torch.overrides.handle_torch_function(
|
||||
constant_, (tensor,), tensor=tensor, val=val
|
||||
)
|
||||
return _no_grad_fill_(tensor, val)
|
||||
|
||||
|
||||
def ones_(tensor: Tensor) -> Tensor:
|
||||
r"""Fill the input Tensor with the scalar value `1`.
|
||||
|
||||
Args:
|
||||
tensor: an n-dimensional `torch.Tensor`
|
||||
|
||||
Examples:
|
||||
>>> w = torch.empty(3, 5)
|
||||
>>> nn.init.ones_(w)
|
||||
"""
|
||||
return _no_grad_fill_(tensor, 1.0)
|
||||
|
||||
|
||||
def zeros_(tensor: Tensor) -> Tensor:
|
||||
r"""Fill the input Tensor with the scalar value `0`.
|
||||
|
||||
Args:
|
||||
tensor: an n-dimensional `torch.Tensor`
|
||||
|
||||
Examples:
|
||||
>>> w = torch.empty(3, 5)
|
||||
>>> nn.init.zeros_(w)
|
||||
"""
|
||||
return _no_grad_zero_(tensor)
|
||||
|
||||
|
||||
def eye_(tensor: Tensor) -> Tensor:
|
||||
r"""Fill the 2-dimensional input `Tensor` with the identity matrix.
|
||||
|
||||
Preserves the identity of the inputs in `Linear` layers, where as
|
||||
many inputs are preserved as possible.
|
||||
|
||||
Args:
|
||||
tensor: a 2-dimensional `torch.Tensor`
|
||||
|
||||
Examples:
|
||||
>>> w = torch.empty(3, 5)
|
||||
>>> nn.init.eye_(w)
|
||||
"""
|
||||
if tensor.ndimension() != 2:
|
||||
raise ValueError("Only tensors with 2 dimensions are supported")
|
||||
|
||||
with torch.no_grad():
|
||||
torch.eye(*tensor.shape, out=tensor, requires_grad=tensor.requires_grad)
|
||||
return tensor
|
||||
|
||||
|
||||
def dirac_(tensor: Tensor, groups: int = 1) -> Tensor:
|
||||
r"""Fill the {3, 4, 5}-dimensional input `Tensor` with the Dirac delta function.
|
||||
|
||||
Preserves the identity of the inputs in `Convolutional`
|
||||
layers, where as many input channels are preserved as possible. In case
|
||||
of groups>1, each group of channels preserves identity
|
||||
|
||||
Args:
|
||||
tensor: a {3, 4, 5}-dimensional `torch.Tensor`
|
||||
groups (int, optional): number of groups in the conv layer (default: 1)
|
||||
Examples:
|
||||
>>> w = torch.empty(3, 16, 5, 5)
|
||||
>>> nn.init.dirac_(w)
|
||||
>>> w = torch.empty(3, 24, 5, 5)
|
||||
>>> nn.init.dirac_(w, 3)
|
||||
"""
|
||||
dimensions = tensor.ndimension()
|
||||
if dimensions not in [3, 4, 5]:
|
||||
raise ValueError("Only tensors with 3, 4, or 5 dimensions are supported")
|
||||
|
||||
sizes = tensor.size()
|
||||
|
||||
if sizes[0] % groups != 0:
|
||||
raise ValueError("dim 0 must be divisible by groups")
|
||||
|
||||
if tensor.is_meta:
|
||||
return tensor
|
||||
|
||||
out_chans_per_grp = sizes[0] // groups
|
||||
min_dim = min(out_chans_per_grp, sizes[1])
|
||||
|
||||
with torch.no_grad():
|
||||
tensor.zero_()
|
||||
|
||||
for g in range(groups):
|
||||
for d in range(min_dim):
|
||||
if dimensions == 3: # Temporal convolution
|
||||
tensor[g * out_chans_per_grp + d, d, tensor.size(2) // 2] = 1
|
||||
elif dimensions == 4: # Spatial convolution
|
||||
tensor[
|
||||
g * out_chans_per_grp + d,
|
||||
d,
|
||||
tensor.size(2) // 2,
|
||||
tensor.size(3) // 2,
|
||||
] = 1
|
||||
else: # Volumetric convolution
|
||||
tensor[
|
||||
g * out_chans_per_grp + d,
|
||||
d,
|
||||
tensor.size(2) // 2,
|
||||
tensor.size(3) // 2,
|
||||
tensor.size(4) // 2,
|
||||
] = 1
|
||||
return tensor
|
||||
|
||||
|
||||
def _calculate_fan_in_and_fan_out(tensor: Tensor) -> tuple[int, int]:
|
||||
dimensions = tensor.dim()
|
||||
if dimensions < 2:
|
||||
raise ValueError(
|
||||
"Fan in and fan out can not be computed for tensor with fewer than 2 dimensions"
|
||||
)
|
||||
|
||||
num_input_fmaps = tensor.size(1)
|
||||
num_output_fmaps = tensor.size(0)
|
||||
receptive_field_size = 1
|
||||
if tensor.dim() > 2:
|
||||
# math.prod is not always available, accumulate the product manually
|
||||
# we could use functools.reduce but that is not supported by TorchScript
|
||||
for s in tensor.shape[2:]:
|
||||
receptive_field_size *= s
|
||||
fan_in = num_input_fmaps * receptive_field_size
|
||||
fan_out = num_output_fmaps * receptive_field_size
|
||||
|
||||
return fan_in, fan_out
|
||||
|
||||
|
||||
def xavier_uniform_(
|
||||
tensor: Tensor,
|
||||
gain: float = 1.0,
|
||||
generator: torch.Generator | None = None,
|
||||
) -> Tensor:
|
||||
r"""Fill the input `Tensor` with values using a Xavier uniform distribution.
|
||||
|
||||
The method is described in `Understanding the difficulty of training
|
||||
deep feedforward neural networks` - Glorot, X. & Bengio, Y. (2010).
|
||||
The resulting tensor will have values sampled from
|
||||
:math:`\mathcal{U}(-a, a)` where
|
||||
|
||||
.. math::
|
||||
a = \text{gain} \times \sqrt{\frac{6}{\text{fan\_in} + \text{fan\_out}}}
|
||||
|
||||
Also known as Glorot initialization.
|
||||
|
||||
Args:
|
||||
tensor: an n-dimensional `torch.Tensor`
|
||||
gain: an optional scaling factor
|
||||
generator: the torch Generator to sample from (default: None)
|
||||
|
||||
Examples:
|
||||
>>> w = torch.empty(3, 5)
|
||||
>>> nn.init.xavier_uniform_(w, gain=nn.init.calculate_gain("relu"))
|
||||
"""
|
||||
fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
|
||||
std = gain * math.sqrt(2.0 / float(fan_in + fan_out))
|
||||
a = math.sqrt(3.0) * std # Calculate uniform bounds from standard deviation
|
||||
|
||||
return _no_grad_uniform_(tensor, -a, a, generator)
|
||||
|
||||
|
||||
def xavier_normal_(
|
||||
tensor: Tensor,
|
||||
gain: float = 1.0,
|
||||
generator: torch.Generator | None = None,
|
||||
) -> Tensor:
|
||||
r"""Fill the input `Tensor` with values using a Xavier normal distribution.
|
||||
|
||||
The method is described in `Understanding the difficulty of training deep feedforward
|
||||
neural networks` - Glorot, X. & Bengio, Y. (2010). The resulting tensor
|
||||
will have values sampled from :math:`\mathcal{N}(0, \text{std}^2)` where
|
||||
|
||||
.. math::
|
||||
\text{std} = \text{gain} \times \sqrt{\frac{2}{\text{fan\_in} + \text{fan\_out}}}
|
||||
|
||||
Also known as Glorot initialization.
|
||||
|
||||
Args:
|
||||
tensor: an n-dimensional `torch.Tensor`
|
||||
gain: an optional scaling factor
|
||||
generator: the torch Generator to sample from (default: None)
|
||||
|
||||
Examples:
|
||||
>>> w = torch.empty(3, 5)
|
||||
>>> nn.init.xavier_normal_(w)
|
||||
"""
|
||||
fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
|
||||
std = gain * math.sqrt(2.0 / float(fan_in + fan_out))
|
||||
|
||||
return _no_grad_normal_(tensor, 0.0, std, generator)
|
||||
|
||||
|
||||
def _calculate_correct_fan(tensor: Tensor, mode: _FanMode) -> int:
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
mode = mode.lower()
|
||||
valid_modes = ["fan_in", "fan_out"]
|
||||
if mode not in valid_modes:
|
||||
raise ValueError(f"Mode {mode} not supported, please use one of {valid_modes}")
|
||||
|
||||
fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
|
||||
return fan_in if mode == "fan_in" else fan_out
|
||||
|
||||
|
||||
def kaiming_uniform_(
|
||||
tensor: Tensor,
|
||||
a: float = 0,
|
||||
mode: _FanMode = "fan_in",
|
||||
nonlinearity: _NonlinearityType = "leaky_relu",
|
||||
generator: torch.Generator | None = None,
|
||||
) -> Tensor:
|
||||
r"""Fill the input `Tensor` with values using a Kaiming uniform distribution.
|
||||
|
||||
The method is described in `Delving deep into rectifiers: Surpassing
|
||||
human-level performance on ImageNet classification` - He, K. et al. (2015).
|
||||
The resulting tensor will have values sampled from
|
||||
:math:`\mathcal{U}(-\text{bound}, \text{bound})` where
|
||||
|
||||
.. math::
|
||||
\text{bound} = \text{gain} \times \sqrt{\frac{3}{\text{fan\_mode}}}
|
||||
|
||||
Also known as He initialization.
|
||||
|
||||
Args:
|
||||
tensor: an n-dimensional `torch.Tensor`
|
||||
a: the negative slope of the rectifier used after this layer (only
|
||||
used with ``'leaky_relu'``)
|
||||
mode: either ``'fan_in'`` (default) or ``'fan_out'``. Choosing ``'fan_in'``
|
||||
preserves the magnitude of the variance of the weights in the
|
||||
forward pass. Choosing ``'fan_out'`` preserves the magnitudes in the
|
||||
backwards pass.
|
||||
nonlinearity: the non-linear function (`nn.functional` name),
|
||||
recommended to use only with ``'relu'`` or ``'leaky_relu'`` (default).
|
||||
generator: the torch Generator to sample from (default: None)
|
||||
|
||||
Examples:
|
||||
>>> w = torch.empty(3, 5)
|
||||
>>> nn.init.kaiming_uniform_(w, mode="fan_in", nonlinearity="relu")
|
||||
|
||||
Note:
|
||||
Be aware that ``fan_in`` and ``fan_out`` are calculated assuming
|
||||
that the weight matrix is used in a transposed manner,
|
||||
(i.e., ``x @ w.T`` in ``Linear`` layers, where ``w.shape = [fan_out, fan_in]``).
|
||||
This is important for correct initialization.
|
||||
If you plan to use ``x @ w``, where ``w.shape = [fan_in, fan_out]``,
|
||||
pass in a transposed weight matrix, i.e. ``nn.init.kaiming_uniform_(w.T, ...)``.
|
||||
"""
|
||||
if torch.overrides.has_torch_function_variadic(tensor):
|
||||
return torch.overrides.handle_torch_function(
|
||||
kaiming_uniform_,
|
||||
(tensor,),
|
||||
tensor=tensor,
|
||||
a=a,
|
||||
mode=mode,
|
||||
nonlinearity=nonlinearity,
|
||||
generator=generator,
|
||||
)
|
||||
|
||||
if 0 in tensor.shape:
|
||||
warnings.warn("Initializing zero-element tensors is a no-op", stacklevel=2)
|
||||
return tensor
|
||||
fan = _calculate_correct_fan(tensor, mode)
|
||||
gain = calculate_gain(nonlinearity, a)
|
||||
std = gain / math.sqrt(fan)
|
||||
bound = math.sqrt(3.0) * std # Calculate uniform bounds from standard deviation
|
||||
with torch.no_grad():
|
||||
return tensor.uniform_(-bound, bound, generator=generator)
|
||||
|
||||
|
||||
def kaiming_normal_(
|
||||
tensor: Tensor,
|
||||
a: float = 0,
|
||||
mode: _FanMode = "fan_in",
|
||||
nonlinearity: _NonlinearityType = "leaky_relu",
|
||||
generator: torch.Generator | None = None,
|
||||
) -> Tensor:
|
||||
r"""Fill the input `Tensor` with values using a Kaiming normal distribution.
|
||||
|
||||
The method is described in `Delving deep into rectifiers: Surpassing
|
||||
human-level performance on ImageNet classification` - He, K. et al. (2015).
|
||||
The resulting tensor will have values sampled from
|
||||
:math:`\mathcal{N}(0, \text{std}^2)` where
|
||||
|
||||
.. math::
|
||||
\text{std} = \frac{\text{gain}}{\sqrt{\text{fan\_mode}}}
|
||||
|
||||
Also known as He initialization.
|
||||
|
||||
Args:
|
||||
tensor: an n-dimensional `torch.Tensor`
|
||||
a: the negative slope of the rectifier used after this layer (only
|
||||
used with ``'leaky_relu'``)
|
||||
mode: either ``'fan_in'`` (default) or ``'fan_out'``. Choosing ``'fan_in'``
|
||||
preserves the magnitude of the variance of the weights in the
|
||||
forward pass. Choosing ``'fan_out'`` preserves the magnitudes in the
|
||||
backwards pass.
|
||||
nonlinearity: the non-linear function (`nn.functional` name),
|
||||
recommended to use only with ``'relu'`` or ``'leaky_relu'`` (default).
|
||||
generator: the torch Generator to sample from (default: None)
|
||||
|
||||
Examples:
|
||||
>>> w = torch.empty(3, 5)
|
||||
>>> nn.init.kaiming_normal_(w, mode="fan_out", nonlinearity="relu")
|
||||
|
||||
Note:
|
||||
Be aware that ``fan_in`` and ``fan_out`` are calculated assuming
|
||||
that the weight matrix is used in a transposed manner,
|
||||
(i.e., ``x @ w.T`` in ``Linear`` layers, where ``w.shape = [fan_out, fan_in]``).
|
||||
This is important for correct initialization.
|
||||
If you plan to use ``x @ w``, where ``w.shape = [fan_in, fan_out]``,
|
||||
pass in a transposed weight matrix, i.e. ``nn.init.kaiming_normal_(w.T, ...)``.
|
||||
"""
|
||||
if 0 in tensor.shape:
|
||||
warnings.warn("Initializing zero-element tensors is a no-op", stacklevel=2)
|
||||
return tensor
|
||||
fan = _calculate_correct_fan(tensor, mode)
|
||||
gain = calculate_gain(nonlinearity, a)
|
||||
std = gain / math.sqrt(fan)
|
||||
with torch.no_grad():
|
||||
return tensor.normal_(0, std, generator=generator)
|
||||
|
||||
|
||||
def orthogonal_(
|
||||
tensor: Tensor,
|
||||
gain: float = 1,
|
||||
generator: torch.Generator | None = None,
|
||||
) -> Tensor:
|
||||
r"""Fill the input `Tensor` with a (semi) orthogonal matrix.
|
||||
|
||||
Described in `Exact solutions to the nonlinear dynamics of learning in deep
|
||||
linear neural networks` - Saxe, A. et al. (2013). The input tensor must have
|
||||
at least 2 dimensions, and for tensors with more than 2 dimensions the
|
||||
trailing dimensions are flattened.
|
||||
|
||||
Args:
|
||||
tensor: an n-dimensional `torch.Tensor`, where :math:`n \geq 2`
|
||||
gain: optional scaling factor
|
||||
generator: the torch Generator to sample from (default: None)
|
||||
|
||||
Examples:
|
||||
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_LAPACK)
|
||||
>>> w = torch.empty(3, 5)
|
||||
>>> nn.init.orthogonal_(w)
|
||||
"""
|
||||
if tensor.ndimension() < 2:
|
||||
raise ValueError("Only tensors with 2 or more dimensions are supported")
|
||||
|
||||
if tensor.numel() == 0 or tensor.is_meta:
|
||||
# no-op
|
||||
return tensor
|
||||
rows = tensor.size(0)
|
||||
cols = tensor.numel() // rows
|
||||
flattened = tensor.new_empty((rows, cols)).normal_(0, 1, generator=generator)
|
||||
|
||||
if rows < cols:
|
||||
flattened.t_()
|
||||
|
||||
# Compute the qr factorization
|
||||
q, r = torch.linalg.qr(flattened)
|
||||
# Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
|
||||
d = torch.diag(r, 0)
|
||||
ph = d.sign()
|
||||
q *= ph
|
||||
|
||||
if rows < cols:
|
||||
q.t_()
|
||||
|
||||
with torch.no_grad():
|
||||
tensor.view_as(q).copy_(q)
|
||||
tensor.mul_(gain)
|
||||
return tensor
|
||||
|
||||
|
||||
def sparse_(
|
||||
tensor: Tensor,
|
||||
sparsity: float,
|
||||
std: float = 0.01,
|
||||
generator: torch.Generator | None = None,
|
||||
) -> Tensor:
|
||||
r"""Fill the 2D input `Tensor` as a sparse matrix.
|
||||
|
||||
The non-zero elements will be drawn from the normal distribution
|
||||
:math:`\mathcal{N}(0, 0.01)`, as described in `Deep learning via
|
||||
Hessian-free optimization` - Martens, J. (2010).
|
||||
|
||||
Args:
|
||||
tensor: an n-dimensional `torch.Tensor`
|
||||
sparsity: The fraction of elements in each column to be set to zero
|
||||
std: the standard deviation of the normal distribution used to generate
|
||||
the non-zero values
|
||||
generator: the torch Generator to sample from (default: None)
|
||||
|
||||
Examples:
|
||||
>>> w = torch.empty(3, 5)
|
||||
>>> nn.init.sparse_(w, sparsity=0.1)
|
||||
"""
|
||||
if tensor.ndimension() != 2:
|
||||
raise ValueError("Only tensors with 2 dimensions are supported")
|
||||
|
||||
if tensor.is_meta:
|
||||
return tensor
|
||||
|
||||
rows, cols = tensor.shape
|
||||
num_zeros = math.ceil(sparsity * rows)
|
||||
|
||||
with torch.no_grad():
|
||||
tensor.normal_(0, std, generator=generator)
|
||||
for col_idx in range(cols):
|
||||
row_indices = torch.randperm(rows)
|
||||
zero_indices = row_indices[:num_zeros]
|
||||
tensor[zero_indices, col_idx] = 0
|
||||
return tensor
|
||||
|
||||
|
||||
# for backward compatibility
|
||||
def _make_deprecate(meth: Callable[_P, _R]) -> Callable[_P, _R]:
|
||||
new_name = meth.__name__
|
||||
old_name = new_name[:-1]
|
||||
|
||||
def deprecated_init(*args: _P.args, **kwargs: _P.kwargs) -> _R:
|
||||
warnings.warn(
|
||||
f"`nn.init.{old_name}` is now deprecated in favor of `nn.init.{new_name}`.",
|
||||
FutureWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
return meth(*args, **kwargs)
|
||||
|
||||
deprecated_init.__doc__ = rf"""
|
||||
{old_name}(...)
|
||||
|
||||
.. warning::
|
||||
This method is now deprecated in favor of :func:`torch.nn.init.{new_name}`.
|
||||
|
||||
See :func:`~torch.nn.init.{new_name}` for details."""
|
||||
deprecated_init.__name__ = old_name
|
||||
return deprecated_init
|
||||
|
||||
|
||||
uniform = _make_deprecate(uniform_)
|
||||
normal = _make_deprecate(normal_)
|
||||
constant = _make_deprecate(constant_)
|
||||
eye = _make_deprecate(eye_)
|
||||
dirac = _make_deprecate(dirac_)
|
||||
xavier_uniform = _make_deprecate(xavier_uniform_)
|
||||
xavier_normal = _make_deprecate(xavier_normal_)
|
||||
kaiming_uniform = _make_deprecate(kaiming_uniform_)
|
||||
kaiming_normal = _make_deprecate(kaiming_normal_)
|
||||
orthogonal = _make_deprecate(orthogonal_)
|
||||
sparse = _make_deprecate(sparse_)
|
||||
@@ -0,0 +1,36 @@
|
||||
from torch.ao.nn.intrinsic import (
|
||||
BNReLU2d,
|
||||
BNReLU3d,
|
||||
ConvBn1d,
|
||||
ConvBn2d,
|
||||
ConvBn3d,
|
||||
ConvBnReLU1d,
|
||||
ConvBnReLU2d,
|
||||
ConvBnReLU3d,
|
||||
ConvReLU1d,
|
||||
ConvReLU2d,
|
||||
ConvReLU3d,
|
||||
LinearBn1d,
|
||||
LinearReLU,
|
||||
)
|
||||
from torch.ao.nn.intrinsic.modules.fused import _FusedModule # noqa: F401
|
||||
|
||||
# Include the subpackages in case user imports from it directly
|
||||
from torch.nn.intrinsic import modules, qat, quantized # noqa: F401
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ConvBn1d",
|
||||
"ConvBn2d",
|
||||
"ConvBn3d",
|
||||
"ConvBnReLU1d",
|
||||
"ConvBnReLU2d",
|
||||
"ConvBnReLU3d",
|
||||
"ConvReLU1d",
|
||||
"ConvReLU2d",
|
||||
"ConvReLU3d",
|
||||
"LinearReLU",
|
||||
"BNReLU2d",
|
||||
"BNReLU3d",
|
||||
"LinearBn1d",
|
||||
]
|
||||
@@ -0,0 +1,33 @@
|
||||
from torch.nn.intrinsic.modules.fused import (
|
||||
_FusedModule,
|
||||
BNReLU2d,
|
||||
BNReLU3d,
|
||||
ConvBn1d,
|
||||
ConvBn2d,
|
||||
ConvBn3d,
|
||||
ConvBnReLU1d,
|
||||
ConvBnReLU2d,
|
||||
ConvBnReLU3d,
|
||||
ConvReLU1d,
|
||||
ConvReLU2d,
|
||||
ConvReLU3d,
|
||||
LinearBn1d,
|
||||
LinearReLU,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BNReLU2d",
|
||||
"BNReLU3d",
|
||||
"ConvBn1d",
|
||||
"ConvBn2d",
|
||||
"ConvBn3d",
|
||||
"ConvBnReLU1d",
|
||||
"ConvBnReLU2d",
|
||||
"ConvBnReLU3d",
|
||||
"ConvReLU1d",
|
||||
"ConvReLU2d",
|
||||
"ConvReLU3d",
|
||||
"LinearBn1d",
|
||||
"LinearReLU",
|
||||
]
|
||||
@@ -0,0 +1,33 @@
|
||||
from torch.ao.nn.intrinsic import (
|
||||
BNReLU2d,
|
||||
BNReLU3d,
|
||||
ConvBn1d,
|
||||
ConvBn2d,
|
||||
ConvBn3d,
|
||||
ConvBnReLU1d,
|
||||
ConvBnReLU2d,
|
||||
ConvBnReLU3d,
|
||||
ConvReLU1d,
|
||||
ConvReLU2d,
|
||||
ConvReLU3d,
|
||||
LinearBn1d,
|
||||
LinearReLU,
|
||||
)
|
||||
from torch.ao.nn.intrinsic.modules.fused import _FusedModule # noqa: F401
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BNReLU2d",
|
||||
"BNReLU3d",
|
||||
"ConvBn1d",
|
||||
"ConvBn2d",
|
||||
"ConvBn3d",
|
||||
"ConvBnReLU1d",
|
||||
"ConvBnReLU2d",
|
||||
"ConvBnReLU3d",
|
||||
"ConvReLU1d",
|
||||
"ConvReLU2d",
|
||||
"ConvReLU3d",
|
||||
"LinearBn1d",
|
||||
"LinearReLU",
|
||||
]
|
||||
@@ -0,0 +1 @@
|
||||
from torch.nn.intrinsic.qat.modules import * # noqa: F403
|
||||
@@ -0,0 +1,32 @@
|
||||
from torch.nn.intrinsic.qat.modules.conv_fused import (
|
||||
ConvBn1d,
|
||||
ConvBn2d,
|
||||
ConvBn3d,
|
||||
ConvBnReLU1d,
|
||||
ConvBnReLU2d,
|
||||
ConvBnReLU3d,
|
||||
ConvReLU1d,
|
||||
ConvReLU2d,
|
||||
ConvReLU3d,
|
||||
freeze_bn_stats,
|
||||
update_bn_stats,
|
||||
)
|
||||
from torch.nn.intrinsic.qat.modules.linear_fused import LinearBn1d
|
||||
from torch.nn.intrinsic.qat.modules.linear_relu import LinearReLU
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LinearReLU",
|
||||
"LinearBn1d",
|
||||
"ConvReLU1d",
|
||||
"ConvReLU2d",
|
||||
"ConvReLU3d",
|
||||
"ConvBn1d",
|
||||
"ConvBn2d",
|
||||
"ConvBn3d",
|
||||
"ConvBnReLU1d",
|
||||
"ConvBnReLU2d",
|
||||
"ConvBnReLU3d",
|
||||
"update_bn_stats",
|
||||
"freeze_bn_stats",
|
||||
]
|
||||
@@ -0,0 +1,39 @@
|
||||
r"""Intrinsic QAT Modules.
|
||||
|
||||
This file is in the process of migration to `torch/ao/nn/intrinsic/qat`, and
|
||||
is kept here for compatibility while the migration process is ongoing.
|
||||
If you are adding a new entry/functionality, please, add it to the
|
||||
appropriate file under the `torch/ao/nn/intrinsic/qat/modules`,
|
||||
while adding an import statement here.
|
||||
"""
|
||||
|
||||
from torch.ao.nn.intrinsic.qat import (
|
||||
ConvBn1d,
|
||||
ConvBn2d,
|
||||
ConvBn3d,
|
||||
ConvBnReLU1d,
|
||||
ConvBnReLU2d,
|
||||
ConvBnReLU3d,
|
||||
ConvReLU1d,
|
||||
ConvReLU2d,
|
||||
ConvReLU3d,
|
||||
freeze_bn_stats,
|
||||
update_bn_stats,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
# Modules
|
||||
"ConvBn1d",
|
||||
"ConvBnReLU1d",
|
||||
"ConvReLU1d",
|
||||
"ConvBn2d",
|
||||
"ConvBnReLU2d",
|
||||
"ConvReLU2d",
|
||||
"ConvBn3d",
|
||||
"ConvBnReLU3d",
|
||||
"ConvReLU3d",
|
||||
# Utilities
|
||||
"freeze_bn_stats",
|
||||
"update_bn_stats",
|
||||
]
|
||||
+15
@@ -0,0 +1,15 @@
|
||||
r"""Intrinsic QAT Modules.
|
||||
|
||||
This file is in the process of migration to `torch/ao/nn/intrinsic/qat`, and
|
||||
is kept here for compatibility while the migration process is ongoing.
|
||||
If you are adding a new entry/functionality, please, add it to the
|
||||
appropriate file under the `torch/ao/nn/intrinsic/qat/modules`,
|
||||
while adding an import statement here.
|
||||
"""
|
||||
|
||||
from torch.ao.nn.intrinsic.qat import LinearBn1d
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LinearBn1d",
|
||||
]
|
||||
+15
@@ -0,0 +1,15 @@
|
||||
r"""Intrinsic QAT Modules.
|
||||
|
||||
This file is in the process of migration to `torch/ao/nn/intrinsic/qat`, and
|
||||
is kept here for compatibility while the migration process is ongoing.
|
||||
If you are adding a new entry/functionality, please, add it to the
|
||||
appropriate file under the `torch/ao/nn/intrinsic/qat/modules`,
|
||||
while adding an import statement here.
|
||||
"""
|
||||
|
||||
from torch.ao.nn.intrinsic.qat import LinearReLU
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LinearReLU",
|
||||
]
|
||||
@@ -0,0 +1,14 @@
|
||||
# to ensure customers can use the module below
|
||||
# without importing it directly
|
||||
from torch.nn.intrinsic.quantized import dynamic, modules # noqa: F401
|
||||
from torch.nn.intrinsic.quantized.modules import * # noqa: F403
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BNReLU2d",
|
||||
"BNReLU3d",
|
||||
"ConvReLU1d",
|
||||
"ConvReLU2d",
|
||||
"ConvReLU3d",
|
||||
"LinearReLU",
|
||||
]
|
||||
+1
@@ -0,0 +1 @@
|
||||
from torch.nn.intrinsic.quantized.dynamic.modules import * # noqa: F403
|
||||
+6
@@ -0,0 +1,6 @@
|
||||
from torch.nn.intrinsic.quantized.dynamic.modules.linear_relu import LinearReLU
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LinearReLU",
|
||||
]
|
||||
+6
@@ -0,0 +1,6 @@
|
||||
from torch.ao.nn.intrinsic.quantized.dynamic import LinearReLU
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LinearReLU",
|
||||
]
|
||||
+17
@@ -0,0 +1,17 @@
|
||||
from torch.nn.intrinsic.quantized.modules.bn_relu import BNReLU2d, BNReLU3d
|
||||
from torch.nn.intrinsic.quantized.modules.conv_relu import (
|
||||
ConvReLU1d,
|
||||
ConvReLU2d,
|
||||
ConvReLU3d,
|
||||
)
|
||||
from torch.nn.intrinsic.quantized.modules.linear_relu import LinearReLU
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LinearReLU",
|
||||
"ConvReLU1d",
|
||||
"ConvReLU2d",
|
||||
"ConvReLU3d",
|
||||
"BNReLU2d",
|
||||
"BNReLU3d",
|
||||
]
|
||||
+7
@@ -0,0 +1,7 @@
|
||||
from torch.ao.nn.intrinsic.quantized import BNReLU2d, BNReLU3d
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BNReLU2d",
|
||||
"BNReLU3d",
|
||||
]
|
||||
+8
@@ -0,0 +1,8 @@
|
||||
from torch.ao.nn.intrinsic.quantized import ConvReLU1d, ConvReLU2d, ConvReLU3d
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ConvReLU1d",
|
||||
"ConvReLU2d",
|
||||
"ConvReLU3d",
|
||||
]
|
||||
+6
@@ -0,0 +1,6 @@
|
||||
from torch.ao.nn.intrinsic.quantized import LinearReLU
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LinearReLU",
|
||||
]
|
||||
@@ -0,0 +1,335 @@
|
||||
from .module import Module # usort: skip
|
||||
from .linear import Bilinear, Identity, LazyLinear, Linear # usort: skip
|
||||
from .activation import (
|
||||
CELU,
|
||||
ELU,
|
||||
GELU,
|
||||
GLU,
|
||||
Hardshrink,
|
||||
Hardsigmoid,
|
||||
Hardswish,
|
||||
Hardtanh,
|
||||
LeakyReLU,
|
||||
LogSigmoid,
|
||||
LogSoftmax,
|
||||
Mish,
|
||||
MultiheadAttention,
|
||||
PReLU,
|
||||
ReLU,
|
||||
ReLU6,
|
||||
RReLU,
|
||||
SELU,
|
||||
Sigmoid,
|
||||
SiLU,
|
||||
Softmax,
|
||||
Softmax2d,
|
||||
Softmin,
|
||||
Softplus,
|
||||
Softshrink,
|
||||
Softsign,
|
||||
Tanh,
|
||||
Tanhshrink,
|
||||
Threshold,
|
||||
)
|
||||
from .adaptive import AdaptiveLogSoftmaxWithLoss
|
||||
from .batchnorm import (
|
||||
BatchNorm1d,
|
||||
BatchNorm2d,
|
||||
BatchNorm3d,
|
||||
LazyBatchNorm1d,
|
||||
LazyBatchNorm2d,
|
||||
LazyBatchNorm3d,
|
||||
SyncBatchNorm,
|
||||
)
|
||||
from .channelshuffle import ChannelShuffle
|
||||
from .container import (
|
||||
Container,
|
||||
ModuleDict,
|
||||
ModuleList,
|
||||
ParameterDict,
|
||||
ParameterList,
|
||||
Sequential,
|
||||
)
|
||||
from .conv import (
|
||||
Conv1d,
|
||||
Conv2d,
|
||||
Conv3d,
|
||||
ConvTranspose1d,
|
||||
ConvTranspose2d,
|
||||
ConvTranspose3d,
|
||||
LazyConv1d,
|
||||
LazyConv2d,
|
||||
LazyConv3d,
|
||||
LazyConvTranspose1d,
|
||||
LazyConvTranspose2d,
|
||||
LazyConvTranspose3d,
|
||||
)
|
||||
from .distance import CosineSimilarity, PairwiseDistance
|
||||
from .dropout import (
|
||||
AlphaDropout,
|
||||
Dropout,
|
||||
Dropout1d,
|
||||
Dropout2d,
|
||||
Dropout3d,
|
||||
FeatureAlphaDropout,
|
||||
)
|
||||
from .flatten import Flatten, Unflatten
|
||||
from .fold import Fold, Unfold
|
||||
from .instancenorm import (
|
||||
InstanceNorm1d,
|
||||
InstanceNorm2d,
|
||||
InstanceNorm3d,
|
||||
LazyInstanceNorm1d,
|
||||
LazyInstanceNorm2d,
|
||||
LazyInstanceNorm3d,
|
||||
)
|
||||
from .loss import (
|
||||
BCELoss,
|
||||
BCEWithLogitsLoss,
|
||||
CosineEmbeddingLoss,
|
||||
CrossEntropyLoss,
|
||||
CTCLoss,
|
||||
GaussianNLLLoss,
|
||||
HingeEmbeddingLoss,
|
||||
HuberLoss,
|
||||
KLDivLoss,
|
||||
L1Loss,
|
||||
MarginRankingLoss,
|
||||
MSELoss,
|
||||
MultiLabelMarginLoss,
|
||||
MultiLabelSoftMarginLoss,
|
||||
MultiMarginLoss,
|
||||
NLLLoss,
|
||||
NLLLoss2d,
|
||||
PoissonNLLLoss,
|
||||
SmoothL1Loss,
|
||||
SoftMarginLoss,
|
||||
TripletMarginLoss,
|
||||
TripletMarginWithDistanceLoss,
|
||||
)
|
||||
from .normalization import (
|
||||
CrossMapLRN2d,
|
||||
GroupNorm,
|
||||
LayerNorm,
|
||||
LocalResponseNorm,
|
||||
RMSNorm,
|
||||
)
|
||||
from .padding import (
|
||||
CircularPad1d,
|
||||
CircularPad2d,
|
||||
CircularPad3d,
|
||||
ConstantPad1d,
|
||||
ConstantPad2d,
|
||||
ConstantPad3d,
|
||||
ReflectionPad1d,
|
||||
ReflectionPad2d,
|
||||
ReflectionPad3d,
|
||||
ReplicationPad1d,
|
||||
ReplicationPad2d,
|
||||
ReplicationPad3d,
|
||||
ZeroPad1d,
|
||||
ZeroPad2d,
|
||||
ZeroPad3d,
|
||||
)
|
||||
from .pixelshuffle import PixelShuffle, PixelUnshuffle
|
||||
from .pooling import (
|
||||
AdaptiveAvgPool1d,
|
||||
AdaptiveAvgPool2d,
|
||||
AdaptiveAvgPool3d,
|
||||
AdaptiveMaxPool1d,
|
||||
AdaptiveMaxPool2d,
|
||||
AdaptiveMaxPool3d,
|
||||
AvgPool1d,
|
||||
AvgPool2d,
|
||||
AvgPool3d,
|
||||
FractionalMaxPool2d,
|
||||
FractionalMaxPool3d,
|
||||
LPPool1d,
|
||||
LPPool2d,
|
||||
LPPool3d,
|
||||
MaxPool1d,
|
||||
MaxPool2d,
|
||||
MaxPool3d,
|
||||
MaxUnpool1d,
|
||||
MaxUnpool2d,
|
||||
MaxUnpool3d,
|
||||
)
|
||||
from .rnn import GRU, GRUCell, LSTM, LSTMCell, RNN, RNNBase, RNNCell, RNNCellBase
|
||||
from .sparse import Embedding, EmbeddingBag
|
||||
from .transformer import (
|
||||
Transformer,
|
||||
TransformerDecoder,
|
||||
TransformerDecoderLayer,
|
||||
TransformerEncoder,
|
||||
TransformerEncoderLayer,
|
||||
)
|
||||
from .upsampling import Upsample, UpsamplingBilinear2d, UpsamplingNearest2d
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AdaptiveAvgPool1d",
|
||||
"AdaptiveAvgPool2d",
|
||||
"AdaptiveAvgPool3d",
|
||||
"AdaptiveLogSoftmaxWithLoss",
|
||||
"AdaptiveMaxPool1d",
|
||||
"AdaptiveMaxPool2d",
|
||||
"AdaptiveMaxPool3d",
|
||||
"AlphaDropout",
|
||||
"AvgPool1d",
|
||||
"AvgPool2d",
|
||||
"AvgPool3d",
|
||||
"BCELoss",
|
||||
"BCEWithLogitsLoss",
|
||||
"BatchNorm1d",
|
||||
"BatchNorm2d",
|
||||
"BatchNorm3d",
|
||||
"Bilinear",
|
||||
"CELU",
|
||||
"CTCLoss",
|
||||
"ChannelShuffle",
|
||||
"CircularPad1d",
|
||||
"CircularPad2d",
|
||||
"CircularPad3d",
|
||||
"ConstantPad1d",
|
||||
"ConstantPad2d",
|
||||
"ConstantPad3d",
|
||||
"Container",
|
||||
"Conv1d",
|
||||
"Conv2d",
|
||||
"Conv3d",
|
||||
"ConvTranspose1d",
|
||||
"ConvTranspose2d",
|
||||
"ConvTranspose3d",
|
||||
"CosineEmbeddingLoss",
|
||||
"CosineSimilarity",
|
||||
"CrossEntropyLoss",
|
||||
"CrossMapLRN2d",
|
||||
"Dropout",
|
||||
"Dropout1d",
|
||||
"Dropout2d",
|
||||
"Dropout3d",
|
||||
"ELU",
|
||||
"Embedding",
|
||||
"EmbeddingBag",
|
||||
"FeatureAlphaDropout",
|
||||
"Flatten",
|
||||
"Fold",
|
||||
"FractionalMaxPool2d",
|
||||
"FractionalMaxPool3d",
|
||||
"GELU",
|
||||
"GLU",
|
||||
"GRU",
|
||||
"GRUCell",
|
||||
"GaussianNLLLoss",
|
||||
"GroupNorm",
|
||||
"Hardshrink",
|
||||
"Hardsigmoid",
|
||||
"Hardswish",
|
||||
"Hardtanh",
|
||||
"HingeEmbeddingLoss",
|
||||
"HuberLoss",
|
||||
"Identity",
|
||||
"InstanceNorm1d",
|
||||
"InstanceNorm2d",
|
||||
"InstanceNorm3d",
|
||||
"KLDivLoss",
|
||||
"L1Loss",
|
||||
"LPPool1d",
|
||||
"LPPool2d",
|
||||
"LPPool3d",
|
||||
"LSTM",
|
||||
"LSTMCell",
|
||||
"LayerNorm",
|
||||
"LazyBatchNorm1d",
|
||||
"LazyBatchNorm2d",
|
||||
"LazyBatchNorm3d",
|
||||
"LazyConv1d",
|
||||
"LazyConv2d",
|
||||
"LazyConv3d",
|
||||
"LazyConvTranspose1d",
|
||||
"LazyConvTranspose2d",
|
||||
"LazyConvTranspose3d",
|
||||
"LazyInstanceNorm1d",
|
||||
"LazyInstanceNorm2d",
|
||||
"LazyInstanceNorm3d",
|
||||
"LazyLinear",
|
||||
"LeakyReLU",
|
||||
"Linear",
|
||||
"LocalResponseNorm",
|
||||
"LogSigmoid",
|
||||
"LogSoftmax",
|
||||
"MSELoss",
|
||||
"MarginRankingLoss",
|
||||
"MaxPool1d",
|
||||
"MaxPool2d",
|
||||
"MaxPool3d",
|
||||
"MaxUnpool1d",
|
||||
"MaxUnpool2d",
|
||||
"MaxUnpool3d",
|
||||
"Mish",
|
||||
"Module",
|
||||
"ModuleDict",
|
||||
"ModuleList",
|
||||
"MultiLabelMarginLoss",
|
||||
"MultiLabelSoftMarginLoss",
|
||||
"MultiMarginLoss",
|
||||
"MultiheadAttention",
|
||||
"NLLLoss",
|
||||
"NLLLoss2d",
|
||||
"PReLU",
|
||||
"PairwiseDistance",
|
||||
"ParameterDict",
|
||||
"ParameterList",
|
||||
"PixelShuffle",
|
||||
"PixelUnshuffle",
|
||||
"PoissonNLLLoss",
|
||||
"RMSNorm",
|
||||
"RNN",
|
||||
"RNNBase",
|
||||
"RNNCell",
|
||||
"RNNCellBase",
|
||||
"RReLU",
|
||||
"ReLU",
|
||||
"ReLU6",
|
||||
"ReflectionPad1d",
|
||||
"ReflectionPad2d",
|
||||
"ReflectionPad3d",
|
||||
"ReplicationPad1d",
|
||||
"ReplicationPad2d",
|
||||
"ReplicationPad3d",
|
||||
"SELU",
|
||||
"Sequential",
|
||||
"SiLU",
|
||||
"Sigmoid",
|
||||
"SmoothL1Loss",
|
||||
"SoftMarginLoss",
|
||||
"Softmax",
|
||||
"Softmax2d",
|
||||
"Softmin",
|
||||
"Softplus",
|
||||
"Softshrink",
|
||||
"Softsign",
|
||||
"SyncBatchNorm",
|
||||
"Tanh",
|
||||
"Tanhshrink",
|
||||
"Threshold",
|
||||
"Transformer",
|
||||
"TransformerDecoder",
|
||||
"TransformerDecoderLayer",
|
||||
"TransformerEncoder",
|
||||
"TransformerEncoderLayer",
|
||||
"TripletMarginLoss",
|
||||
"TripletMarginWithDistanceLoss",
|
||||
"Unflatten",
|
||||
"Unfold",
|
||||
"Upsample",
|
||||
"UpsamplingBilinear2d",
|
||||
"UpsamplingNearest2d",
|
||||
"ZeroPad1d",
|
||||
"ZeroPad2d",
|
||||
"ZeroPad3d",
|
||||
]
|
||||
|
||||
# Please keep this list sorted
|
||||
if __all__ != sorted(__all__):
|
||||
raise AssertionError("__all__ must be sorted")
|
||||
@@ -0,0 +1,320 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from torch.autograd.function import Function
|
||||
|
||||
|
||||
class SyncBatchNorm(Function):
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def forward(
|
||||
self,
|
||||
input,
|
||||
weight,
|
||||
bias,
|
||||
running_mean,
|
||||
running_var,
|
||||
eps,
|
||||
momentum,
|
||||
process_group,
|
||||
world_size,
|
||||
):
|
||||
if not (
|
||||
input.is_contiguous(memory_format=torch.channels_last)
|
||||
or input.is_contiguous(memory_format=torch.channels_last_3d)
|
||||
):
|
||||
input = input.contiguous()
|
||||
if weight is not None:
|
||||
weight = weight.contiguous()
|
||||
|
||||
size = int(input.numel() // input.size(1))
|
||||
if size == 1 and world_size < 2:
|
||||
raise ValueError(
|
||||
f"Expected more than 1 value per channel when training, got input size {size}"
|
||||
)
|
||||
|
||||
num_channels = input.shape[1]
|
||||
if input.numel() > 0:
|
||||
# calculate mean/invstd for input.
|
||||
mean, invstd = torch.batch_norm_stats(input, eps)
|
||||
|
||||
count = torch.full(
|
||||
(1,),
|
||||
input.numel() // input.size(1),
|
||||
dtype=mean.dtype,
|
||||
device=mean.device,
|
||||
)
|
||||
|
||||
# C, C, 1 -> (2C + 1)
|
||||
combined = torch.cat([mean, invstd, count], dim=0)
|
||||
else:
|
||||
# for empty input, set stats and the count to zero. The stats with
|
||||
# zero count will be filtered out later when computing global mean
|
||||
# & invstd, but they still needs to participate the all_gather
|
||||
# collective communication to unblock other peer processes.
|
||||
combined = torch.zeros(
|
||||
2 * num_channels + 1, dtype=input.dtype, device=input.device
|
||||
)
|
||||
|
||||
# Use allgather instead of allreduce because count could be different across
|
||||
# ranks, simple all reduce op can not give correct results.
|
||||
# batch_norm_gather_stats_with_counts calculates global mean & invstd based on
|
||||
# all gathered mean, invstd and count.
|
||||
# for nccl backend, use the optimized version of all gather.
|
||||
# The Gloo backend does not support `all_gather_into_tensor`.
|
||||
if process_group._get_backend_name() != "gloo":
|
||||
# world_size * (2C + 1)
|
||||
combined_size = combined.numel()
|
||||
combined_flat = torch.empty(
|
||||
1,
|
||||
combined_size * world_size,
|
||||
dtype=combined.dtype,
|
||||
device=combined.device,
|
||||
)
|
||||
dist.all_gather_into_tensor(
|
||||
combined_flat, combined, process_group, async_op=False
|
||||
)
|
||||
combined = torch.reshape(combined_flat, (world_size, combined_size))
|
||||
# world_size * (2C + 1) -> world_size * C, world_size * C, world_size * 1
|
||||
mean_all, invstd_all, count_all = torch.split(combined, num_channels, dim=1)
|
||||
else:
|
||||
# world_size * (2C + 1)
|
||||
combined_list = [torch.empty_like(combined) for _ in range(world_size)]
|
||||
dist.all_gather(combined_list, combined, process_group, async_op=False)
|
||||
combined = torch.stack(combined_list, dim=0)
|
||||
# world_size * (2C + 1) -> world_size * C, world_size * C, world_size * 1
|
||||
mean_all, invstd_all, count_all = torch.split(combined, num_channels, dim=1)
|
||||
|
||||
if not (torch.cuda.is_available() and torch.cuda.is_current_stream_capturing()):
|
||||
# The lines below force a synchronization between CUDA and CPU, because
|
||||
# the shape of the result count_all depends on the values in mask tensor.
|
||||
# Such synchronizations break CUDA Graph capturing.
|
||||
# See https://github.com/pytorch/pytorch/issues/78549
|
||||
# FIXME: https://github.com/pytorch/pytorch/issues/78656 describes
|
||||
# a better longer-term solution.
|
||||
|
||||
# remove stats from empty inputs
|
||||
mask = count_all.squeeze(-1) >= 1
|
||||
count_all = count_all[mask]
|
||||
mean_all = mean_all[mask]
|
||||
invstd_all = invstd_all[mask]
|
||||
|
||||
# calculate global mean & invstd
|
||||
counts = count_all.view(-1)
|
||||
if running_mean is not None and counts.dtype != running_mean.dtype:
|
||||
counts = counts.to(running_mean.dtype)
|
||||
mean, invstd = torch.batch_norm_gather_stats_with_counts(
|
||||
input,
|
||||
mean_all,
|
||||
invstd_all,
|
||||
running_mean,
|
||||
running_var,
|
||||
momentum,
|
||||
eps,
|
||||
counts,
|
||||
)
|
||||
|
||||
self.save_for_backward(input, weight, mean, invstd, count_all.to(torch.int32))
|
||||
self.process_group = process_group
|
||||
|
||||
# apply element-wise normalization
|
||||
if input.numel() > 0:
|
||||
return torch.batch_norm_elemt(input, weight, bias, mean, invstd, eps)
|
||||
else:
|
||||
return torch.empty_like(input)
|
||||
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def backward(self, grad_output):
|
||||
if not (
|
||||
grad_output.is_contiguous(memory_format=torch.channels_last)
|
||||
or grad_output.is_contiguous(memory_format=torch.channels_last_3d)
|
||||
):
|
||||
grad_output = grad_output.contiguous()
|
||||
saved_input, weight, mean, invstd, count_tensor = self.saved_tensors
|
||||
grad_input = grad_weight = grad_bias = None
|
||||
process_group = self.process_group
|
||||
|
||||
if saved_input.numel() > 0:
|
||||
# calculate local stats as well as grad_weight / grad_bias
|
||||
(
|
||||
sum_dy,
|
||||
sum_dy_xmu,
|
||||
grad_weight,
|
||||
grad_bias,
|
||||
) = torch.batch_norm_backward_reduce(
|
||||
grad_output,
|
||||
saved_input,
|
||||
mean,
|
||||
invstd,
|
||||
weight,
|
||||
self.needs_input_grad[0],
|
||||
self.needs_input_grad[1],
|
||||
self.needs_input_grad[2],
|
||||
)
|
||||
|
||||
if self.needs_input_grad[0]:
|
||||
# synchronizing stats used to calculate input gradient.
|
||||
num_channels = sum_dy.shape[0]
|
||||
combined = torch.cat([sum_dy, sum_dy_xmu], dim=0)
|
||||
torch.distributed.all_reduce(
|
||||
combined,
|
||||
torch.distributed.ReduceOp.SUM,
|
||||
process_group,
|
||||
async_op=False,
|
||||
)
|
||||
sum_dy, sum_dy_xmu = torch.split(combined, num_channels)
|
||||
|
||||
# backward pass for gradient calculation
|
||||
if weight is not None and weight.dtype != mean.dtype:
|
||||
weight = weight.to(mean.dtype)
|
||||
grad_input = torch.batch_norm_backward_elemt(
|
||||
grad_output,
|
||||
saved_input,
|
||||
mean,
|
||||
invstd,
|
||||
weight,
|
||||
sum_dy,
|
||||
sum_dy_xmu,
|
||||
count_tensor,
|
||||
)
|
||||
# synchronizing of grad_weight / grad_bias is not needed as distributed
|
||||
# training would handle all reduce.
|
||||
if weight is None or not self.needs_input_grad[1]:
|
||||
grad_weight = None
|
||||
|
||||
if weight is None or not self.needs_input_grad[2]:
|
||||
grad_bias = None
|
||||
else:
|
||||
# This process got an empty input tensor in the forward pass.
|
||||
# Although this process can directly set grad_input as an empty
|
||||
# tensor of zeros, it still needs to participate in the collective
|
||||
# communication to unblock its peers, as other peer processes might
|
||||
# have received non-empty inputs.
|
||||
num_channels = saved_input.shape[1]
|
||||
if self.needs_input_grad[0]:
|
||||
# launch all_reduce to unblock other peer processes
|
||||
combined = torch.zeros(
|
||||
2 * num_channels, dtype=saved_input.dtype, device=saved_input.device
|
||||
)
|
||||
torch.distributed.all_reduce(
|
||||
combined,
|
||||
torch.distributed.ReduceOp.SUM,
|
||||
process_group,
|
||||
async_op=False,
|
||||
)
|
||||
|
||||
# Leave grad_input, grad_weight and grad_bias as None, which will be
|
||||
# interpreted by the autograd engine as Tensors full of zeros.
|
||||
|
||||
return grad_input, grad_weight, grad_bias, None, None, None, None, None, None
|
||||
|
||||
|
||||
class CrossMapLRN2d(Function):
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def forward(ctx, input, size, alpha=1e-4, beta=0.75, k=1):
|
||||
ctx.size = size
|
||||
ctx.alpha = alpha
|
||||
ctx.beta = beta
|
||||
ctx.k = k
|
||||
ctx.scale = None
|
||||
|
||||
if input.dim() != 4:
|
||||
raise ValueError(
|
||||
f"CrossMapLRN2d: Expected input to be 4D, got {input.dim()}D instead."
|
||||
)
|
||||
|
||||
ctx.scale = ctx.scale or input.new()
|
||||
output = input.new()
|
||||
channels = input.size(1)
|
||||
|
||||
output.resize_as_(input)
|
||||
ctx.scale.resize_as_(input)
|
||||
|
||||
# use output storage as temporary buffer
|
||||
input_square = output
|
||||
torch.pow(input, 2, out=input_square)
|
||||
|
||||
pre_pad = int((ctx.size - 1) / 2 + 1)
|
||||
pre_pad_crop = min(pre_pad, channels)
|
||||
|
||||
scale_first = ctx.scale.select(1, 0)
|
||||
scale_first.zero_()
|
||||
# compute first feature map normalization
|
||||
for c in range(pre_pad_crop):
|
||||
scale_first.add_(input_square.select(1, c))
|
||||
|
||||
# reuse computations for next feature maps normalization
|
||||
# by adding the next feature map and removing the previous
|
||||
for c in range(1, channels):
|
||||
scale_previous = ctx.scale.select(1, c - 1)
|
||||
scale_current = ctx.scale.select(1, c)
|
||||
scale_current.copy_(scale_previous)
|
||||
if c < channels - pre_pad + 1:
|
||||
square_next = input_square.select(1, c + pre_pad - 1)
|
||||
scale_current.add_(square_next, alpha=1)
|
||||
|
||||
if c > pre_pad:
|
||||
square_previous = input_square.select(1, c - pre_pad)
|
||||
scale_current.add_(square_previous, alpha=-1)
|
||||
|
||||
ctx.scale.mul_(ctx.alpha / ctx.size).add_(ctx.k)
|
||||
|
||||
torch.pow(ctx.scale, -ctx.beta, out=output)
|
||||
output.mul_(input)
|
||||
|
||||
ctx.save_for_backward(input, output)
|
||||
return output
|
||||
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def backward(ctx, grad_output):
|
||||
input, output = ctx.saved_tensors
|
||||
grad_input = grad_output.new()
|
||||
|
||||
batch_size = input.size(0)
|
||||
channels = input.size(1)
|
||||
input_height = input.size(2)
|
||||
input_width = input.size(3)
|
||||
|
||||
paddded_ratio = input.new(channels + ctx.size - 1, input_height, input_width)
|
||||
accum_ratio = input.new(input_height, input_width)
|
||||
|
||||
cache_ratio_value = 2 * ctx.alpha * ctx.beta / ctx.size
|
||||
inversePrePad = int(ctx.size - (ctx.size - 1) / 2)
|
||||
|
||||
grad_input.resize_as_(input)
|
||||
torch.pow(ctx.scale, -ctx.beta, out=grad_input).mul_(grad_output)
|
||||
|
||||
paddded_ratio.zero_()
|
||||
padded_ratio_center = paddded_ratio.narrow(0, inversePrePad, channels)
|
||||
for n in range(batch_size):
|
||||
torch.mul(grad_output[n], output[n], out=padded_ratio_center)
|
||||
padded_ratio_center.div_(ctx.scale[n])
|
||||
torch.sum(
|
||||
paddded_ratio.narrow(0, 0, ctx.size - 1),
|
||||
0,
|
||||
keepdim=False,
|
||||
out=accum_ratio,
|
||||
)
|
||||
for c in range(channels):
|
||||
accum_ratio.add_(paddded_ratio[c + ctx.size - 1])
|
||||
grad_input[n][c].addcmul_(
|
||||
input[n][c], accum_ratio, value=-cache_ratio_value
|
||||
)
|
||||
accum_ratio.add_(paddded_ratio[c], alpha=-1)
|
||||
|
||||
return grad_input, None, None, None, None
|
||||
|
||||
|
||||
class BackwardHookFunction(torch.autograd.Function):
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def forward(ctx, *args):
|
||||
ctx.mark_non_differentiable(*[arg for arg in args if not arg.requires_grad])
|
||||
return args
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args):
|
||||
return args
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,339 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import itertools
|
||||
from collections import namedtuple
|
||||
from collections.abc import Sequence
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
|
||||
from .container import ModuleList, Sequential
|
||||
from .linear import Linear
|
||||
from .module import Module
|
||||
|
||||
|
||||
__all__ = ["AdaptiveLogSoftmaxWithLoss"]
|
||||
|
||||
_ASMoutput = namedtuple("_ASMoutput", ["output", "loss"])
|
||||
|
||||
|
||||
class AdaptiveLogSoftmaxWithLoss(Module):
|
||||
(
|
||||
"""Efficient softmax approximation.
|
||||
|
||||
As described in
|
||||
`Efficient softmax approximation for GPUs by Edouard Grave, Armand Joulin,
|
||||
Moustapha Ciss\u00e9, David Grangier, and Herv\u00e9 J\u00e9gou
|
||||
<https://arxiv.org/abs/1609.04309>`__.
|
||||
"""
|
||||
r"""
|
||||
Adaptive softmax is an approximate strategy for training models with large
|
||||
output spaces. It is most effective when the label distribution is highly
|
||||
imbalanced, for example in natural language modelling, where the word
|
||||
frequency distribution approximately follows the `Zipf's law`_.
|
||||
|
||||
Adaptive softmax partitions the labels into several clusters, according to
|
||||
their frequency. These clusters may contain different number of targets
|
||||
each.
|
||||
Additionally, clusters containing less frequent labels assign lower
|
||||
dimensional embeddings to those labels, which speeds up the computation.
|
||||
For each minibatch, only clusters for which at least one target is
|
||||
present are evaluated.
|
||||
|
||||
The idea is that the clusters which are accessed frequently
|
||||
(like the first one, containing most frequent labels), should also be cheap
|
||||
to compute -- that is, contain a small number of assigned labels.
|
||||
|
||||
We highly recommend taking a look at the original paper for more details.
|
||||
|
||||
* :attr:`cutoffs` should be an ordered Sequence of integers sorted
|
||||
in the increasing order.
|
||||
It controls number of clusters and the partitioning of targets into
|
||||
clusters. For example setting ``cutoffs = [10, 100, 1000]``
|
||||
means that first `10` targets will be assigned
|
||||
to the 'head' of the adaptive softmax, targets `11, 12, ..., 100` will be
|
||||
assigned to the first cluster, and targets `101, 102, ..., 1000` will be
|
||||
assigned to the second cluster, while targets
|
||||
`1001, 1002, ..., n_classes - 1` will be assigned
|
||||
to the last, third cluster.
|
||||
|
||||
* :attr:`div_value` is used to compute the size of each additional cluster,
|
||||
which is given as
|
||||
:math:`\left\lfloor\frac{\texttt{in\_features}}{\texttt{div\_value}^{idx}}\right\rfloor`,
|
||||
where :math:`idx` is the cluster index (with clusters
|
||||
for less frequent words having larger indices,
|
||||
and indices starting from :math:`1`).
|
||||
|
||||
* :attr:`head_bias` if set to True, adds a bias term to the 'head' of the
|
||||
adaptive softmax. See paper for details. Set to False in the official
|
||||
implementation.
|
||||
|
||||
.. warning::
|
||||
Labels passed as inputs to this module should be sorted according to
|
||||
their frequency. This means that the most frequent label should be
|
||||
represented by the index `0`, and the least frequent
|
||||
label should be represented by the index `n_classes - 1`.
|
||||
|
||||
.. note::
|
||||
This module returns a ``NamedTuple`` with ``output``
|
||||
and ``loss`` fields. See further documentation for details.
|
||||
|
||||
.. note::
|
||||
To compute log-probabilities for all classes, the ``log_prob``
|
||||
method can be used.
|
||||
|
||||
Args:
|
||||
in_features (int): Number of features in the input tensor
|
||||
n_classes (int): Number of classes in the dataset
|
||||
cutoffs (Sequence): Cutoffs used to assign targets to their buckets
|
||||
div_value (float, optional): value used as an exponent to compute sizes
|
||||
of the clusters. Default: 4.0
|
||||
head_bias (bool, optional): If ``True``, adds a bias term to the 'head' of the
|
||||
adaptive softmax. Default: ``False``
|
||||
|
||||
Returns:
|
||||
``NamedTuple`` with ``output`` and ``loss`` fields:
|
||||
* **output** is a Tensor of size ``N`` containing computed target
|
||||
log probabilities for each example
|
||||
* **loss** is a Scalar representing the computed negative
|
||||
log likelihood loss
|
||||
|
||||
Shape:
|
||||
- input: :math:`(N, \texttt{in\_features})` or :math:`(\texttt{in\_features})`
|
||||
- target: :math:`(N)` or :math:`()` where each value satisfies :math:`0 <= \texttt{target[i]} <= \texttt{n\_classes}`
|
||||
- output1: :math:`(N)` or :math:`()`
|
||||
- output2: ``Scalar``
|
||||
|
||||
.. _Zipf's law: https://en.wikipedia.org/wiki/Zipf%27s_law
|
||||
"""
|
||||
)
|
||||
|
||||
in_features: int
|
||||
n_classes: int
|
||||
cutoffs: list[int]
|
||||
div_value: float
|
||||
head_bias: bool
|
||||
head: Linear
|
||||
tail: ModuleList
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
n_classes: int,
|
||||
cutoffs: Sequence[int],
|
||||
div_value: float = 4.0,
|
||||
head_bias: bool = False,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
|
||||
cutoffs = list(cutoffs)
|
||||
|
||||
if len(cutoffs) == 0:
|
||||
raise ValueError("cutoffs should be a sequence of length larger than 0")
|
||||
|
||||
if (
|
||||
(cutoffs != sorted(cutoffs))
|
||||
or (min(cutoffs) <= 0)
|
||||
or (max(cutoffs) > (n_classes - 1))
|
||||
or (len(set(cutoffs)) != len(cutoffs))
|
||||
or any(int(c) != c for c in cutoffs)
|
||||
):
|
||||
raise ValueError(
|
||||
"cutoffs should be a sequence of unique, positive "
|
||||
"integers sorted in an increasing order, where "
|
||||
"each value is between 1 and n_classes-1"
|
||||
)
|
||||
|
||||
self.in_features = in_features
|
||||
self.n_classes = n_classes
|
||||
self.cutoffs = cutoffs + [n_classes]
|
||||
self.div_value = div_value
|
||||
self.head_bias = head_bias
|
||||
|
||||
self.shortlist_size = self.cutoffs[0]
|
||||
self.n_clusters = len(self.cutoffs) - 1
|
||||
self.head_size = self.shortlist_size + self.n_clusters
|
||||
|
||||
self.head = Linear(
|
||||
self.in_features, self.head_size, bias=self.head_bias, **factory_kwargs
|
||||
)
|
||||
self.tail = ModuleList()
|
||||
|
||||
for i in range(self.n_clusters):
|
||||
hsz = int(self.in_features // (self.div_value ** (i + 1)))
|
||||
osz = self.cutoffs[i + 1] - self.cutoffs[i]
|
||||
|
||||
projection = Sequential(
|
||||
Linear(self.in_features, hsz, bias=False, **factory_kwargs),
|
||||
Linear(hsz, osz, bias=False, **factory_kwargs),
|
||||
)
|
||||
|
||||
self.tail.append(projection)
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
"""
|
||||
Resets parameters based on their initialization used in ``__init__``.
|
||||
"""
|
||||
self.head.reset_parameters()
|
||||
for i2h, h2o in self.tail: # type: ignore[misc]
|
||||
i2h.reset_parameters() # type: ignore[has-type]
|
||||
h2o.reset_parameters() # type: ignore[has-type]
|
||||
|
||||
def forward(self, input_: Tensor, target_: Tensor) -> _ASMoutput:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
targ_dim = target_.dim()
|
||||
|
||||
if targ_dim == 1:
|
||||
if input_.size(0) != target_.size(0):
|
||||
raise RuntimeError(
|
||||
"Input and target should have the same size in the batch dimension."
|
||||
)
|
||||
if input_.dim() != 2:
|
||||
raise RuntimeError(
|
||||
"1D target tensor expects 2D input tensors, "
|
||||
"but found inputs with size",
|
||||
input_.size(),
|
||||
)
|
||||
elif targ_dim == 0:
|
||||
if input_.dim() != 1:
|
||||
raise RuntimeError(
|
||||
"0D target tensor expects 1D input tensors, "
|
||||
"but found inputs with size",
|
||||
input_.size(),
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
"0D or 1D target tensor expected, multi-target not supported"
|
||||
)
|
||||
|
||||
is_batched = targ_dim > 0
|
||||
input = input_ if is_batched else input_.unsqueeze(0)
|
||||
target = target_ if is_batched else target_.unsqueeze(0)
|
||||
|
||||
used_rows = 0
|
||||
batch_size = target.size(0)
|
||||
|
||||
output = input.new_zeros(batch_size)
|
||||
gather_inds = target.new_empty(batch_size)
|
||||
|
||||
cutoff_values = [0] + self.cutoffs
|
||||
for i in range(len(cutoff_values) - 1):
|
||||
low_idx = cutoff_values[i]
|
||||
high_idx = cutoff_values[i + 1]
|
||||
|
||||
target_mask = (target >= low_idx) & (target < high_idx)
|
||||
row_indices = target_mask.nonzero().squeeze()
|
||||
|
||||
if row_indices.numel() == 0:
|
||||
continue
|
||||
|
||||
if i == 0:
|
||||
gather_inds.index_copy_(0, row_indices, target[target_mask])
|
||||
|
||||
else:
|
||||
relative_target = target[target_mask] - low_idx
|
||||
input_subset = input.index_select(0, row_indices)
|
||||
|
||||
cluster_output = self.tail[i - 1](input_subset)
|
||||
cluster_index = self.shortlist_size + i - 1
|
||||
|
||||
gather_inds.index_fill_(0, row_indices, cluster_index)
|
||||
cluster_logprob = F.log_softmax(cluster_output, dim=1)
|
||||
local_logprob = cluster_logprob.gather(1, relative_target.unsqueeze(1))
|
||||
output.index_copy_(0, row_indices, local_logprob.squeeze(1))
|
||||
|
||||
used_rows += row_indices.numel()
|
||||
|
||||
if used_rows != batch_size:
|
||||
raise RuntimeError(
|
||||
f"Target values should be in [0, {self.n_classes - 1}], "
|
||||
f"but values in range [{target.min().item()}, {target.max().item()}] "
|
||||
"were found. "
|
||||
)
|
||||
|
||||
head_output = self.head(input)
|
||||
head_logprob = F.log_softmax(head_output, dim=1)
|
||||
output += head_logprob.gather(1, gather_inds.unsqueeze(1)).squeeze()
|
||||
loss = (-output).mean()
|
||||
|
||||
if not is_batched:
|
||||
output = output.squeeze(0)
|
||||
|
||||
return _ASMoutput(output, loss)
|
||||
|
||||
def _get_full_log_prob(self, input, head_output):
|
||||
"""Given input tensor, and output of ``self.head``, compute the log of the full distribution."""
|
||||
out = input.new_empty((head_output.size(0), self.n_classes))
|
||||
head_logprob = F.log_softmax(head_output, dim=1)
|
||||
|
||||
out[:, : self.shortlist_size] = head_logprob[:, : self.shortlist_size]
|
||||
|
||||
for i, (start_idx, stop_idx) in enumerate(itertools.pairwise(self.cutoffs)):
|
||||
cluster_output = self.tail[i](input)
|
||||
cluster_logprob = F.log_softmax(cluster_output, dim=1)
|
||||
output_logprob = cluster_logprob + head_logprob[
|
||||
:, self.shortlist_size + i
|
||||
].unsqueeze(1)
|
||||
|
||||
out[:, start_idx:stop_idx] = output_logprob
|
||||
|
||||
return out
|
||||
|
||||
def log_prob(self, input: Tensor) -> Tensor:
|
||||
r"""Compute log probabilities for all :math:`\texttt{n\_classes}`.
|
||||
|
||||
Args:
|
||||
input (Tensor): a minibatch of examples
|
||||
|
||||
Returns:
|
||||
log-probabilities of for each class :math:`c`
|
||||
in range :math:`0 <= c <= \texttt{n\_classes}`, where :math:`\texttt{n\_classes}` is a
|
||||
parameter passed to ``AdaptiveLogSoftmaxWithLoss`` constructor.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, \texttt{in\_features})`
|
||||
- Output: :math:`(N, \texttt{n\_classes})`
|
||||
|
||||
"""
|
||||
head_output = self.head(input)
|
||||
return self._get_full_log_prob(input, head_output)
|
||||
|
||||
def predict(self, input: Tensor) -> Tensor:
|
||||
r"""Return the class with the highest probability for each example in the input minibatch.
|
||||
|
||||
This is equivalent to ``self.log_prob(input).argmax(dim=1)``, but is more efficient in some cases.
|
||||
|
||||
Args:
|
||||
input (Tensor): a minibatch of examples
|
||||
|
||||
Returns:
|
||||
output (Tensor): a class with the highest probability for each example
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, \texttt{in\_features})`
|
||||
- Output: :math:`(N)`
|
||||
"""
|
||||
head_output = self.head(input)
|
||||
output = torch.argmax(head_output, dim=1)
|
||||
not_in_shortlist = output >= self.shortlist_size
|
||||
all_in_shortlist = not (not_in_shortlist.any())
|
||||
|
||||
if all_in_shortlist:
|
||||
return output
|
||||
|
||||
elif not_in_shortlist.all():
|
||||
log_prob = self._get_full_log_prob(input, head_output)
|
||||
return torch.argmax(log_prob, dim=1)
|
||||
|
||||
else:
|
||||
log_prob = self._get_full_log_prob(
|
||||
input[not_in_shortlist], head_output[not_in_shortlist]
|
||||
)
|
||||
output[not_in_shortlist] = torch.argmax(log_prob, dim=1)
|
||||
return output
|
||||
@@ -0,0 +1,951 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.nn import functional as F, init
|
||||
from torch.nn.parameter import Parameter, UninitializedBuffer, UninitializedParameter
|
||||
|
||||
from ._functions import SyncBatchNorm as sync_batch_norm
|
||||
from .lazy import LazyModuleMixin
|
||||
from .module import Module
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BatchNorm1d",
|
||||
"LazyBatchNorm1d",
|
||||
"BatchNorm2d",
|
||||
"LazyBatchNorm2d",
|
||||
"BatchNorm3d",
|
||||
"LazyBatchNorm3d",
|
||||
"SyncBatchNorm",
|
||||
]
|
||||
|
||||
|
||||
class _NormBase(Module):
|
||||
"""Common base of _InstanceNorm and _BatchNorm."""
|
||||
|
||||
_version = 2
|
||||
__constants__ = ["track_running_stats", "momentum", "eps", "num_features", "affine"]
|
||||
num_features: int
|
||||
eps: float
|
||||
momentum: float | None
|
||||
affine: bool
|
||||
track_running_stats: bool
|
||||
# WARNING: weight and bias purposely not defined here.
|
||||
# See https://github.com/pytorch/pytorch/issues/39670
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_features: int,
|
||||
eps: float = 1e-5,
|
||||
momentum: float | None = 0.1,
|
||||
affine: bool = True,
|
||||
track_running_stats: bool = True,
|
||||
device=None,
|
||||
dtype=None,
|
||||
*,
|
||||
bias: bool = True,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
self.num_features = num_features
|
||||
self.eps = eps
|
||||
self.momentum = momentum
|
||||
self.affine = affine
|
||||
self.track_running_stats = track_running_stats
|
||||
if self.affine:
|
||||
self.weight = Parameter(torch.empty(num_features, **factory_kwargs))
|
||||
if bias:
|
||||
self.bias = Parameter(torch.empty(num_features, **factory_kwargs))
|
||||
else:
|
||||
self.register_parameter("bias", None)
|
||||
else:
|
||||
self.register_parameter("weight", None)
|
||||
self.register_parameter("bias", None)
|
||||
if self.track_running_stats:
|
||||
self.register_buffer(
|
||||
"running_mean", torch.zeros(num_features, **factory_kwargs)
|
||||
)
|
||||
self.register_buffer(
|
||||
"running_var", torch.ones(num_features, **factory_kwargs)
|
||||
)
|
||||
self.running_mean: Tensor | None
|
||||
self.running_var: Tensor | None
|
||||
self.register_buffer(
|
||||
"num_batches_tracked",
|
||||
torch.tensor(
|
||||
0,
|
||||
dtype=torch.long,
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
**{k: v for k, v in factory_kwargs.items() if k != "dtype"},
|
||||
),
|
||||
)
|
||||
self.num_batches_tracked: Tensor | None
|
||||
else:
|
||||
self.register_buffer("running_mean", None)
|
||||
self.register_buffer("running_var", None)
|
||||
self.register_buffer("num_batches_tracked", None)
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_running_stats(self) -> None:
|
||||
if self.track_running_stats:
|
||||
# running_mean/running_var/num_batches... are registered at runtime depending
|
||||
# if self.track_running_stats is on
|
||||
self.running_mean.zero_() # type: ignore[union-attr]
|
||||
self.running_var.fill_(1) # type: ignore[union-attr]
|
||||
self.num_batches_tracked.zero_() # type: ignore[union-attr,operator]
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
self.reset_running_stats()
|
||||
if self.affine:
|
||||
init.ones_(self.weight)
|
||||
if self.bias is not None:
|
||||
init.zeros_(self.bias)
|
||||
|
||||
def _check_input_dim(self, input):
|
||||
raise NotImplementedError
|
||||
|
||||
def extra_repr(self):
|
||||
return (
|
||||
"{num_features}, eps={eps}, momentum={momentum}, affine={affine}, "
|
||||
"bias={use_bias}, track_running_stats={track_running_stats}".format(
|
||||
**self.__dict__, use_bias=self.bias is not None
|
||||
)
|
||||
)
|
||||
|
||||
def _load_from_state_dict(
|
||||
self,
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
strict,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
) -> None:
|
||||
version = local_metadata.get("version", None)
|
||||
|
||||
if (version is None or version < 2) and self.track_running_stats:
|
||||
# at version 2: added num_batches_tracked buffer
|
||||
# this should have a default value of 0
|
||||
num_batches_tracked_key = prefix + "num_batches_tracked"
|
||||
if num_batches_tracked_key not in state_dict:
|
||||
state_dict[num_batches_tracked_key] = (
|
||||
self.num_batches_tracked
|
||||
if self.num_batches_tracked is not None
|
||||
and self.num_batches_tracked.device != torch.device("meta")
|
||||
else torch.tensor(0, dtype=torch.long)
|
||||
)
|
||||
|
||||
super()._load_from_state_dict(
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
strict,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
)
|
||||
|
||||
|
||||
class _BatchNorm(_NormBase):
|
||||
def __init__(
|
||||
self,
|
||||
num_features: int,
|
||||
eps: float = 1e-5,
|
||||
momentum: float | None = 0.1,
|
||||
affine: bool = True,
|
||||
track_running_stats: bool = True,
|
||||
device=None,
|
||||
dtype=None,
|
||||
*,
|
||||
bias: bool = True,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__(
|
||||
num_features,
|
||||
eps,
|
||||
momentum,
|
||||
affine,
|
||||
track_running_stats,
|
||||
**factory_kwargs,
|
||||
bias=bias,
|
||||
)
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
self._check_input_dim(input)
|
||||
|
||||
# exponential_average_factor is set to self.momentum
|
||||
# (when it is available) only so that it gets updated
|
||||
# in ONNX graph when this node is exported to ONNX.
|
||||
if self.momentum is None:
|
||||
exponential_average_factor = 0.0
|
||||
else:
|
||||
exponential_average_factor = self.momentum
|
||||
|
||||
if self.training and self.track_running_stats:
|
||||
# TODO: if statement only here to tell the jit to skip emitting this when it is None
|
||||
if self.num_batches_tracked is not None: # type: ignore[has-type]
|
||||
self.num_batches_tracked.add_(1) # type: ignore[has-type]
|
||||
if self.momentum is None: # use cumulative moving average
|
||||
exponential_average_factor = 1.0 / float(self.num_batches_tracked)
|
||||
else: # use exponential moving average
|
||||
exponential_average_factor = self.momentum
|
||||
|
||||
r"""
|
||||
Decide whether the mini-batch stats should be used for normalization rather than the buffers.
|
||||
Mini-batch stats are used in training mode, and in eval mode when buffers are None.
|
||||
"""
|
||||
if self.training:
|
||||
bn_training = True
|
||||
else:
|
||||
bn_training = (self.running_mean is None) and (self.running_var is None)
|
||||
|
||||
r"""
|
||||
Buffers are only updated if they are to be tracked and we are in training mode. Thus they only need to be
|
||||
passed when the update should occur (i.e. in training mode when they are tracked), or when buffer stats are
|
||||
used for normalization (i.e. in eval mode when buffers are not None).
|
||||
"""
|
||||
return F.batch_norm(
|
||||
input,
|
||||
# If buffers are not to be tracked, ensure that they won't be updated
|
||||
(
|
||||
self.running_mean
|
||||
if not self.training or self.track_running_stats
|
||||
else None
|
||||
),
|
||||
self.running_var if not self.training or self.track_running_stats else None,
|
||||
self.weight,
|
||||
self.bias,
|
||||
bn_training,
|
||||
exponential_average_factor,
|
||||
self.eps,
|
||||
)
|
||||
|
||||
|
||||
class _LazyNormBase(LazyModuleMixin, _NormBase):
|
||||
weight: UninitializedParameter # type: ignore[assignment]
|
||||
bias: UninitializedParameter # type: ignore[assignment]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
eps=1e-5,
|
||||
momentum=0.1,
|
||||
affine=True,
|
||||
track_running_stats=True,
|
||||
device=None,
|
||||
dtype=None,
|
||||
*,
|
||||
bias=True,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
super().__init__(
|
||||
# affine, bias and track_running_stats are hardcoded to False to
|
||||
# avoid creating tensors that will soon be overwritten.
|
||||
0,
|
||||
eps,
|
||||
momentum,
|
||||
False,
|
||||
False,
|
||||
**factory_kwargs,
|
||||
bias=False,
|
||||
)
|
||||
self.affine = affine
|
||||
self.track_running_stats = track_running_stats
|
||||
if self.affine:
|
||||
# pyrefly: ignore [unexpected-keyword]
|
||||
self.weight = UninitializedParameter(**factory_kwargs)
|
||||
if bias:
|
||||
# pyrefly: ignore # bad-argument-type
|
||||
self.bias = UninitializedParameter(**factory_kwargs)
|
||||
if self.track_running_stats:
|
||||
# pyrefly: ignore [unexpected-keyword]
|
||||
self.running_mean = UninitializedBuffer(**factory_kwargs)
|
||||
# pyrefly: ignore [unexpected-keyword]
|
||||
self.running_var = UninitializedBuffer(**factory_kwargs)
|
||||
self.num_batches_tracked = torch.tensor(
|
||||
0,
|
||||
dtype=torch.long,
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
**{k: v for k, v in factory_kwargs.items() if k != "dtype"},
|
||||
)
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
if not self.has_uninitialized_params() and self.num_features != 0:
|
||||
super().reset_parameters()
|
||||
|
||||
def initialize_parameters(self, input) -> None: # type: ignore[override]
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
if self.has_uninitialized_params():
|
||||
self.num_features = input.shape[1]
|
||||
if self.affine:
|
||||
if not isinstance(self.weight, UninitializedParameter):
|
||||
raise AssertionError(
|
||||
"self.weight must be an UninitializedParameter"
|
||||
)
|
||||
self.weight.materialize((self.num_features,))
|
||||
if self.bias is not None:
|
||||
if not isinstance(self.bias, UninitializedParameter):
|
||||
raise AssertionError(
|
||||
"self.bias must be an UninitializedParameter"
|
||||
)
|
||||
self.bias.materialize((self.num_features,))
|
||||
if self.track_running_stats:
|
||||
self.running_mean.materialize( # type:ignore[union-attr]
|
||||
(self.num_features,)
|
||||
)
|
||||
self.running_var.materialize( # type:ignore[union-attr]
|
||||
(self.num_features,)
|
||||
)
|
||||
self.reset_parameters()
|
||||
|
||||
|
||||
class BatchNorm1d(_BatchNorm):
|
||||
r"""Applies Batch Normalization over a 2D or 3D input.
|
||||
|
||||
Method described in the paper
|
||||
`Batch Normalization: Accelerating Deep Network Training by Reducing
|
||||
Internal Covariate Shift <https://arxiv.org/abs/1502.03167>`__ .
|
||||
|
||||
.. math::
|
||||
|
||||
y = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta
|
||||
|
||||
The mean and standard-deviation are calculated per-dimension over
|
||||
the mini-batches and :math:`\gamma` and :math:`\beta` are learnable parameter vectors
|
||||
of size `C` (where `C` is the number of features or channels of the input). By default, the
|
||||
elements of :math:`\gamma` are set to 1 and the elements of :math:`\beta` are set to 0.
|
||||
At train time in the forward pass, the variance is calculated via the biased estimator,
|
||||
equivalent to ``torch.var(input, correction=0)``. However, the value stored in the
|
||||
moving average of the variance is calculated via the unbiased estimator, equivalent to
|
||||
``torch.var(input, correction=1)``.
|
||||
|
||||
Also by default, during training this layer keeps running estimates of its
|
||||
computed mean and variance, which are then used for normalization during
|
||||
evaluation. The running estimates are kept with a default :attr:`momentum`
|
||||
of 0.1.
|
||||
|
||||
If :attr:`track_running_stats` is set to ``False``, this layer then does not
|
||||
keep running estimates, and batch statistics are instead used during
|
||||
evaluation time as well.
|
||||
|
||||
.. note::
|
||||
This :attr:`momentum` argument is different from one used in optimizer
|
||||
classes and the conventional notion of momentum. Mathematically, the
|
||||
update rule for running statistics here is
|
||||
:math:`\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t`,
|
||||
where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the
|
||||
new observed value.
|
||||
|
||||
Because the Batch Normalization is done over the `C` dimension, computing statistics
|
||||
on `(N, L)` slices, it's common terminology to call this Temporal Batch Normalization.
|
||||
|
||||
Args:
|
||||
num_features: number of features or channels :math:`C` of the input
|
||||
eps: a value added to the denominator for numerical stability.
|
||||
Default: 1e-5
|
||||
momentum: the value used for the running_mean and running_var
|
||||
computation. Can be set to ``None`` for cumulative moving average
|
||||
(i.e. simple average). Default: 0.1
|
||||
affine: a boolean value that when set to ``True``, this module has
|
||||
learnable affine parameters. Default: ``True``
|
||||
track_running_stats: a boolean value that when set to ``True``, this
|
||||
module tracks the running mean and variance, and when set to ``False``,
|
||||
this module does not track such statistics, and initializes statistics
|
||||
buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
|
||||
When these buffers are ``None``, this module always uses batch statistics.
|
||||
in both training and eval modes. Default: ``True``
|
||||
bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
|
||||
:attr:`affine` is ``True``). Default: ``True``
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C)` or :math:`(N, C, L)`, where :math:`N` is the batch size,
|
||||
:math:`C` is the number of features or channels, and :math:`L` is the sequence length
|
||||
- Output: :math:`(N, C)` or :math:`(N, C, L)` (same shape as input)
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # With Learnable Parameters
|
||||
>>> m = nn.BatchNorm1d(100)
|
||||
>>> # Without Learnable Parameters
|
||||
>>> m = nn.BatchNorm1d(100, affine=False)
|
||||
>>> input = torch.randn(20, 100)
|
||||
>>> output = m(input)
|
||||
"""
|
||||
|
||||
def _check_input_dim(self, input) -> None:
|
||||
if input.dim() != 2 and input.dim() != 3:
|
||||
raise ValueError(f"expected 2D or 3D input (got {input.dim()}D input)")
|
||||
|
||||
|
||||
class LazyBatchNorm1d(_LazyNormBase, _BatchNorm):
|
||||
r"""A :class:`torch.nn.BatchNorm1d` module with lazy initialization.
|
||||
|
||||
Lazy initialization based on the ``num_features`` argument of the :class:`BatchNorm1d` that is inferred
|
||||
from the ``input.size(1)``.
|
||||
The attributes that will be lazily initialized are `weight`, `bias`,
|
||||
`running_mean` and `running_var`.
|
||||
|
||||
Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation
|
||||
on lazy modules and their limitations.
|
||||
|
||||
Args:
|
||||
eps: a value added to the denominator for numerical stability.
|
||||
Default: 1e-5
|
||||
momentum: the value used for the running_mean and running_var
|
||||
computation. Can be set to ``None`` for cumulative moving average
|
||||
(i.e. simple average). Default: 0.1
|
||||
affine: a boolean value that when set to ``True``, this module has
|
||||
learnable affine parameters. Default: ``True``
|
||||
track_running_stats: a boolean value that when set to ``True``, this
|
||||
module tracks the running mean and variance, and when set to ``False``,
|
||||
this module does not track such statistics, and initializes statistics
|
||||
buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
|
||||
When these buffers are ``None``, this module always uses batch statistics.
|
||||
in both training and eval modes. Default: ``True``
|
||||
bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
|
||||
:attr:`affine` is ``True``). Default: ``True``
|
||||
"""
|
||||
|
||||
cls_to_become = BatchNorm1d # type: ignore[assignment]
|
||||
|
||||
def _check_input_dim(self, input) -> None:
|
||||
if input.dim() != 2 and input.dim() != 3:
|
||||
raise ValueError(f"expected 2D or 3D input (got {input.dim()}D input)")
|
||||
|
||||
|
||||
class BatchNorm2d(_BatchNorm):
|
||||
r"""Applies Batch Normalization over a 4D input.
|
||||
|
||||
4D is a mini-batch of 2D inputs
|
||||
with additional channel dimension. Method described in the paper
|
||||
`Batch Normalization: Accelerating Deep Network Training by Reducing
|
||||
Internal Covariate Shift <https://arxiv.org/abs/1502.03167>`__ .
|
||||
|
||||
.. math::
|
||||
|
||||
y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta
|
||||
|
||||
The mean and standard-deviation are calculated per-dimension over
|
||||
the mini-batches and :math:`\gamma` and :math:`\beta` are learnable parameter vectors
|
||||
of size `C` (where `C` is the input size). By default, the elements of :math:`\gamma` are set
|
||||
to 1 and the elements of :math:`\beta` are set to 0. At train time in the forward pass, the
|
||||
standard-deviation is calculated via the biased estimator, equivalent to
|
||||
``torch.var(input, correction=0)``. However, the value stored in the moving average of the
|
||||
standard-deviation is calculated via the unbiased estimator, equivalent to
|
||||
``torch.var(input, correction=1)``.
|
||||
|
||||
Also by default, during training this layer keeps running estimates of its
|
||||
computed mean and variance, which are then used for normalization during
|
||||
evaluation. The running estimates are kept with a default :attr:`momentum`
|
||||
of 0.1.
|
||||
|
||||
If :attr:`track_running_stats` is set to ``False``, this layer then does not
|
||||
keep running estimates, and batch statistics are instead used during
|
||||
evaluation time as well.
|
||||
|
||||
.. note::
|
||||
This :attr:`momentum` argument is different from one used in optimizer
|
||||
classes and the conventional notion of momentum. Mathematically, the
|
||||
update rule for running statistics here is
|
||||
:math:`\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t`,
|
||||
where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the
|
||||
new observed value.
|
||||
|
||||
Because the Batch Normalization is done over the `C` dimension, computing statistics
|
||||
on `(N, H, W)` slices, it's common terminology to call this Spatial Batch Normalization.
|
||||
|
||||
Args:
|
||||
num_features: :math:`C` from an expected input of size
|
||||
:math:`(N, C, H, W)`
|
||||
eps: a value added to the denominator for numerical stability.
|
||||
Default: 1e-5
|
||||
momentum: the value used for the running_mean and running_var
|
||||
computation. Can be set to ``None`` for cumulative moving average
|
||||
(i.e. simple average). Default: 0.1
|
||||
affine: a boolean value that when set to ``True``, this module has
|
||||
learnable affine parameters. Default: ``True``
|
||||
track_running_stats: a boolean value that when set to ``True``, this
|
||||
module tracks the running mean and variance, and when set to ``False``,
|
||||
this module does not track such statistics, and initializes statistics
|
||||
buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
|
||||
When these buffers are ``None``, this module always uses batch statistics.
|
||||
in both training and eval modes. Default: ``True``
|
||||
bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
|
||||
:attr:`affine` is ``True``). Default: ``True``
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, H, W)`
|
||||
- Output: :math:`(N, C, H, W)` (same shape as input)
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # With Learnable Parameters
|
||||
>>> m = nn.BatchNorm2d(100)
|
||||
>>> # Without Learnable Parameters
|
||||
>>> m = nn.BatchNorm2d(100, affine=False)
|
||||
>>> input = torch.randn(20, 100, 35, 45)
|
||||
>>> output = m(input)
|
||||
"""
|
||||
|
||||
def _check_input_dim(self, input) -> None:
|
||||
if input.dim() != 4:
|
||||
raise ValueError(f"expected 4D input (got {input.dim()}D input)")
|
||||
|
||||
|
||||
class LazyBatchNorm2d(_LazyNormBase, _BatchNorm):
|
||||
r"""A :class:`torch.nn.BatchNorm2d` module with lazy initialization.
|
||||
|
||||
Lazy initialization is done for the ``num_features`` argument of the :class:`BatchNorm2d` that is inferred
|
||||
from the ``input.size(1)``.
|
||||
The attributes that will be lazily initialized are `weight`, `bias`,
|
||||
`running_mean` and `running_var`.
|
||||
|
||||
Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation
|
||||
on lazy modules and their limitations.
|
||||
|
||||
Args:
|
||||
eps: a value added to the denominator for numerical stability.
|
||||
Default: 1e-5
|
||||
momentum: the value used for the running_mean and running_var
|
||||
computation. Can be set to ``None`` for cumulative moving average
|
||||
(i.e. simple average). Default: 0.1
|
||||
affine: a boolean value that when set to ``True``, this module has
|
||||
learnable affine parameters. Default: ``True``
|
||||
track_running_stats: a boolean value that when set to ``True``, this
|
||||
module tracks the running mean and variance, and when set to ``False``,
|
||||
this module does not track such statistics, and initializes statistics
|
||||
buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
|
||||
When these buffers are ``None``, this module always uses batch statistics.
|
||||
in both training and eval modes. Default: ``True``
|
||||
bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
|
||||
:attr:`affine` is ``True``). Default: ``True``
|
||||
"""
|
||||
|
||||
cls_to_become = BatchNorm2d # type: ignore[assignment]
|
||||
|
||||
def _check_input_dim(self, input) -> None:
|
||||
if input.dim() != 4:
|
||||
raise ValueError(f"expected 4D input (got {input.dim()}D input)")
|
||||
|
||||
|
||||
class BatchNorm3d(_BatchNorm):
|
||||
r"""Applies Batch Normalization over a 5D input.
|
||||
|
||||
5D is a mini-batch of 3D inputs with additional channel dimension as described in the paper
|
||||
`Batch Normalization: Accelerating Deep Network Training by Reducing
|
||||
Internal Covariate Shift <https://arxiv.org/abs/1502.03167>`__ .
|
||||
|
||||
.. math::
|
||||
|
||||
y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta
|
||||
|
||||
The mean and standard-deviation are calculated per-dimension over
|
||||
the mini-batches and :math:`\gamma` and :math:`\beta` are learnable parameter vectors
|
||||
of size `C` (where `C` is the input size). By default, the elements of :math:`\gamma` are set
|
||||
to 1 and the elements of :math:`\beta` are set to 0. At train time in the forward pass, the
|
||||
standard-deviation is calculated via the biased estimator, equivalent to
|
||||
``torch.var(input, correction=0)``. However, the value stored in the moving average of the
|
||||
standard-deviation is calculated via the unbiased estimator, equivalent to
|
||||
``torch.var(input, correction=1)``.
|
||||
|
||||
Also by default, during training this layer keeps running estimates of its
|
||||
computed mean and variance, which are then used for normalization during
|
||||
evaluation. The running estimates are kept with a default :attr:`momentum`
|
||||
of 0.1.
|
||||
|
||||
If :attr:`track_running_stats` is set to ``False``, this layer then does not
|
||||
keep running estimates, and batch statistics are instead used during
|
||||
evaluation time as well.
|
||||
|
||||
.. note::
|
||||
This :attr:`momentum` argument is different from one used in optimizer
|
||||
classes and the conventional notion of momentum. Mathematically, the
|
||||
update rule for running statistics here is
|
||||
:math:`\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t`,
|
||||
where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the
|
||||
new observed value.
|
||||
|
||||
Because the Batch Normalization is done over the `C` dimension, computing statistics
|
||||
on `(N, D, H, W)` slices, it's common terminology to call this Volumetric Batch Normalization
|
||||
or Spatio-temporal Batch Normalization.
|
||||
|
||||
Args:
|
||||
num_features: :math:`C` from an expected input of size
|
||||
:math:`(N, C, D, H, W)`
|
||||
eps: a value added to the denominator for numerical stability.
|
||||
Default: 1e-5
|
||||
momentum: the value used for the running_mean and running_var
|
||||
computation. Can be set to ``None`` for cumulative moving average
|
||||
(i.e. simple average). Default: 0.1
|
||||
affine: a boolean value that when set to ``True``, this module has
|
||||
learnable affine parameters. Default: ``True``
|
||||
track_running_stats: a boolean value that when set to ``True``, this
|
||||
module tracks the running mean and variance, and when set to ``False``,
|
||||
this module does not track such statistics, and initializes statistics
|
||||
buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
|
||||
When these buffers are ``None``, this module always uses batch statistics.
|
||||
in both training and eval modes. Default: ``True``
|
||||
bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
|
||||
:attr:`affine` is ``True``). Default: ``True``
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, D, H, W)`
|
||||
- Output: :math:`(N, C, D, H, W)` (same shape as input)
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # With Learnable Parameters
|
||||
>>> m = nn.BatchNorm3d(100)
|
||||
>>> # Without Learnable Parameters
|
||||
>>> m = nn.BatchNorm3d(100, affine=False)
|
||||
>>> input = torch.randn(20, 100, 35, 45, 10)
|
||||
>>> output = m(input)
|
||||
"""
|
||||
|
||||
def _check_input_dim(self, input) -> None:
|
||||
if input.dim() != 5:
|
||||
raise ValueError(f"expected 5D input (got {input.dim()}D input)")
|
||||
|
||||
|
||||
class LazyBatchNorm3d(_LazyNormBase, _BatchNorm):
|
||||
r"""A :class:`torch.nn.BatchNorm3d` module with lazy initialization.
|
||||
|
||||
Lazy initialization is done for the ``num_features`` argument of the :class:`BatchNorm3d` that is inferred
|
||||
from the ``input.size(1)``.
|
||||
The attributes that will be lazily initialized are `weight`, `bias`,
|
||||
`running_mean` and `running_var`.
|
||||
|
||||
Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation
|
||||
on lazy modules and their limitations.
|
||||
|
||||
Args:
|
||||
eps: a value added to the denominator for numerical stability.
|
||||
Default: 1e-5
|
||||
momentum: the value used for the running_mean and running_var
|
||||
computation. Can be set to ``None`` for cumulative moving average
|
||||
(i.e. simple average). Default: 0.1
|
||||
affine: a boolean value that when set to ``True``, this module has
|
||||
learnable affine parameters. Default: ``True``
|
||||
track_running_stats: a boolean value that when set to ``True``, this
|
||||
module tracks the running mean and variance, and when set to ``False``,
|
||||
this module does not track such statistics, and initializes statistics
|
||||
buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
|
||||
When these buffers are ``None``, this module always uses batch statistics.
|
||||
in both training and eval modes. Default: ``True``
|
||||
bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
|
||||
:attr:`affine` is ``True``). Default: ``True``
|
||||
"""
|
||||
|
||||
cls_to_become = BatchNorm3d # type: ignore[assignment]
|
||||
|
||||
def _check_input_dim(self, input) -> None:
|
||||
if input.dim() != 5:
|
||||
raise ValueError(f"expected 5D input (got {input.dim()}D input)")
|
||||
|
||||
|
||||
class SyncBatchNorm(_BatchNorm):
|
||||
r"""Applies Batch Normalization over a N-Dimensional input.
|
||||
|
||||
The N-D input is a mini-batch of [N-2]D inputs with additional channel dimension) as described in the paper
|
||||
`Batch Normalization: Accelerating Deep Network Training by Reducing
|
||||
Internal Covariate Shift <https://arxiv.org/abs/1502.03167>`__ .
|
||||
|
||||
.. math::
|
||||
|
||||
y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta
|
||||
|
||||
The mean and standard-deviation are calculated per-dimension over all
|
||||
mini-batches of the same process groups. :math:`\gamma` and :math:`\beta`
|
||||
are learnable parameter vectors of size `C` (where `C` is the input size).
|
||||
By default, the elements of :math:`\gamma` are sampled from
|
||||
:math:`\mathcal{U}(0, 1)` and the elements of :math:`\beta` are set to 0.
|
||||
The standard-deviation is calculated via the biased estimator, equivalent to
|
||||
`torch.var(input, correction=0)`.
|
||||
|
||||
Also by default, during training this layer keeps running estimates of its
|
||||
computed mean and variance, which are then used for normalization during
|
||||
evaluation. The running estimates are kept with a default :attr:`momentum`
|
||||
of 0.1.
|
||||
|
||||
If :attr:`track_running_stats` is set to ``False``, this layer then does not
|
||||
keep running estimates, and batch statistics are instead used during
|
||||
evaluation time as well.
|
||||
|
||||
.. note::
|
||||
This :attr:`momentum` argument is different from one used in optimizer
|
||||
classes and the conventional notion of momentum. Mathematically, the
|
||||
update rule for running statistics here is
|
||||
:math:`\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t`,
|
||||
where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the
|
||||
new observed value.
|
||||
|
||||
Because the Batch Normalization is done for each channel in the ``C`` dimension, computing
|
||||
statistics on ``(N, +)`` slices, it's common terminology to call this Volumetric Batch
|
||||
Normalization or Spatio-temporal Batch Normalization.
|
||||
|
||||
Currently :class:`SyncBatchNorm` only supports
|
||||
:class:`~torch.nn.DistributedDataParallel` (DDP) with single GPU per process. Use
|
||||
:meth:`torch.nn.SyncBatchNorm.convert_sync_batchnorm()` to convert
|
||||
:attr:`BatchNorm*D` layer to :class:`SyncBatchNorm` before wrapping
|
||||
Network with DDP.
|
||||
|
||||
Args:
|
||||
num_features: :math:`C` from an expected input of size
|
||||
:math:`(N, C, +)`
|
||||
eps: a value added to the denominator for numerical stability.
|
||||
Default: ``1e-5``
|
||||
momentum: the value used for the running_mean and running_var
|
||||
computation. Can be set to ``None`` for cumulative moving average
|
||||
(i.e. simple average). Default: 0.1
|
||||
affine: a boolean value that when set to ``True``, this module has
|
||||
learnable affine parameters. Default: ``True``
|
||||
track_running_stats: a boolean value that when set to ``True``, this
|
||||
module tracks the running mean and variance, and when set to ``False``,
|
||||
this module does not track such statistics, and initializes statistics
|
||||
buffers :attr:`running_mean` and :attr:`running_var` as ``None``.
|
||||
When these buffers are ``None``, this module always uses batch statistics.
|
||||
in both training and eval modes. Default: ``True``
|
||||
process_group: synchronization of stats happen within each process group
|
||||
individually. Default behavior is synchronization across the whole
|
||||
world
|
||||
bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
|
||||
:attr:`affine` is ``True``). Default: ``True``
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, +)`
|
||||
- Output: :math:`(N, C, +)` (same shape as input)
|
||||
|
||||
.. note::
|
||||
Synchronization of batchnorm statistics occurs only while training, i.e.
|
||||
synchronization is disabled when ``model.eval()`` is set or if
|
||||
``self.training`` is otherwise ``False``.
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> # With Learnable Parameters
|
||||
>>> m = nn.SyncBatchNorm(100)
|
||||
>>> # creating process group (optional)
|
||||
>>> # ranks is a list of int identifying rank ids.
|
||||
>>> ranks = list(range(8))
|
||||
>>> r1, r2 = ranks[:4], ranks[4:]
|
||||
>>> # Note: every rank calls into new_group for every
|
||||
>>> # process group created, even if that rank is not
|
||||
>>> # part of the group.
|
||||
>>> process_groups = [torch.distributed.new_group(pids) for pids in [r1, r2]]
|
||||
>>> process_group = process_groups[0 if dist.get_rank() <= 3 else 1]
|
||||
>>> # Without Learnable Parameters
|
||||
>>> m = nn.BatchNorm3d(100, affine=False, process_group=process_group)
|
||||
>>> input = torch.randn(20, 100, 35, 45, 10)
|
||||
>>> output = m(input)
|
||||
|
||||
>>> # network is nn.BatchNorm layer
|
||||
>>> sync_bn_network = nn.SyncBatchNorm.convert_sync_batchnorm(network, process_group)
|
||||
>>> # only single gpu per process is currently supported
|
||||
>>> ddp_sync_bn_network = torch.nn.parallel.DistributedDataParallel(
|
||||
>>> sync_bn_network,
|
||||
>>> device_ids=[args.local_rank],
|
||||
>>> output_device=args.local_rank)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_features: int,
|
||||
eps: float = 1e-5,
|
||||
momentum: float | None = 0.1,
|
||||
affine: bool = True,
|
||||
track_running_stats: bool = True,
|
||||
process_group: Any | None = None,
|
||||
device=None,
|
||||
dtype=None,
|
||||
*,
|
||||
bias: bool = True,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__(
|
||||
num_features,
|
||||
eps,
|
||||
momentum,
|
||||
affine,
|
||||
track_running_stats,
|
||||
**factory_kwargs,
|
||||
bias=bias,
|
||||
)
|
||||
self.process_group = process_group
|
||||
|
||||
def _check_input_dim(self, input) -> None:
|
||||
if input.dim() < 2:
|
||||
raise ValueError(f"expected at least 2D input (got {input.dim()}D input)")
|
||||
|
||||
def _check_non_zero_input_channels(self, input) -> None:
|
||||
if input.size(1) == 0:
|
||||
raise ValueError(
|
||||
"SyncBatchNorm number of input channels should be non-zero"
|
||||
)
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
self._check_input_dim(input)
|
||||
self._check_non_zero_input_channels(input)
|
||||
|
||||
# exponential_average_factor is set to self.momentum
|
||||
# (when it is available) only so that it gets updated
|
||||
# in ONNX graph when this node is exported to ONNX.
|
||||
if self.momentum is None:
|
||||
exponential_average_factor = 0.0
|
||||
else:
|
||||
exponential_average_factor = self.momentum
|
||||
|
||||
if self.training and self.track_running_stats:
|
||||
if self.num_batches_tracked is None:
|
||||
raise AssertionError("num_batches_tracked must not be None")
|
||||
self.num_batches_tracked.add_(1)
|
||||
if self.momentum is None: # use cumulative moving average
|
||||
exponential_average_factor = 1.0 / self.num_batches_tracked.item()
|
||||
else: # use exponential moving average
|
||||
exponential_average_factor = self.momentum
|
||||
|
||||
r"""
|
||||
Decide whether the mini-batch stats should be used for normalization rather than the buffers.
|
||||
Mini-batch stats are used in training mode, and in eval mode when buffers are None.
|
||||
"""
|
||||
if self.training:
|
||||
bn_training = True
|
||||
else:
|
||||
bn_training = (self.running_mean is None) and (self.running_var is None)
|
||||
|
||||
r"""
|
||||
Buffers are only updated if they are to be tracked and we are in training mode. Thus they only need to be
|
||||
passed when the update should occur (i.e. in training mode when they are tracked), or when buffer stats are
|
||||
used for normalization (i.e. in eval mode when buffers are not None).
|
||||
"""
|
||||
# If buffers are not to be tracked, ensure that they won't be updated
|
||||
running_mean = (
|
||||
self.running_mean if not self.training or self.track_running_stats else None
|
||||
)
|
||||
running_var = (
|
||||
self.running_var if not self.training or self.track_running_stats else None
|
||||
)
|
||||
|
||||
# Don't sync batchnorm stats in inference mode (model.eval()).
|
||||
need_sync = (
|
||||
bn_training
|
||||
and self.training
|
||||
and torch.distributed.is_available()
|
||||
and torch.distributed.is_initialized()
|
||||
)
|
||||
if need_sync:
|
||||
# currently only GPU/PrivateUse1 input is supported
|
||||
if input.device.type not in [
|
||||
"cuda",
|
||||
"hpu",
|
||||
"xpu",
|
||||
torch._C._get_privateuse1_backend_name(),
|
||||
]:
|
||||
raise ValueError(
|
||||
"SyncBatchNorm expected input tensor to be on GPU or XPU or "
|
||||
f"{torch._C._get_privateuse1_backend_name()}"
|
||||
)
|
||||
|
||||
process_group = torch.distributed.group.WORLD
|
||||
if self.process_group:
|
||||
process_group = self.process_group
|
||||
world_size = torch.distributed.get_world_size(process_group)
|
||||
need_sync = world_size > 1
|
||||
|
||||
# fallback to framework BN when synchronization is not necessary
|
||||
if not need_sync:
|
||||
return F.batch_norm(
|
||||
input,
|
||||
running_mean,
|
||||
running_var,
|
||||
self.weight,
|
||||
self.bias,
|
||||
bn_training,
|
||||
exponential_average_factor,
|
||||
self.eps,
|
||||
)
|
||||
else:
|
||||
if not bn_training:
|
||||
raise AssertionError("bn_training must be True")
|
||||
return sync_batch_norm.apply(
|
||||
input,
|
||||
self.weight,
|
||||
self.bias,
|
||||
running_mean,
|
||||
running_var,
|
||||
self.eps,
|
||||
exponential_average_factor,
|
||||
process_group, # type: ignore[possibly-undefined]
|
||||
world_size, # type: ignore[possibly-undefined]
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def convert_sync_batchnorm(cls, module, process_group=None):
|
||||
r"""Converts all :attr:`BatchNorm*D` layers in the model to :class:`torch.nn.SyncBatchNorm` layers.
|
||||
|
||||
Args:
|
||||
module (nn.Module): module containing one or more :attr:`BatchNorm*D` layers
|
||||
process_group (optional): process group to scope synchronization,
|
||||
default is the whole world
|
||||
|
||||
Returns:
|
||||
The original :attr:`module` with the converted :class:`torch.nn.SyncBatchNorm`
|
||||
layers. If the original :attr:`module` is a :attr:`BatchNorm*D` layer,
|
||||
a new :class:`torch.nn.SyncBatchNorm` layer object will be returned
|
||||
instead.
|
||||
|
||||
Example::
|
||||
|
||||
>>> # Network with nn.BatchNorm layer
|
||||
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
|
||||
>>> module = torch.nn.Sequential(
|
||||
>>> torch.nn.Linear(20, 100),
|
||||
>>> torch.nn.BatchNorm1d(100),
|
||||
>>> ).cuda()
|
||||
>>> # creating process group (optional)
|
||||
>>> # ranks is a list of int identifying rank ids.
|
||||
>>> ranks = list(range(8))
|
||||
>>> r1, r2 = ranks[:4], ranks[4:]
|
||||
>>> # Note: every rank calls into new_group for every
|
||||
>>> # process group created, even if that rank is not
|
||||
>>> # part of the group.
|
||||
>>> # xdoctest: +SKIP("distributed")
|
||||
>>> process_groups = [torch.distributed.new_group(pids) for pids in [r1, r2]]
|
||||
>>> process_group = process_groups[0 if dist.get_rank() <= 3 else 1]
|
||||
>>> sync_bn_module = torch.nn.SyncBatchNorm.convert_sync_batchnorm(module, process_group)
|
||||
|
||||
"""
|
||||
module_output = module
|
||||
if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
|
||||
module_output = torch.nn.SyncBatchNorm(
|
||||
module.num_features,
|
||||
module.eps,
|
||||
module.momentum,
|
||||
module.affine,
|
||||
module.track_running_stats,
|
||||
process_group,
|
||||
bias=module.bias is not None,
|
||||
)
|
||||
if module.affine:
|
||||
with torch.no_grad():
|
||||
module_output.weight = module.weight
|
||||
module_output.bias = module.bias
|
||||
module_output.running_mean = module.running_mean
|
||||
module_output.running_var = module.running_var
|
||||
module_output.num_batches_tracked = module.num_batches_tracked
|
||||
module_output.training = module.training
|
||||
if hasattr(module, "qconfig"):
|
||||
module_output.qconfig = module.qconfig
|
||||
for name, child in module.named_children():
|
||||
module_output.add_module(
|
||||
name, cls.convert_sync_batchnorm(child, process_group)
|
||||
)
|
||||
del module
|
||||
return module_output
|
||||
@@ -0,0 +1,62 @@
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
|
||||
from .module import Module
|
||||
|
||||
|
||||
__all__ = ["ChannelShuffle"]
|
||||
|
||||
|
||||
class ChannelShuffle(Module):
|
||||
r"""Divides and rearranges the channels in a tensor.
|
||||
|
||||
This operation divides the channels in a tensor of shape :math:`(N, C, *)`
|
||||
into g groups as :math:`(N, \frac{C}{g}, g, *)` and shuffles them,
|
||||
while retaining the original tensor shape in the final output.
|
||||
|
||||
Args:
|
||||
groups (int): number of groups to divide channels in.
|
||||
|
||||
Examples::
|
||||
|
||||
>>> channel_shuffle = nn.ChannelShuffle(2)
|
||||
>>> input = torch.arange(1, 17, dtype=torch.float32).view(1, 4, 2, 2)
|
||||
>>> input
|
||||
tensor([[[[ 1., 2.],
|
||||
[ 3., 4.]],
|
||||
[[ 5., 6.],
|
||||
[ 7., 8.]],
|
||||
[[ 9., 10.],
|
||||
[11., 12.]],
|
||||
[[13., 14.],
|
||||
[15., 16.]]]])
|
||||
>>> output = channel_shuffle(input)
|
||||
>>> output
|
||||
tensor([[[[ 1., 2.],
|
||||
[ 3., 4.]],
|
||||
[[ 9., 10.],
|
||||
[11., 12.]],
|
||||
[[ 5., 6.],
|
||||
[ 7., 8.]],
|
||||
[[13., 14.],
|
||||
[15., 16.]]]])
|
||||
"""
|
||||
|
||||
__constants__ = ["groups"]
|
||||
groups: int
|
||||
|
||||
def __init__(self, groups: int) -> None:
|
||||
super().__init__()
|
||||
self.groups = groups
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
return F.channel_shuffle(input, self.groups)
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
"""
|
||||
Return the extra representation of the module.
|
||||
"""
|
||||
return f"groups={self.groups}"
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,100 @@
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
|
||||
from .module import Module
|
||||
|
||||
|
||||
__all__ = ["PairwiseDistance", "CosineSimilarity"]
|
||||
|
||||
|
||||
class PairwiseDistance(Module):
|
||||
r"""
|
||||
Computes the pairwise distance between input vectors, or between columns of input matrices.
|
||||
|
||||
Distances are computed using ``p``-norm, with constant ``eps`` added to avoid division by zero
|
||||
if ``p`` is negative, i.e.:
|
||||
|
||||
.. math ::
|
||||
\mathrm{dist}\left(x, y\right) = \left\Vert x-y + \epsilon e \right\Vert_p,
|
||||
|
||||
where :math:`e` is the vector of ones and the ``p``-norm is given by.
|
||||
|
||||
.. math ::
|
||||
\Vert x \Vert _p = \left( \sum_{i=1}^n \vert x_i \vert ^ p \right) ^ {1/p}.
|
||||
|
||||
Args:
|
||||
p (real, optional): the norm degree. Can be negative. Default: 2
|
||||
eps (float, optional): Small value to avoid division by zero.
|
||||
Default: 1e-6
|
||||
keepdim (bool, optional): Determines whether or not to keep the vector dimension.
|
||||
Default: False
|
||||
Shape:
|
||||
- Input1: :math:`(N, D)` or :math:`(D)` where `N = batch dimension` and `D = vector dimension`
|
||||
- Input2: :math:`(N, D)` or :math:`(D)`, same shape as the Input1
|
||||
- Output: :math:`(N)` or :math:`()` based on input dimension.
|
||||
If :attr:`keepdim` is ``True``, then :math:`(N, 1)` or :math:`(1)` based on input dimension.
|
||||
|
||||
Examples:
|
||||
>>> pdist = nn.PairwiseDistance(p=2)
|
||||
>>> input1 = torch.randn(100, 128)
|
||||
>>> input2 = torch.randn(100, 128)
|
||||
>>> output = pdist(input1, input2)
|
||||
"""
|
||||
|
||||
__constants__ = ["norm", "eps", "keepdim"]
|
||||
norm: float
|
||||
eps: float
|
||||
keepdim: bool
|
||||
|
||||
def __init__(
|
||||
self, p: float = 2.0, eps: float = 1e-6, keepdim: bool = False
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.norm = p
|
||||
self.eps = eps
|
||||
self.keepdim = keepdim
|
||||
|
||||
def forward(self, x1: Tensor, x2: Tensor) -> Tensor:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
return F.pairwise_distance(x1, x2, self.norm, self.eps, self.keepdim)
|
||||
|
||||
|
||||
class CosineSimilarity(Module):
|
||||
r"""Returns cosine similarity between :math:`x_1` and :math:`x_2`, computed along `dim`.
|
||||
|
||||
.. math ::
|
||||
\text{similarity} = \dfrac{x_1 \cdot x_2}{\max(\Vert x_1 \Vert _2 \cdot \Vert x_2 \Vert _2, \epsilon)}.
|
||||
|
||||
Args:
|
||||
dim (int, optional): Dimension where cosine similarity is computed. Default: 1
|
||||
eps (float, optional): Small value to avoid division by zero.
|
||||
Default: 1e-8
|
||||
Shape:
|
||||
- Input1: :math:`(\ast_1, D, \ast_2)` where D is at position `dim`
|
||||
- Input2: :math:`(\ast_1, D, \ast_2)`, same number of dimensions as x1, matching x1 size at dimension `dim`,
|
||||
and broadcastable with x1 at other dimensions.
|
||||
- Output: :math:`(\ast_1, \ast_2)`
|
||||
|
||||
Examples:
|
||||
>>> input1 = torch.randn(100, 128)
|
||||
>>> input2 = torch.randn(100, 128)
|
||||
>>> cos = nn.CosineSimilarity(dim=1, eps=1e-6)
|
||||
>>> output = cos(input1, input2)
|
||||
"""
|
||||
|
||||
__constants__ = ["dim", "eps"]
|
||||
dim: int
|
||||
eps: float
|
||||
|
||||
def __init__(self, dim: int = 1, eps: float = 1e-8) -> None:
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, x1: Tensor, x2: Tensor) -> Tensor:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
return F.cosine_similarity(x1, x2, self.dim, self.eps)
|
||||
@@ -0,0 +1,323 @@
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
|
||||
from .module import Module
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Dropout",
|
||||
"Dropout1d",
|
||||
"Dropout2d",
|
||||
"Dropout3d",
|
||||
"AlphaDropout",
|
||||
"FeatureAlphaDropout",
|
||||
]
|
||||
|
||||
|
||||
class _DropoutNd(Module):
|
||||
__constants__ = ["p", "inplace"]
|
||||
p: float
|
||||
inplace: bool
|
||||
|
||||
def __init__(self, p: float = 0.5, inplace: bool = False) -> None:
|
||||
super().__init__()
|
||||
if p < 0 or p > 1:
|
||||
raise ValueError(
|
||||
f"dropout probability has to be between 0 and 1, but got {p}"
|
||||
)
|
||||
self.p = p
|
||||
self.inplace = inplace
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
return f"p={self.p}, inplace={self.inplace}"
|
||||
|
||||
|
||||
class Dropout(_DropoutNd):
|
||||
r"""During training, randomly zeroes some of the elements of the input tensor with probability :attr:`p`.
|
||||
|
||||
The zeroed elements are chosen independently for each forward call and are sampled from a Bernoulli distribution.
|
||||
|
||||
Each channel will be zeroed out independently on every forward call.
|
||||
|
||||
This has proven to be an effective technique for regularization and
|
||||
preventing the co-adaptation of neurons as described in the paper
|
||||
`Improving neural networks by preventing co-adaptation of feature
|
||||
detectors`_ .
|
||||
|
||||
Furthermore, the outputs are scaled by a factor of :math:`\frac{1}{1-p}` during
|
||||
training. This means that during evaluation the module simply computes an
|
||||
identity function.
|
||||
|
||||
Args:
|
||||
p: probability of an element to be zeroed. Default: 0.5
|
||||
inplace: If set to ``True``, will do this operation in-place. Default: ``False``
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(*)`. Input can be of any shape
|
||||
- Output: :math:`(*)`. Output is of the same shape as input
|
||||
|
||||
Examples::
|
||||
|
||||
>>> m = nn.Dropout(p=0.2)
|
||||
>>> input = torch.randn(20, 16)
|
||||
>>> output = m(input)
|
||||
|
||||
.. _Improving neural networks by preventing co-adaptation of feature
|
||||
detectors: https://arxiv.org/abs/1207.0580
|
||||
"""
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
return F.dropout(input, self.p, self.training, self.inplace)
|
||||
|
||||
|
||||
class Dropout1d(_DropoutNd):
|
||||
r"""Randomly zero out entire channels.
|
||||
|
||||
A channel is a 1D feature map,
|
||||
e.g., the :math:`j`-th channel of the :math:`i`-th sample in the
|
||||
batched input is a 1D tensor :math:`\text{input}[i, j]`.
|
||||
|
||||
Each channel will be zeroed out independently on every forward call with
|
||||
probability :attr:`p` using samples from a Bernoulli distribution.
|
||||
|
||||
Usually the input comes from :class:`nn.Conv1d` modules.
|
||||
|
||||
As described in the paper
|
||||
`Efficient Object Localization Using Convolutional Networks`_ ,
|
||||
if adjacent pixels within feature maps are strongly correlated
|
||||
(as is normally the case in early convolution layers) then i.i.d. dropout
|
||||
will not regularize the activations and will otherwise just result
|
||||
in an effective learning rate decrease.
|
||||
|
||||
In this case, :func:`nn.Dropout1d` will help promote independence between
|
||||
feature maps and should be used instead.
|
||||
|
||||
Args:
|
||||
p (float, optional): probability of an element to be zero-ed.
|
||||
inplace (bool, optional): If set to ``True``, will do this operation
|
||||
in-place
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, L)` or :math:`(C, L)`.
|
||||
- Output: :math:`(N, C, L)` or :math:`(C, L)` (same shape as input).
|
||||
|
||||
Examples::
|
||||
|
||||
>>> m = nn.Dropout1d(p=0.2)
|
||||
>>> input = torch.randn(20, 16, 32)
|
||||
>>> output = m(input)
|
||||
|
||||
.. _Efficient Object Localization Using Convolutional Networks:
|
||||
https://arxiv.org/abs/1411.4280
|
||||
"""
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
return F.dropout1d(input, self.p, self.training, self.inplace)
|
||||
|
||||
|
||||
class Dropout2d(_DropoutNd):
|
||||
r"""Randomly zero out entire channels.
|
||||
|
||||
A channel is a 2D feature map,
|
||||
e.g., the :math:`j`-th channel of the :math:`i`-th sample in the
|
||||
batched input is a 2D tensor :math:`\text{input}[i, j]`.
|
||||
|
||||
Each channel will be zeroed out independently on every forward call with
|
||||
probability :attr:`p` using samples from a Bernoulli distribution.
|
||||
|
||||
Usually the input comes from :class:`nn.Conv2d` modules.
|
||||
|
||||
As described in the paper
|
||||
`Efficient Object Localization Using Convolutional Networks`_ ,
|
||||
if adjacent pixels within feature maps are strongly correlated
|
||||
(as is normally the case in early convolution layers) then i.i.d. dropout
|
||||
will not regularize the activations and will otherwise just result
|
||||
in an effective learning rate decrease.
|
||||
|
||||
In this case, :func:`nn.Dropout2d` will help promote independence between
|
||||
feature maps and should be used instead.
|
||||
|
||||
Args:
|
||||
p (float, optional): probability of an element to be zero-ed.
|
||||
inplace (bool, optional): If set to ``True``, will do this operation
|
||||
in-place
|
||||
|
||||
.. warning ::
|
||||
Due to historical reasons, this class will perform 1D channel-wise dropout
|
||||
for 3D inputs (as done by :class:`nn.Dropout1d`). Thus, it currently does NOT
|
||||
support inputs without a batch dimension of shape :math:`(C, H, W)`. This
|
||||
behavior will change in a future release to interpret 3D inputs as no-batch-dim
|
||||
inputs. To maintain the old behavior, switch to :class:`nn.Dropout1d`.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, H, W)` or :math:`(N, C, L)`.
|
||||
- Output: :math:`(N, C, H, W)` or :math:`(N, C, L)` (same shape as input).
|
||||
|
||||
Examples::
|
||||
|
||||
>>> m = nn.Dropout2d(p=0.2)
|
||||
>>> input = torch.randn(20, 16, 32, 32)
|
||||
>>> output = m(input)
|
||||
|
||||
.. _Efficient Object Localization Using Convolutional Networks:
|
||||
https://arxiv.org/abs/1411.4280
|
||||
"""
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
return F.dropout2d(input, self.p, self.training, self.inplace)
|
||||
|
||||
|
||||
class Dropout3d(_DropoutNd):
|
||||
r"""Randomly zero out entire channels.
|
||||
|
||||
A channel is a 3D feature map,
|
||||
e.g., the :math:`j`-th channel of the :math:`i`-th sample in the
|
||||
batched input is a 3D tensor :math:`\text{input}[i, j]`.
|
||||
|
||||
Each channel will be zeroed out independently on every forward call with
|
||||
probability :attr:`p` using samples from a Bernoulli distribution.
|
||||
|
||||
Usually the input comes from :class:`nn.Conv3d` modules.
|
||||
|
||||
As described in the paper
|
||||
`Efficient Object Localization Using Convolutional Networks`_ ,
|
||||
if adjacent pixels within feature maps are strongly correlated
|
||||
(as is normally the case in early convolution layers) then i.i.d. dropout
|
||||
will not regularize the activations and will otherwise just result
|
||||
in an effective learning rate decrease.
|
||||
|
||||
In this case, :func:`nn.Dropout3d` will help promote independence between
|
||||
feature maps and should be used instead.
|
||||
|
||||
Args:
|
||||
p (float, optional): probability of an element to be zeroed.
|
||||
inplace (bool, optional): If set to ``True``, will do this operation
|
||||
in-place
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)`.
|
||||
- Output: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)` (same shape as input).
|
||||
|
||||
Examples::
|
||||
|
||||
>>> m = nn.Dropout3d(p=0.2)
|
||||
>>> input = torch.randn(20, 16, 4, 32, 32)
|
||||
>>> output = m(input)
|
||||
|
||||
.. _Efficient Object Localization Using Convolutional Networks:
|
||||
https://arxiv.org/abs/1411.4280
|
||||
"""
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
return F.dropout3d(input, self.p, self.training, self.inplace)
|
||||
|
||||
|
||||
class AlphaDropout(_DropoutNd):
|
||||
r"""Applies Alpha Dropout over the input.
|
||||
|
||||
Alpha Dropout is a type of Dropout that maintains the self-normalizing
|
||||
property.
|
||||
For an input with zero mean and unit standard deviation, the output of
|
||||
Alpha Dropout maintains the original mean and standard deviation of the
|
||||
input.
|
||||
Alpha Dropout goes hand-in-hand with SELU activation function, which ensures
|
||||
that the outputs have zero mean and unit standard deviation.
|
||||
|
||||
During training, it randomly masks some of the elements of the input
|
||||
tensor with probability *p* using samples from a bernoulli distribution.
|
||||
The elements to masked are randomized on every forward call, and scaled
|
||||
and shifted to maintain zero mean and unit standard deviation.
|
||||
|
||||
During evaluation the module simply computes an identity function.
|
||||
|
||||
More details can be found in the paper `Self-Normalizing Neural Networks`_ .
|
||||
|
||||
Args:
|
||||
p (float): probability of an element to be dropped. Default: 0.5
|
||||
inplace (bool, optional): If set to ``True``, will do this operation
|
||||
in-place
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(*)`. Input can be of any shape
|
||||
- Output: :math:`(*)`. Output is of the same shape as input
|
||||
|
||||
Examples::
|
||||
|
||||
>>> m = nn.AlphaDropout(p=0.2)
|
||||
>>> input = torch.randn(20, 16)
|
||||
>>> output = m(input)
|
||||
|
||||
.. _Self-Normalizing Neural Networks: https://arxiv.org/abs/1706.02515
|
||||
"""
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
return F.alpha_dropout(input, self.p, self.training)
|
||||
|
||||
|
||||
class FeatureAlphaDropout(_DropoutNd):
|
||||
r"""Randomly masks out entire channels.
|
||||
|
||||
A channel is a feature map,
|
||||
e.g. the :math:`j`-th channel of the :math:`i`-th sample in the batch input
|
||||
is a tensor :math:`\text{input}[i, j]` of the input tensor). Instead of
|
||||
setting activations to zero, as in regular Dropout, the activations are set
|
||||
to the negative saturation value of the SELU activation function. More details
|
||||
can be found in the paper `Self-Normalizing Neural Networks`_ .
|
||||
|
||||
Each element will be masked independently for each sample on every forward
|
||||
call with probability :attr:`p` using samples from a Bernoulli distribution.
|
||||
The elements to be masked are randomized on every forward call, and scaled
|
||||
and shifted to maintain zero mean and unit variance.
|
||||
|
||||
Usually the input comes from :class:`nn.AlphaDropout` modules.
|
||||
|
||||
As described in the paper
|
||||
`Efficient Object Localization Using Convolutional Networks`_ ,
|
||||
if adjacent pixels within feature maps are strongly correlated
|
||||
(as is normally the case in early convolution layers) then i.i.d. dropout
|
||||
will not regularize the activations and will otherwise just result
|
||||
in an effective learning rate decrease.
|
||||
|
||||
In this case, :func:`nn.AlphaDropout` will help promote independence between
|
||||
feature maps and should be used instead.
|
||||
|
||||
Args:
|
||||
p (float, optional): probability of an element to be zeroed. Default: 0.5
|
||||
inplace (bool, optional): If set to ``True``, will do this operation
|
||||
in-place
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)`.
|
||||
- Output: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)` (same shape as input).
|
||||
|
||||
Examples::
|
||||
|
||||
>>> m = nn.FeatureAlphaDropout(p=0.2)
|
||||
>>> input = torch.randn(20, 16, 4, 32, 32)
|
||||
>>> output = m(input)
|
||||
|
||||
.. _Self-Normalizing Neural Networks: https://arxiv.org/abs/1706.02515
|
||||
.. _Efficient Object Localization Using Convolutional Networks:
|
||||
https://arxiv.org/abs/1411.4280
|
||||
"""
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
return F.feature_alpha_dropout(input, self.p, self.training)
|
||||
@@ -0,0 +1,167 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
from torch import Tensor
|
||||
from torch.types import _size
|
||||
|
||||
from .module import Module
|
||||
|
||||
|
||||
__all__ = ["Flatten", "Unflatten"]
|
||||
|
||||
|
||||
class Flatten(Module):
|
||||
r"""
|
||||
Flattens a contiguous range of dims into a tensor.
|
||||
|
||||
For use with :class:`~nn.Sequential`, see :meth:`torch.flatten` for details.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(*, S_{\text{start}},..., S_{i}, ..., S_{\text{end}}, *)`,'
|
||||
where :math:`S_{i}` is the size at dimension :math:`i` and :math:`*` means any
|
||||
number of dimensions including none.
|
||||
- Output: :math:`(*, \prod_{i=\text{start}}^{\text{end}} S_{i}, *)`.
|
||||
|
||||
Args:
|
||||
start_dim: first dim to flatten (default = 1).
|
||||
end_dim: last dim to flatten (default = -1).
|
||||
|
||||
Examples::
|
||||
>>> input = torch.randn(32, 1, 5, 5)
|
||||
>>> # With default parameters
|
||||
>>> m = nn.Flatten()
|
||||
>>> output = m(input)
|
||||
>>> output.size()
|
||||
torch.Size([32, 25])
|
||||
>>> # With non-default parameters
|
||||
>>> m = nn.Flatten(0, 2)
|
||||
>>> output = m(input)
|
||||
>>> output.size()
|
||||
torch.Size([160, 5])
|
||||
"""
|
||||
|
||||
__constants__ = ["start_dim", "end_dim"]
|
||||
start_dim: int
|
||||
end_dim: int
|
||||
|
||||
def __init__(self, start_dim: int = 1, end_dim: int = -1) -> None:
|
||||
super().__init__()
|
||||
self.start_dim = start_dim
|
||||
self.end_dim = end_dim
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
return input.flatten(self.start_dim, self.end_dim)
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
"""
|
||||
Returns the extra representation of the module.
|
||||
"""
|
||||
return f"start_dim={self.start_dim}, end_dim={self.end_dim}"
|
||||
|
||||
|
||||
class Unflatten(Module):
|
||||
r"""
|
||||
Unflattens a tensor dim expanding it to a desired shape. For use with :class:`~nn.Sequential`.
|
||||
|
||||
* :attr:`dim` specifies the dimension of the input tensor to be unflattened, and it can
|
||||
be either `int` or `str` when `Tensor` or `NamedTensor` is used, respectively.
|
||||
|
||||
* :attr:`unflattened_size` is the new shape of the unflattened dimension of the tensor and it can be
|
||||
a `tuple` of ints or a `list` of ints or `torch.Size` for `Tensor` input; a `NamedShape`
|
||||
(tuple of `(name, size)` tuples) for `NamedTensor` input.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(*, S_{\text{dim}}, *)`, where :math:`S_{\text{dim}}` is the size at
|
||||
dimension :attr:`dim` and :math:`*` means any number of dimensions including none.
|
||||
- Output: :math:`(*, U_1, ..., U_n, *)`, where :math:`U` = :attr:`unflattened_size` and
|
||||
:math:`\prod_{i=1}^n U_i = S_{\text{dim}}`.
|
||||
|
||||
Args:
|
||||
dim (Union[int, str]): Dimension to be unflattened
|
||||
unflattened_size (Union[torch.Size, Tuple, List, NamedShape]): New shape of the unflattened dimension
|
||||
|
||||
Examples:
|
||||
>>> input = torch.randn(2, 50)
|
||||
>>> # With tuple of ints
|
||||
>>> m = nn.Sequential(
|
||||
>>> nn.Linear(50, 50),
|
||||
>>> nn.Unflatten(1, (2, 5, 5))
|
||||
>>> )
|
||||
>>> output = m(input)
|
||||
>>> output.size()
|
||||
torch.Size([2, 2, 5, 5])
|
||||
>>> # With torch.Size
|
||||
>>> m = nn.Sequential(
|
||||
>>> nn.Linear(50, 50),
|
||||
>>> nn.Unflatten(1, torch.Size([2, 5, 5]))
|
||||
>>> )
|
||||
>>> output = m(input)
|
||||
>>> output.size()
|
||||
torch.Size([2, 2, 5, 5])
|
||||
>>> # With namedshape (tuple of tuples)
|
||||
>>> input = torch.randn(2, 50, names=("N", "features"))
|
||||
>>> unflatten = nn.Unflatten("features", (("C", 2), ("H", 5), ("W", 5)))
|
||||
>>> output = unflatten(input)
|
||||
>>> output.size()
|
||||
torch.Size([2, 2, 5, 5])
|
||||
"""
|
||||
|
||||
NamedShape = tuple[tuple[str, int]]
|
||||
|
||||
__constants__ = ["dim", "unflattened_size"]
|
||||
dim: int | str
|
||||
unflattened_size: _size | NamedShape
|
||||
|
||||
def __init__(self, dim: int | str, unflattened_size: _size | NamedShape) -> None:
|
||||
super().__init__()
|
||||
|
||||
if isinstance(dim, int):
|
||||
self._require_tuple_int(unflattened_size)
|
||||
elif isinstance(dim, str):
|
||||
self._require_tuple_tuple(unflattened_size)
|
||||
else:
|
||||
raise TypeError("invalid argument type for dim parameter")
|
||||
|
||||
self.dim = dim
|
||||
self.unflattened_size = unflattened_size
|
||||
|
||||
def _require_tuple_tuple(self, input) -> None:
|
||||
if isinstance(input, tuple):
|
||||
for idx, elem in enumerate(input):
|
||||
if not isinstance(elem, tuple):
|
||||
raise TypeError(
|
||||
"unflattened_size must be tuple of tuples, "
|
||||
+ f"but found element of type {type(elem).__name__} at pos {idx}"
|
||||
)
|
||||
return
|
||||
raise TypeError(
|
||||
"unflattened_size must be a tuple of tuples, "
|
||||
+ f"but found type {type(input).__name__}"
|
||||
)
|
||||
|
||||
def _require_tuple_int(self, input) -> None:
|
||||
if isinstance(input, (tuple, list)):
|
||||
for idx, elem in enumerate(input):
|
||||
if not isinstance(elem, int):
|
||||
raise TypeError(
|
||||
"unflattened_size must be tuple of ints, "
|
||||
+ f"but found element of type {type(elem).__name__} at pos {idx}"
|
||||
)
|
||||
return
|
||||
raise TypeError(
|
||||
f"unflattened_size must be a tuple of ints, but found type {type(input).__name__}"
|
||||
)
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
return input.unflatten(self.dim, self.unflattened_size)
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
"""
|
||||
Returns the extra representation of the module.
|
||||
"""
|
||||
return f"dim={self.dim}, unflattened_size={self.unflattened_size}"
|
||||
@@ -0,0 +1,335 @@
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
from torch.nn.common_types import _size_any_t
|
||||
|
||||
from .module import Module
|
||||
|
||||
|
||||
__all__ = ["Fold", "Unfold"]
|
||||
|
||||
|
||||
class Fold(Module):
|
||||
(
|
||||
r"""Combines an array of sliding local blocks into a large containing tensor.
|
||||
|
||||
Consider a batched :attr:`input` tensor containing sliding local blocks,
|
||||
e.g., patches of images, of shape :math:`(N, C \times \prod(\text{kernel\_size}), L)`,
|
||||
where :math:`N` is batch dimension, :math:`C \times \prod(\text{kernel\_size})`
|
||||
is the number of values within a block (a block has :math:`\prod(\text{kernel\_size})`
|
||||
spatial locations each containing a :math:`C`-channeled vector), and
|
||||
:math:`L` is the total number of blocks. (This is exactly the
|
||||
same specification as the output shape of :class:`~torch.nn.Unfold`.) This
|
||||
operation combines these local blocks into the large :attr:`output` tensor
|
||||
of shape :math:`(N, C, \text{output\_size}[0], \text{output\_size}[1], \dots)`
|
||||
by summing the overlapping values. Similar to :class:`~torch.nn.Unfold`, the
|
||||
arguments must satisfy
|
||||
|
||||
.. math::
|
||||
L = \prod_d \left\lfloor\frac{\text{output\_size}[d] + 2 \times \text{padding}[d] %
|
||||
- \text{dilation}[d] \times (\text{kernel\_size}[d] - 1) - 1}{\text{stride}[d]} + 1\right\rfloor,
|
||||
|
||||
where :math:`d` is over all spatial dimensions.
|
||||
|
||||
* :attr:`output_size` describes the spatial shape of the large containing
|
||||
tensor of the sliding local blocks. It is useful to resolve the ambiguity
|
||||
when multiple input shapes map to same number of sliding blocks, e.g.,
|
||||
with ``stride > 0``.
|
||||
|
||||
The :attr:`padding`, :attr:`stride` and :attr:`dilation` arguments specify
|
||||
how the sliding blocks are retrieved.
|
||||
|
||||
* :attr:`stride` controls the stride for the sliding blocks.
|
||||
|
||||
* :attr:`padding` controls the amount of implicit zero-paddings on both
|
||||
sides for :attr:`padding` number of points for each dimension before
|
||||
reshaping.
|
||||
"""
|
||||
"""
|
||||
* :attr:`dilation` controls the spacing between the kernel points; also known as the \u00e0 trous algorithm.
|
||||
It is harder to describe, but this `link`_ has a nice visualization of what :attr:`dilation` does.
|
||||
"""
|
||||
r"""
|
||||
Args:
|
||||
output_size (int or tuple): the shape of the spatial dimensions of the
|
||||
output (i.e., ``output.sizes()[2:]``)
|
||||
kernel_size (int or tuple): the size of the sliding blocks
|
||||
dilation (int or tuple, optional): a parameter that controls the
|
||||
stride of elements within the
|
||||
neighborhood. Default: 1
|
||||
padding (int or tuple, optional): implicit zero padding to be added on
|
||||
both sides of input. Default: 0
|
||||
stride (int or tuple): the stride of the sliding blocks in the input
|
||||
spatial dimensions. Default: 1
|
||||
|
||||
* If :attr:`output_size`, :attr:`kernel_size`, :attr:`dilation`,
|
||||
:attr:`padding` or :attr:`stride` is an int or a tuple of length 1 then
|
||||
their values will be replicated across all spatial dimensions.
|
||||
|
||||
* For the case of two output spatial dimensions this operation is sometimes
|
||||
called ``col2im``.
|
||||
|
||||
.. note::
|
||||
:class:`~torch.nn.Fold` calculates each combined value in the resulting
|
||||
large tensor by summing all values from all containing blocks.
|
||||
:class:`~torch.nn.Unfold` extracts the values in the local blocks by
|
||||
copying from the large tensor. So, if the blocks overlap, they are not
|
||||
inverses of each other.
|
||||
|
||||
In general, folding and unfolding operations are related as
|
||||
follows. Consider :class:`~torch.nn.Fold` and
|
||||
:class:`~torch.nn.Unfold` instances created with the same
|
||||
parameters:
|
||||
|
||||
>>> fold_params = dict(kernel_size=..., dilation=..., padding=..., stride=...)
|
||||
>>> fold = nn.Fold(output_size=..., **fold_params)
|
||||
>>> unfold = nn.Unfold(**fold_params)
|
||||
|
||||
Then for any (supported) ``input`` tensor the following
|
||||
equality holds:
|
||||
|
||||
::
|
||||
|
||||
fold(unfold(input)) == divisor * input
|
||||
|
||||
where ``divisor`` is a tensor that depends only on the shape
|
||||
and dtype of the ``input``:
|
||||
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> input_ones = torch.ones(input.shape, dtype=input.dtype)
|
||||
>>> divisor = fold(unfold(input_ones))
|
||||
|
||||
When the ``divisor`` tensor contains no zero elements, then
|
||||
``fold`` and ``unfold`` operations are inverses of each
|
||||
other (up to constant divisor).
|
||||
|
||||
.. warning::
|
||||
Currently, only unbatched (3D) or batched (4D) image-like output tensors are supported.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C \times \prod(\text{kernel\_size}), L)` or :math:`(C \times \prod(\text{kernel\_size}), L)`
|
||||
- Output: :math:`(N, C, \text{output\_size}[0], \text{output\_size}[1], \dots)`
|
||||
or :math:`(C, \text{output\_size}[0], \text{output\_size}[1], \dots)` as described above
|
||||
|
||||
Examples::
|
||||
|
||||
>>> fold = nn.Fold(output_size=(4, 5), kernel_size=(2, 2))
|
||||
>>> input = torch.randn(1, 3 * 2 * 2, 12)
|
||||
>>> output = fold(input)
|
||||
>>> output.size()
|
||||
torch.Size([1, 3, 4, 5])
|
||||
|
||||
.. _link:
|
||||
https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md
|
||||
|
||||
"""
|
||||
)
|
||||
|
||||
__constants__ = ["output_size", "kernel_size", "dilation", "padding", "stride"]
|
||||
output_size: _size_any_t
|
||||
kernel_size: _size_any_t
|
||||
dilation: _size_any_t
|
||||
padding: _size_any_t
|
||||
stride: _size_any_t
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
output_size: _size_any_t,
|
||||
kernel_size: _size_any_t,
|
||||
dilation: _size_any_t = 1,
|
||||
padding: _size_any_t = 0,
|
||||
stride: _size_any_t = 1,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.output_size = output_size
|
||||
self.kernel_size = kernel_size
|
||||
self.dilation = dilation
|
||||
self.padding = padding
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
return F.fold(
|
||||
input,
|
||||
self.output_size,
|
||||
self.kernel_size,
|
||||
self.dilation,
|
||||
self.padding,
|
||||
self.stride,
|
||||
)
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
"""
|
||||
Return the extra representation of the module.
|
||||
"""
|
||||
return (
|
||||
"output_size={output_size}, kernel_size={kernel_size}, "
|
||||
"dilation={dilation}, padding={padding}, stride={stride}".format(
|
||||
**self.__dict__
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class Unfold(Module):
|
||||
(
|
||||
r"""Extracts sliding local blocks from a batched input tensor.
|
||||
|
||||
Consider a batched :attr:`input` tensor of shape :math:`(N, C, *)`,
|
||||
where :math:`N` is the batch dimension, :math:`C` is the channel dimension,
|
||||
and :math:`*` represent arbitrary spatial dimensions. This operation flattens
|
||||
each sliding :attr:`kernel_size`-sized block within the spatial dimensions
|
||||
of :attr:`input` into a column (i.e., last dimension) of a 3-D :attr:`output`
|
||||
tensor of shape :math:`(N, C \times \prod(\text{kernel\_size}), L)`, where
|
||||
:math:`C \times \prod(\text{kernel\_size})` is the total number of values
|
||||
within each block (a block has :math:`\prod(\text{kernel\_size})` spatial
|
||||
locations each containing a :math:`C`-channeled vector), and :math:`L` is
|
||||
the total number of such blocks:
|
||||
|
||||
.. math::
|
||||
L = \prod_d \left\lfloor\frac{\text{spatial\_size}[d] + 2 \times \text{padding}[d] %
|
||||
- \text{dilation}[d] \times (\text{kernel\_size}[d] - 1) - 1}{\text{stride}[d]} + 1\right\rfloor,
|
||||
|
||||
where :math:`\text{spatial\_size}` is formed by the spatial dimensions
|
||||
of :attr:`input` (:math:`*` above), and :math:`d` is over all spatial
|
||||
dimensions.
|
||||
|
||||
Therefore, indexing :attr:`output` at the last dimension (column dimension)
|
||||
gives all values within a certain block.
|
||||
|
||||
The :attr:`padding`, :attr:`stride` and :attr:`dilation` arguments specify
|
||||
how the sliding blocks are retrieved.
|
||||
|
||||
* :attr:`stride` controls the stride for the sliding blocks.
|
||||
|
||||
* :attr:`padding` controls the amount of implicit zero-paddings on both
|
||||
sides for :attr:`padding` number of points for each dimension before
|
||||
reshaping.
|
||||
"""
|
||||
"""
|
||||
* :attr:`dilation` controls the spacing between the kernel points; also known as the \u00e0 trous algorithm.
|
||||
It is harder to describe, but this `link`_ has a nice visualization of what :attr:`dilation` does.
|
||||
"""
|
||||
r"""
|
||||
Args:
|
||||
kernel_size (int or tuple): the size of the sliding blocks
|
||||
dilation (int or tuple, optional): a parameter that controls the
|
||||
stride of elements within the
|
||||
neighborhood. Default: 1
|
||||
padding (int or tuple, optional): implicit zero padding to be added on
|
||||
both sides of input. Default: 0
|
||||
stride (int or tuple, optional): the stride of the sliding blocks in the input
|
||||
spatial dimensions. Default: 1
|
||||
|
||||
* If :attr:`kernel_size`, :attr:`dilation`, :attr:`padding` or
|
||||
:attr:`stride` is an int or a tuple of length 1, their values will be
|
||||
replicated across all spatial dimensions.
|
||||
|
||||
* For the case of two input spatial dimensions this operation is sometimes
|
||||
called ``im2col``.
|
||||
|
||||
.. note::
|
||||
:class:`~torch.nn.Fold` calculates each combined value in the resulting
|
||||
large tensor by summing all values from all containing blocks.
|
||||
:class:`~torch.nn.Unfold` extracts the values in the local blocks by
|
||||
copying from the large tensor. So, if the blocks overlap, they are not
|
||||
inverses of each other.
|
||||
|
||||
In general, folding and unfolding operations are related as
|
||||
follows. Consider :class:`~torch.nn.Fold` and
|
||||
:class:`~torch.nn.Unfold` instances created with the same
|
||||
parameters:
|
||||
|
||||
>>> fold_params = dict(kernel_size=..., dilation=..., padding=..., stride=...)
|
||||
>>> fold = nn.Fold(output_size=..., **fold_params)
|
||||
>>> unfold = nn.Unfold(**fold_params)
|
||||
|
||||
Then for any (supported) ``input`` tensor the following
|
||||
equality holds:
|
||||
|
||||
::
|
||||
|
||||
fold(unfold(input)) == divisor * input
|
||||
|
||||
where ``divisor`` is a tensor that depends only on the shape
|
||||
and dtype of the ``input``:
|
||||
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> input_ones = torch.ones(input.shape, dtype=input.dtype)
|
||||
>>> divisor = fold(unfold(input_ones))
|
||||
|
||||
When the ``divisor`` tensor contains no zero elements, then
|
||||
``fold`` and ``unfold`` operations are inverses of each
|
||||
other (up to constant divisor).
|
||||
|
||||
.. warning::
|
||||
Currently, only 4-D input tensors (batched image-like tensors) are
|
||||
supported.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, *)`
|
||||
- Output: :math:`(N, C \times \prod(\text{kernel\_size}), L)` as described above
|
||||
|
||||
Examples::
|
||||
|
||||
>>> unfold = nn.Unfold(kernel_size=(2, 3))
|
||||
>>> input = torch.randn(2, 5, 3, 4)
|
||||
>>> output = unfold(input)
|
||||
>>> # each patch contains 30 values (2x3=6 vectors, each of 5 channels)
|
||||
>>> # 4 blocks (2x3 kernels) in total in the 3x4 input
|
||||
>>> output.size()
|
||||
torch.Size([2, 30, 4])
|
||||
|
||||
>>> # xdoctest: +IGNORE_WANT
|
||||
>>> # Convolution is equivalent with Unfold + Matrix Multiplication + Fold (or view to output shape)
|
||||
>>> inp = torch.randn(1, 3, 10, 12)
|
||||
>>> w = torch.randn(2, 3, 4, 5)
|
||||
>>> inp_unf = torch.nn.functional.unfold(inp, (4, 5))
|
||||
>>> out_unf = inp_unf.transpose(1, 2).matmul(w.view(w.size(0), -1).t()).transpose(1, 2)
|
||||
>>> out = torch.nn.functional.fold(out_unf, (7, 8), (1, 1))
|
||||
>>> # or equivalently (and avoiding a copy),
|
||||
>>> # out = out_unf.view(1, 2, 7, 8)
|
||||
>>> (torch.nn.functional.conv2d(inp, w) - out).abs().max()
|
||||
tensor(1.9073e-06)
|
||||
|
||||
.. _link:
|
||||
https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md
|
||||
|
||||
"""
|
||||
)
|
||||
|
||||
__constants__ = ["kernel_size", "dilation", "padding", "stride"]
|
||||
kernel_size: _size_any_t
|
||||
dilation: _size_any_t
|
||||
padding: _size_any_t
|
||||
stride: _size_any_t
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
kernel_size: _size_any_t,
|
||||
dilation: _size_any_t = 1,
|
||||
padding: _size_any_t = 0,
|
||||
stride: _size_any_t = 1,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.kernel_size = kernel_size
|
||||
self.dilation = dilation
|
||||
self.padding = padding
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
return F.unfold(
|
||||
input, self.kernel_size, self.dilation, self.padding, self.stride
|
||||
)
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
"""
|
||||
Return the extra representation of the module.
|
||||
"""
|
||||
return (
|
||||
"kernel_size={kernel_size}, dilation={dilation}, padding={padding},"
|
||||
" stride={stride}".format(**self.__dict__)
|
||||
)
|
||||
@@ -0,0 +1,492 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import warnings
|
||||
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
|
||||
from .batchnorm import _LazyNormBase, _NormBase
|
||||
|
||||
|
||||
__all__ = [
|
||||
"InstanceNorm1d",
|
||||
"InstanceNorm2d",
|
||||
"InstanceNorm3d",
|
||||
"LazyInstanceNorm1d",
|
||||
"LazyInstanceNorm2d",
|
||||
"LazyInstanceNorm3d",
|
||||
]
|
||||
|
||||
|
||||
class _InstanceNorm(_NormBase):
|
||||
def __init__(
|
||||
self,
|
||||
num_features: int,
|
||||
eps: float = 1e-5,
|
||||
momentum: float = 0.1,
|
||||
affine: bool = False,
|
||||
track_running_stats: bool = False,
|
||||
device=None,
|
||||
dtype=None,
|
||||
*,
|
||||
bias: bool = True, # for backward compatibility
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__(
|
||||
num_features,
|
||||
eps,
|
||||
momentum,
|
||||
affine,
|
||||
track_running_stats,
|
||||
**factory_kwargs,
|
||||
bias=bias,
|
||||
)
|
||||
|
||||
def _check_input_dim(self, input):
|
||||
raise NotImplementedError
|
||||
|
||||
def _get_no_batch_dim(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def _handle_no_batch_input(self, input):
|
||||
return self._apply_instance_norm(input.unsqueeze(0)).squeeze(0)
|
||||
|
||||
def _apply_instance_norm(self, input):
|
||||
return F.instance_norm(
|
||||
input,
|
||||
self.running_mean,
|
||||
self.running_var,
|
||||
self.weight,
|
||||
self.bias,
|
||||
self.training or not self.track_running_stats,
|
||||
self.momentum if self.momentum is not None else 0.0,
|
||||
self.eps,
|
||||
)
|
||||
|
||||
def _load_from_state_dict(
|
||||
self,
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
strict,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
) -> None:
|
||||
version = local_metadata.get("version", None)
|
||||
# at version 1: removed running_mean and running_var when
|
||||
# track_running_stats=False (default)
|
||||
if version is None and not self.track_running_stats:
|
||||
running_stats_keys = []
|
||||
for name in ("running_mean", "running_var"):
|
||||
key = prefix + name
|
||||
if key in state_dict:
|
||||
running_stats_keys.append(key)
|
||||
if len(running_stats_keys) > 0:
|
||||
error_msgs.append(
|
||||
"Unexpected running stats buffer(s) {names} for {klass} "
|
||||
"with track_running_stats=False. If state_dict is a "
|
||||
"checkpoint saved before 0.4.0, this may be expected "
|
||||
"because {klass} does not track running stats by default "
|
||||
"since 0.4.0. Please remove these keys from state_dict. If "
|
||||
"the running stats are actually needed, instead set "
|
||||
"track_running_stats=True in {klass} to enable them. See "
|
||||
"the documentation of {klass} for details.".format(
|
||||
names=" and ".join(f'"{k}"' for k in running_stats_keys),
|
||||
klass=self.__class__.__name__,
|
||||
)
|
||||
)
|
||||
for key in running_stats_keys:
|
||||
state_dict.pop(key)
|
||||
|
||||
super()._load_from_state_dict(
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
strict,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
)
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
self._check_input_dim(input)
|
||||
|
||||
feature_dim = input.dim() - self._get_no_batch_dim()
|
||||
if input.size(feature_dim) != self.num_features:
|
||||
if self.affine:
|
||||
raise ValueError(
|
||||
f"expected input's size at dim={feature_dim} to match num_features"
|
||||
f" ({self.num_features}), but got: {input.size(feature_dim)}."
|
||||
)
|
||||
else:
|
||||
warnings.warn(
|
||||
f"input's size at dim={feature_dim} does not match num_features. "
|
||||
"You can silence this warning by not passing in num_features, "
|
||||
"which is not used because affine=False",
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
if input.dim() == self._get_no_batch_dim():
|
||||
return self._handle_no_batch_input(input)
|
||||
|
||||
return self._apply_instance_norm(input)
|
||||
|
||||
|
||||
class InstanceNorm1d(_InstanceNorm):
|
||||
r"""Applies Instance Normalization.
|
||||
|
||||
This operation applies Instance Normalization
|
||||
over a 2D (unbatched) or 3D (batched) input as described in the paper
|
||||
`Instance Normalization: The Missing Ingredient for Fast Stylization
|
||||
<https://arxiv.org/abs/1607.08022>`__.
|
||||
|
||||
.. math::
|
||||
|
||||
y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta
|
||||
|
||||
The mean and standard-deviation are calculated per-dimension separately
|
||||
for each object in a mini-batch. :math:`\gamma` and :math:`\beta` are learnable parameter vectors
|
||||
of size `C` (where `C` is the number of features or channels of the input) if :attr:`affine` is ``True``.
|
||||
The variance is calculated via the biased estimator, equivalent to
|
||||
`torch.var(input, correction=0)`.
|
||||
|
||||
By default, this layer uses instance statistics computed from input data in
|
||||
both training and evaluation modes.
|
||||
|
||||
If :attr:`track_running_stats` is set to ``True``, during training this
|
||||
layer keeps running estimates of its computed mean and variance, which are
|
||||
then used for normalization during evaluation. The running estimates are
|
||||
kept with a default :attr:`momentum` of 0.1.
|
||||
|
||||
.. note::
|
||||
This :attr:`momentum` argument is different from one used in optimizer
|
||||
classes and the conventional notion of momentum. Mathematically, the
|
||||
update rule for running statistics here is
|
||||
:math:`\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t`,
|
||||
where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the
|
||||
new observed value.
|
||||
|
||||
.. note::
|
||||
:class:`InstanceNorm1d` and :class:`LayerNorm` are very similar, but
|
||||
have some subtle differences. :class:`InstanceNorm1d` is applied
|
||||
on each channel of channeled data like multidimensional time series, but
|
||||
:class:`LayerNorm` is usually applied on entire sample and often in NLP
|
||||
tasks. Additionally, :class:`LayerNorm` applies elementwise affine
|
||||
transform, while :class:`InstanceNorm1d` usually don't apply affine
|
||||
transform.
|
||||
|
||||
Args:
|
||||
num_features: number of features or channels :math:`C` of the input
|
||||
eps: a value added to the denominator for numerical stability. Default: 1e-5
|
||||
momentum: the value used for the running_mean and running_var computation. Default: 0.1
|
||||
affine: a boolean value that when set to ``True``, this module has
|
||||
learnable affine parameters, initialized the same way as done for batch normalization.
|
||||
Default: ``False``
|
||||
track_running_stats: a boolean value that when set to ``True``, this
|
||||
module tracks the running mean and variance, and when set to ``False``,
|
||||
this module does not track such statistics and always uses batch
|
||||
statistics in both training and eval modes. Default: ``False``
|
||||
bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
|
||||
:attr:`affine` is ``True``). Default: ``True``
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, L)` or :math:`(C, L)`
|
||||
- Output: :math:`(N, C, L)` or :math:`(C, L)` (same shape as input)
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # Without Learnable Parameters
|
||||
>>> m = nn.InstanceNorm1d(100)
|
||||
>>> # With Learnable Parameters
|
||||
>>> m = nn.InstanceNorm1d(100, affine=True)
|
||||
>>> input = torch.randn(20, 100, 40)
|
||||
>>> output = m(input)
|
||||
"""
|
||||
|
||||
def _get_no_batch_dim(self) -> int:
|
||||
return 2
|
||||
|
||||
def _check_input_dim(self, input) -> None:
|
||||
if input.dim() not in (2, 3):
|
||||
raise ValueError(f"expected 2D or 3D input (got {input.dim()}D input)")
|
||||
|
||||
|
||||
class LazyInstanceNorm1d(_LazyNormBase, _InstanceNorm):
|
||||
r"""A :class:`torch.nn.InstanceNorm1d` module with lazy initialization of the ``num_features`` argument.
|
||||
|
||||
The ``num_features`` argument of the :class:`InstanceNorm1d` is inferred from the ``input.size(1)``.
|
||||
The attributes that will be lazily initialized are `weight`, `bias`, `running_mean` and `running_var`.
|
||||
|
||||
Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation
|
||||
on lazy modules and their limitations.
|
||||
|
||||
Args:
|
||||
num_features: :math:`C` from an expected input of size
|
||||
:math:`(N, C, L)` or :math:`(C, L)`
|
||||
eps: a value added to the denominator for numerical stability. Default: 1e-5
|
||||
momentum: the value used for the running_mean and running_var computation. Default: 0.1
|
||||
affine: a boolean value that when set to ``True``, this module has
|
||||
learnable affine parameters, initialized the same way as done for batch normalization.
|
||||
Default: ``False``
|
||||
track_running_stats: a boolean value that when set to ``True``, this
|
||||
module tracks the running mean and variance, and when set to ``False``,
|
||||
this module does not track such statistics and always uses batch
|
||||
statistics in both training and eval modes. Default: ``False``
|
||||
bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
|
||||
:attr:`affine` is ``True``). Default: ``True``
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, L)` or :math:`(C, L)`
|
||||
- Output: :math:`(N, C, L)` or :math:`(C, L)` (same shape as input)
|
||||
"""
|
||||
|
||||
cls_to_become = InstanceNorm1d # type: ignore[assignment]
|
||||
|
||||
def _get_no_batch_dim(self) -> int:
|
||||
return 2
|
||||
|
||||
def _check_input_dim(self, input) -> None:
|
||||
if input.dim() not in (2, 3):
|
||||
raise ValueError(f"expected 2D or 3D input (got {input.dim()}D input)")
|
||||
|
||||
|
||||
class InstanceNorm2d(_InstanceNorm):
|
||||
r"""Applies Instance Normalization.
|
||||
|
||||
This operation applies Instance Normalization
|
||||
over a 4D input (a mini-batch of 2D inputs
|
||||
with additional channel dimension) as described in the paper
|
||||
`Instance Normalization: The Missing Ingredient for Fast Stylization
|
||||
<https://arxiv.org/abs/1607.08022>`__.
|
||||
|
||||
.. math::
|
||||
|
||||
y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta
|
||||
|
||||
The mean and standard-deviation are calculated per-dimension separately
|
||||
for each object in a mini-batch. :math:`\gamma` and :math:`\beta` are learnable parameter vectors
|
||||
of size `C` (where `C` is the input size) if :attr:`affine` is ``True``.
|
||||
The standard-deviation is calculated via the biased estimator, equivalent to
|
||||
`torch.var(input, correction=0)`.
|
||||
|
||||
By default, this layer uses instance statistics computed from input data in
|
||||
both training and evaluation modes.
|
||||
|
||||
If :attr:`track_running_stats` is set to ``True``, during training this
|
||||
layer keeps running estimates of its computed mean and variance, which are
|
||||
then used for normalization during evaluation. The running estimates are
|
||||
kept with a default :attr:`momentum` of 0.1.
|
||||
|
||||
.. note::
|
||||
This :attr:`momentum` argument is different from one used in optimizer
|
||||
classes and the conventional notion of momentum. Mathematically, the
|
||||
update rule for running statistics here is
|
||||
:math:`\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t`,
|
||||
where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the
|
||||
new observed value.
|
||||
|
||||
.. note::
|
||||
:class:`InstanceNorm2d` and :class:`LayerNorm` are very similar, but
|
||||
have some subtle differences. :class:`InstanceNorm2d` is applied
|
||||
on each channel of channeled data like RGB images, but
|
||||
:class:`LayerNorm` is usually applied on entire sample and often in NLP
|
||||
tasks. Additionally, :class:`LayerNorm` applies elementwise affine
|
||||
transform, while :class:`InstanceNorm2d` usually don't apply affine
|
||||
transform.
|
||||
|
||||
Args:
|
||||
num_features: :math:`C` from an expected input of size
|
||||
:math:`(N, C, H, W)` or :math:`(C, H, W)`
|
||||
eps: a value added to the denominator for numerical stability. Default: 1e-5
|
||||
momentum: the value used for the running_mean and running_var computation. Default: 0.1
|
||||
affine: a boolean value that when set to ``True``, this module has
|
||||
learnable affine parameters, initialized the same way as done for batch normalization.
|
||||
Default: ``False``
|
||||
track_running_stats: a boolean value that when set to ``True``, this
|
||||
module tracks the running mean and variance, and when set to ``False``,
|
||||
this module does not track such statistics and always uses batch
|
||||
statistics in both training and eval modes. Default: ``False``
|
||||
bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
|
||||
:attr:`affine` is ``True``). Default: ``True``
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, H, W)` or :math:`(C, H, W)`
|
||||
- Output: :math:`(N, C, H, W)` or :math:`(C, H, W)` (same shape as input)
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # Without Learnable Parameters
|
||||
>>> m = nn.InstanceNorm2d(100)
|
||||
>>> # With Learnable Parameters
|
||||
>>> m = nn.InstanceNorm2d(100, affine=True)
|
||||
>>> input = torch.randn(20, 100, 35, 45)
|
||||
>>> output = m(input)
|
||||
"""
|
||||
|
||||
def _get_no_batch_dim(self) -> int:
|
||||
return 3
|
||||
|
||||
def _check_input_dim(self, input) -> None:
|
||||
if input.dim() not in (3, 4):
|
||||
raise ValueError(f"expected 3D or 4D input (got {input.dim()}D input)")
|
||||
|
||||
|
||||
class LazyInstanceNorm2d(_LazyNormBase, _InstanceNorm):
|
||||
r"""A :class:`torch.nn.InstanceNorm2d` module with lazy initialization of the ``num_features`` argument.
|
||||
|
||||
The ``num_features`` argument of the :class:`InstanceNorm2d` is inferred from the ``input.size(1)``.
|
||||
The attributes that will be lazily initialized are `weight`, `bias`,
|
||||
`running_mean` and `running_var`.
|
||||
|
||||
Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation
|
||||
on lazy modules and their limitations.
|
||||
|
||||
Args:
|
||||
num_features: :math:`C` from an expected input of size
|
||||
:math:`(N, C, H, W)` or :math:`(C, H, W)`
|
||||
eps: a value added to the denominator for numerical stability. Default: 1e-5
|
||||
momentum: the value used for the running_mean and running_var computation. Default: 0.1
|
||||
affine: a boolean value that when set to ``True``, this module has
|
||||
learnable affine parameters, initialized the same way as done for batch normalization.
|
||||
Default: ``False``
|
||||
track_running_stats: a boolean value that when set to ``True``, this
|
||||
module tracks the running mean and variance, and when set to ``False``,
|
||||
this module does not track such statistics and always uses batch
|
||||
statistics in both training and eval modes. Default: ``False``
|
||||
bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
|
||||
:attr:`affine` is ``True``). Default: ``True``
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, H, W)` or :math:`(C, H, W)`
|
||||
- Output: :math:`(N, C, H, W)` or :math:`(C, H, W)` (same shape as input)
|
||||
"""
|
||||
|
||||
cls_to_become = InstanceNorm2d # type: ignore[assignment]
|
||||
|
||||
def _get_no_batch_dim(self) -> int:
|
||||
return 3
|
||||
|
||||
def _check_input_dim(self, input) -> None:
|
||||
if input.dim() not in (3, 4):
|
||||
raise ValueError(f"expected 3D or 4D input (got {input.dim()}D input)")
|
||||
|
||||
|
||||
class InstanceNorm3d(_InstanceNorm):
|
||||
r"""Applies Instance Normalization.
|
||||
|
||||
This operation applies Instance Normalization
|
||||
over a 5D input (a mini-batch of 3D inputs with additional channel dimension) as described in the paper
|
||||
`Instance Normalization: The Missing Ingredient for Fast Stylization
|
||||
<https://arxiv.org/abs/1607.08022>`__.
|
||||
|
||||
.. math::
|
||||
|
||||
y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta
|
||||
|
||||
The mean and standard-deviation are calculated per-dimension separately
|
||||
for each object in a mini-batch. :math:`\gamma` and :math:`\beta` are learnable parameter vectors
|
||||
of size C (where C is the input size) if :attr:`affine` is ``True``.
|
||||
The standard-deviation is calculated via the biased estimator, equivalent to
|
||||
`torch.var(input, correction=0)`.
|
||||
|
||||
By default, this layer uses instance statistics computed from input data in
|
||||
both training and evaluation modes.
|
||||
|
||||
If :attr:`track_running_stats` is set to ``True``, during training this
|
||||
layer keeps running estimates of its computed mean and variance, which are
|
||||
then used for normalization during evaluation. The running estimates are
|
||||
kept with a default :attr:`momentum` of 0.1.
|
||||
|
||||
.. note::
|
||||
This :attr:`momentum` argument is different from one used in optimizer
|
||||
classes and the conventional notion of momentum. Mathematically, the
|
||||
update rule for running statistics here is
|
||||
:math:`\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t`,
|
||||
where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the
|
||||
new observed value.
|
||||
|
||||
.. note::
|
||||
:class:`InstanceNorm3d` and :class:`LayerNorm` are very similar, but
|
||||
have some subtle differences. :class:`InstanceNorm3d` is applied
|
||||
on each channel of channeled data like 3D models with RGB color, but
|
||||
:class:`LayerNorm` is usually applied on entire sample and often in NLP
|
||||
tasks. Additionally, :class:`LayerNorm` applies elementwise affine
|
||||
transform, while :class:`InstanceNorm3d` usually don't apply affine
|
||||
transform.
|
||||
|
||||
Args:
|
||||
num_features: :math:`C` from an expected input of size
|
||||
:math:`(N, C, D, H, W)` or :math:`(C, D, H, W)`
|
||||
eps: a value added to the denominator for numerical stability. Default: 1e-5
|
||||
momentum: the value used for the running_mean and running_var computation. Default: 0.1
|
||||
affine: a boolean value that when set to ``True``, this module has
|
||||
learnable affine parameters, initialized the same way as done for batch normalization.
|
||||
Default: ``False``
|
||||
track_running_stats: a boolean value that when set to ``True``, this
|
||||
module tracks the running mean and variance, and when set to ``False``,
|
||||
this module does not track such statistics and always uses batch
|
||||
statistics in both training and eval modes. Default: ``False``
|
||||
bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
|
||||
:attr:`affine` is ``True``). Default: ``True``
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)`
|
||||
- Output: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)` (same shape as input)
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # Without Learnable Parameters
|
||||
>>> m = nn.InstanceNorm3d(100)
|
||||
>>> # With Learnable Parameters
|
||||
>>> m = nn.InstanceNorm3d(100, affine=True)
|
||||
>>> input = torch.randn(20, 100, 35, 45, 10)
|
||||
>>> output = m(input)
|
||||
"""
|
||||
|
||||
def _get_no_batch_dim(self) -> int:
|
||||
return 4
|
||||
|
||||
def _check_input_dim(self, input) -> None:
|
||||
if input.dim() not in (4, 5):
|
||||
raise ValueError(f"expected 4D or 5D input (got {input.dim()}D input)")
|
||||
|
||||
|
||||
class LazyInstanceNorm3d(_LazyNormBase, _InstanceNorm):
|
||||
r"""A :class:`torch.nn.InstanceNorm3d` module with lazy initialization of the ``num_features`` argument.
|
||||
|
||||
The ``num_features`` argument of the :class:`InstanceNorm3d` is inferred from the ``input.size(1)``.
|
||||
The attributes that will be lazily initialized are `weight`, `bias`,
|
||||
`running_mean` and `running_var`.
|
||||
|
||||
Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation
|
||||
on lazy modules and their limitations.
|
||||
|
||||
Args:
|
||||
num_features: :math:`C` from an expected input of size
|
||||
:math:`(N, C, D, H, W)` or :math:`(C, D, H, W)`
|
||||
eps: a value added to the denominator for numerical stability. Default: 1e-5
|
||||
momentum: the value used for the running_mean and running_var computation. Default: 0.1
|
||||
affine: a boolean value that when set to ``True``, this module has
|
||||
learnable affine parameters, initialized the same way as done for batch normalization.
|
||||
Default: ``False``
|
||||
track_running_stats: a boolean value that when set to ``True``, this
|
||||
module tracks the running mean and variance, and when set to ``False``,
|
||||
this module does not track such statistics and always uses batch
|
||||
statistics in both training and eval modes. Default: ``False``
|
||||
bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
|
||||
:attr:`affine` is ``True``). Default: ``True``
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)`
|
||||
- Output: :math:`(N, C, D, H, W)` or :math:`(C, D, H, W)` (same shape as input)
|
||||
"""
|
||||
|
||||
cls_to_become = InstanceNorm3d # type: ignore[assignment]
|
||||
|
||||
def _get_no_batch_dim(self) -> int:
|
||||
return 4
|
||||
|
||||
def _check_input_dim(self, input) -> None:
|
||||
if input.dim() not in (4, 5):
|
||||
raise ValueError(f"expected 4D or 5D input (got {input.dim()}D input)")
|
||||
@@ -0,0 +1,278 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import itertools
|
||||
from typing import Any, Protocol
|
||||
|
||||
import torch
|
||||
from torch.nn.parameter import is_lazy
|
||||
|
||||
|
||||
__all__ = ["LazyModuleMixin"]
|
||||
|
||||
|
||||
class _LazyProtocol(Protocol):
|
||||
"""This class is used to avoid errors with mypy checks for the attributes in a mixin.
|
||||
|
||||
https://mypy.readthedocs.io/en/latest/more_types.html#mixin-classes
|
||||
"""
|
||||
|
||||
def _register_load_state_dict_pre_hook(self, hook): ...
|
||||
|
||||
def register_forward_pre_hook(self, hook, *, prepend=False, with_kwargs=False): ...
|
||||
|
||||
def _lazy_load_hook(
|
||||
self,
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
strict,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
): ...
|
||||
|
||||
def _get_name(self): ...
|
||||
|
||||
def _infer_parameters(self, module, input): ...
|
||||
|
||||
@property
|
||||
def _parameters(self): ...
|
||||
|
||||
@property
|
||||
def _buffers(self): ...
|
||||
|
||||
@property
|
||||
def _non_persistent_buffers_set(self): ...
|
||||
|
||||
@property
|
||||
def _load_hook(self): ...
|
||||
|
||||
@property
|
||||
def _initialize_hook(self): ...
|
||||
|
||||
|
||||
class LazyModuleMixin:
|
||||
r"""A mixin for modules that lazily initialize parameters, also known as "lazy modules".
|
||||
|
||||
.. warning:
|
||||
Lazy modules are an experimental new feature under active development,
|
||||
and their API is likely to change.
|
||||
|
||||
Modules that lazily initialize parameters, or "lazy modules",
|
||||
derive the shapes of their parameters from the first input(s)
|
||||
to their forward method. Until that first forward they contain
|
||||
:class:`torch.nn.UninitializedParameter` s that should not be accessed
|
||||
or used, and afterward they contain regular :class:`torch.nn.Parameter` s.
|
||||
Lazy modules are convenient since they don't require computing some
|
||||
module arguments, like the :attr:`in_features` argument of a
|
||||
typical :class:`torch.nn.Linear`.
|
||||
|
||||
After construction, networks with lazy modules should first
|
||||
be converted to the desired dtype and placed on the expected device.
|
||||
This is because lazy modules only perform shape inference so the usual dtype
|
||||
and device placement behavior applies.
|
||||
The lazy modules should then perform "dry runs" to initialize all the components in the module.
|
||||
These "dry runs" send inputs of the correct size, dtype, and device through
|
||||
the network and to each one of its lazy modules. After this the network can be used as usual.
|
||||
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> class LazyMLP(torch.nn.Module):
|
||||
... def __init__(self) -> None:
|
||||
... super().__init__()
|
||||
... self.fc1 = torch.nn.LazyLinear(10)
|
||||
... self.relu1 = torch.nn.ReLU()
|
||||
... self.fc2 = torch.nn.LazyLinear(1)
|
||||
... self.relu2 = torch.nn.ReLU()
|
||||
...
|
||||
... def forward(self, input):
|
||||
... x = self.relu1(self.fc1(input))
|
||||
... y = self.relu2(self.fc2(x))
|
||||
... return y
|
||||
>>> # constructs a network with lazy modules
|
||||
>>> lazy_mlp = LazyMLP()
|
||||
>>> # transforms the network's device and dtype
|
||||
>>> # NOTE: these transforms can and should be applied after construction and before any 'dry runs'
|
||||
>>> lazy_mlp = lazy_mlp.cuda()
|
||||
>>> lazy_mlp
|
||||
LazyMLP( (fc1): LazyLinear(in_features=0, out_features=10, bias=True)
|
||||
(relu1): ReLU()
|
||||
(fc2): LazyLinear(in_features=0, out_features=1, bias=True)
|
||||
(relu2): ReLU()
|
||||
)
|
||||
>>> # performs a dry run to initialize the network's lazy modules
|
||||
>>> lazy_mlp(torch.ones(10, 10).cuda())
|
||||
>>> # after initialization, LazyLinear modules become regular Linear modules
|
||||
>>> lazy_mlp
|
||||
LazyMLP(
|
||||
(fc1): Linear(in_features=10, out_features=10, bias=True)
|
||||
(relu1): ReLU()
|
||||
(fc2): Linear(in_features=10, out_features=1, bias=True)
|
||||
(relu2): ReLU()
|
||||
)
|
||||
>>> # attaches an optimizer, since parameters can now be used as usual
|
||||
>>> optim = torch.optim.SGD(lazy_mlp.parameters(), lr=0.01)
|
||||
|
||||
A final caveat when using lazy modules is that the order of initialization of a network's
|
||||
parameters may change, since the lazy modules are always initialized after other modules.
|
||||
For example, if the LazyMLP class defined above had a :class:`torch.nn.LazyLinear` module
|
||||
first and then a regular :class:`torch.nn.Linear` second, the second module would be
|
||||
initialized on construction and the first module would be initialized during the first dry run.
|
||||
This can cause the parameters of a network using lazy modules to be initialized differently
|
||||
than the parameters of a network without lazy modules as the order of parameter initializations,
|
||||
which often depends on a stateful random number generator, is different.
|
||||
Check :doc:`/notes/randomness` for more details.
|
||||
|
||||
Lazy modules can be serialized with a state dict like other modules. For example:
|
||||
|
||||
>>> lazy_mlp = LazyMLP()
|
||||
>>> # The state dict shows the uninitialized parameters
|
||||
>>> lazy_mlp.state_dict()
|
||||
OrderedDict({'fc1.weight': <UninitializedParameter>,
|
||||
'fc1.bias': <UninitializedParameter>,
|
||||
'fc2.weight': <UninitializedParameter>,
|
||||
'fc2.bias': <UninitializedParameter>})
|
||||
|
||||
Lazy modules can load regular :class:`torch.nn.Parameter` s (i.e. you can serialize/deserialize
|
||||
initialized LazyModules and they will remain initialized)
|
||||
|
||||
|
||||
>>> full_mlp = LazyMLP()
|
||||
>>> # Dry run to initialize another module
|
||||
>>> full_mlp.forward(torch.ones(10, 1))
|
||||
>>> # Load an initialized state into a lazy module
|
||||
>>> lazy_mlp.load_state_dict(full_mlp.state_dict())
|
||||
>>> # The state dict now holds valid values
|
||||
>>> lazy_mlp.state_dict()
|
||||
OrderedDict([('fc1.weight',
|
||||
tensor([[-0.3837],
|
||||
[ 0.0907],
|
||||
[ 0.6708],
|
||||
[-0.5223],
|
||||
[-0.9028],
|
||||
[ 0.2851],
|
||||
[-0.4537],
|
||||
[ 0.6813],
|
||||
[ 0.5766],
|
||||
[-0.8678]])),
|
||||
('fc1.bias',
|
||||
tensor([-1.8832e+25, 4.5636e-41, -1.8832e+25, 4.5636e-41, -6.1598e-30,
|
||||
4.5637e-41, -1.8788e+22, 4.5636e-41, -2.0042e-31, 4.5637e-41])),
|
||||
('fc2.weight',
|
||||
tensor([[ 0.1320, 0.2938, 0.0679, 0.2793, 0.1088, -0.1795, -0.2301, 0.2807,
|
||||
0.2479, 0.1091]])),
|
||||
('fc2.bias', tensor([0.0019]))])
|
||||
|
||||
Note, however, that the loaded parameters will not be replaced when doing a "dry run" if they are initialized
|
||||
when the state is loaded. This prevents using initialized modules in different contexts.
|
||||
"""
|
||||
|
||||
# modules inheriting from this will change their __class__ to the specified
|
||||
# one after they are fully initialized
|
||||
cls_to_become: type[Any] | None = None
|
||||
|
||||
def __init__(self: _LazyProtocol, *args, **kwargs):
|
||||
# Mypy doesn't like this super call in a mixin
|
||||
super().__init__(*args, **kwargs) # type: ignore[misc]
|
||||
# pyrefly: ignore [read-only]
|
||||
self._load_hook = self._register_load_state_dict_pre_hook(self._lazy_load_hook)
|
||||
# pyrefly: ignore [read-only]
|
||||
self._initialize_hook = self.register_forward_pre_hook(
|
||||
self._infer_parameters, with_kwargs=True
|
||||
)
|
||||
|
||||
def _save_to_state_dict(self: _LazyProtocol, destination, prefix, keep_vars):
|
||||
# This should be ideally implemented as a hook,
|
||||
# but we should override `detach` in the UninitializedParameter to return itself
|
||||
# which is not clean
|
||||
for name, param in self._parameters.items():
|
||||
if param is not None:
|
||||
if not (is_lazy(param) or keep_vars):
|
||||
param = param.detach()
|
||||
destination[prefix + name] = param
|
||||
for name, buf in self._buffers.items():
|
||||
if buf is not None and name not in self._non_persistent_buffers_set:
|
||||
if not (is_lazy(buf) or keep_vars):
|
||||
buf = buf.detach()
|
||||
destination[prefix + name] = buf
|
||||
|
||||
def _lazy_load_hook(
|
||||
self: _LazyProtocol,
|
||||
state_dict,
|
||||
prefix,
|
||||
local_metadata,
|
||||
strict,
|
||||
missing_keys,
|
||||
unexpected_keys,
|
||||
error_msgs,
|
||||
):
|
||||
"""load_state_dict pre-hook function for lazy buffers and parameters.
|
||||
|
||||
The purpose of this hook is to adjust the current state and/or
|
||||
``state_dict`` being loaded so that a module instance serialized in
|
||||
both un/initialized state can be deserialized onto both un/initialized
|
||||
module instance.
|
||||
See comment in ``torch.nn.Module._register_load_state_dict_pre_hook``
|
||||
for the details of the hook specification.
|
||||
"""
|
||||
for name, param in itertools.chain(
|
||||
self._parameters.items(), self._buffers.items()
|
||||
):
|
||||
key = prefix + name
|
||||
if key in state_dict and param is not None:
|
||||
input_param = state_dict[key]
|
||||
if is_lazy(param):
|
||||
# The current parameter is not initialized but the one being loaded one is
|
||||
# create a new parameter based on the uninitialized one
|
||||
if not is_lazy(input_param):
|
||||
with torch.no_grad():
|
||||
param.materialize(input_param.shape)
|
||||
|
||||
def initialize_parameters(self: _LazyProtocol, *args, **kwargs):
|
||||
r"""Initialize parameters according to the input batch properties.
|
||||
|
||||
This adds an interface to isolate parameter initialization from the
|
||||
forward pass when doing parameter shape inference.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
f"initialize_parameters is not implemented for {self.__class__.__name__}"
|
||||
)
|
||||
|
||||
def has_uninitialized_params(self: _LazyProtocol):
|
||||
r"""Check if a module has parameters that are not initialized."""
|
||||
# This is to avoid the JIT to track this parameter and force
|
||||
# custom modules __setstate__ to add it
|
||||
params = self._parameters.values()
|
||||
buffers = self._buffers.values()
|
||||
for param in itertools.chain(params, buffers):
|
||||
if is_lazy(param):
|
||||
return True
|
||||
return False
|
||||
|
||||
# torchrec tests the code consistency with the following code
|
||||
# fmt: off
|
||||
def _infer_parameters(self: _LazyProtocol, module, args, kwargs=None):
|
||||
r"""Infers the size and initializes the parameters according to the provided input batch.
|
||||
|
||||
Given a module that contains parameters that were declared inferable
|
||||
using :class:`torch.nn.parameter.ParameterMode.Infer`, runs a forward pass
|
||||
in the complete module using the provided input to initialize all the parameters
|
||||
as needed.
|
||||
The module is set into evaluation mode before running the forward pass in order
|
||||
to avoid saving statistics or calculating gradients
|
||||
"""
|
||||
kwargs = kwargs if kwargs else {}
|
||||
module.initialize_parameters(*args, **kwargs)
|
||||
if module.has_uninitialized_params():
|
||||
raise RuntimeError(f'module {self._get_name()} has not been fully initialized')
|
||||
module._initialize_hook.remove()
|
||||
module._load_hook.remove()
|
||||
delattr(module, '_initialize_hook')
|
||||
delattr(module, '_load_hook')
|
||||
if module.cls_to_become is not None:
|
||||
module.__class__ = module.cls_to_become
|
||||
# fmt: on
|
||||
|
||||
def _replicate_for_data_parallel(self: _LazyProtocol):
|
||||
raise RuntimeError(
|
||||
"Modules with uninitialized parameters can't be used with `DataParallel`. "
|
||||
"Run a dummy forward pass to correctly initialize the modules"
|
||||
)
|
||||
@@ -0,0 +1,338 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import math
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.nn import functional as F, init
|
||||
from torch.nn.parameter import Parameter, UninitializedParameter
|
||||
|
||||
from .lazy import LazyModuleMixin
|
||||
from .module import Module
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Bilinear",
|
||||
"Identity",
|
||||
"LazyLinear",
|
||||
"Linear",
|
||||
]
|
||||
|
||||
|
||||
class Identity(Module):
|
||||
r"""A placeholder identity operator that is argument-insensitive.
|
||||
|
||||
Args:
|
||||
args: any argument (unused)
|
||||
kwargs: any keyword argument (unused)
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(*)`, where :math:`*` means any number of dimensions.
|
||||
- Output: :math:`(*)`, same shape as the input.
|
||||
|
||||
Examples::
|
||||
|
||||
>>> m = nn.Identity(54, unused_argument1=0.1, unused_argument2=False)
|
||||
>>> input = torch.randn(128, 20)
|
||||
>>> output = m(input)
|
||||
>>> print(output.size())
|
||||
torch.Size([128, 20])
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
||||
super().__init__()
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
return input
|
||||
|
||||
|
||||
class Linear(Module):
|
||||
r"""Applies an affine linear transformation to the incoming data: :math:`y = xA^T + b`.
|
||||
|
||||
This module supports :ref:`TensorFloat32<tf32_on_ampere>`.
|
||||
|
||||
On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision<fp16_on_mi200>` for backward.
|
||||
|
||||
Args:
|
||||
in_features: size of each input sample
|
||||
out_features: size of each output sample
|
||||
bias: If set to ``False``, the layer will not learn an additive bias.
|
||||
Default: ``True``
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(*, H_\text{in})` where :math:`*` means any number of
|
||||
dimensions including none and :math:`H_\text{in} = \text{in\_features}`.
|
||||
- Output: :math:`(*, H_\text{out})` where all but the last dimension
|
||||
are the same shape as the input and :math:`H_\text{out} = \text{out\_features}`.
|
||||
|
||||
Attributes:
|
||||
weight: the learnable weights of the module of shape
|
||||
:math:`(\text{out\_features}, \text{in\_features})`. The values are
|
||||
initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where
|
||||
:math:`k = \frac{1}{\text{in\_features}}`
|
||||
bias: the learnable bias of the module of shape :math:`(\text{out\_features})`.
|
||||
If :attr:`bias` is ``True``, the values are initialized from
|
||||
:math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
|
||||
:math:`k = \frac{1}{\text{in\_features}}`
|
||||
|
||||
Examples::
|
||||
|
||||
>>> m = nn.Linear(20, 30)
|
||||
>>> input = torch.randn(128, 20)
|
||||
>>> output = m(input)
|
||||
>>> print(output.size())
|
||||
torch.Size([128, 30])
|
||||
"""
|
||||
|
||||
__constants__ = ["in_features", "out_features"]
|
||||
in_features: int
|
||||
out_features: int
|
||||
weight: Tensor
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
out_features: int,
|
||||
bias: bool = True,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
self.in_features = in_features
|
||||
self.out_features = out_features
|
||||
self.weight = Parameter(
|
||||
torch.empty((out_features, in_features), **factory_kwargs)
|
||||
)
|
||||
if bias:
|
||||
self.bias = Parameter(torch.empty(out_features, **factory_kwargs))
|
||||
else:
|
||||
self.register_parameter("bias", None)
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
"""
|
||||
Resets parameters based on their initialization used in ``__init__``.
|
||||
"""
|
||||
# Setting a=sqrt(5) in kaiming_uniform is the same as initializing with
|
||||
# uniform(-1/sqrt(in_features), 1/sqrt(in_features)). For details, see
|
||||
# https://github.com/pytorch/pytorch/issues/57109
|
||||
init.kaiming_uniform_(self.weight, a=math.sqrt(5))
|
||||
if self.bias is not None:
|
||||
fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
|
||||
bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
|
||||
init.uniform_(self.bias, -bound, bound)
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
return F.linear(input, self.weight, self.bias)
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
"""
|
||||
Return the extra representation of the module.
|
||||
"""
|
||||
return f"in_features={self.in_features}, out_features={self.out_features}, bias={self.bias is not None}"
|
||||
|
||||
|
||||
# This class exists solely to avoid triggering an obscure error when scripting
|
||||
# an improperly quantized attention layer. See this issue for details:
|
||||
# https://github.com/pytorch/pytorch/issues/58969
|
||||
# TODO: fail fast on quantization API usage error, then remove this class
|
||||
# and replace uses of it with plain Linear
|
||||
class NonDynamicallyQuantizableLinear(Linear):
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
out_features: int,
|
||||
bias: bool = True,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
in_features, out_features, bias=bias, device=device, dtype=dtype
|
||||
)
|
||||
|
||||
|
||||
class Bilinear(Module):
|
||||
r"""Applies a bilinear transformation to the incoming data: :math:`y = x_1^T A x_2 + b`.
|
||||
|
||||
Args:
|
||||
in1_features: size of each first input sample, must be > 0
|
||||
in2_features: size of each second input sample, must be > 0
|
||||
out_features: size of each output sample, must be > 0
|
||||
bias: If set to ``False``, the layer will not learn an additive bias.
|
||||
Default: ``True``
|
||||
|
||||
Shape:
|
||||
- Input1: :math:`(*, H_\text{in1})` where :math:`H_\text{in1}=\text{in1\_features}` and
|
||||
:math:`*` means any number of additional dimensions including none. All but the last dimension
|
||||
of the inputs should be the same.
|
||||
- Input2: :math:`(*, H_\text{in2})` where :math:`H_\text{in2}=\text{in2\_features}`.
|
||||
- Output: :math:`(*, H_\text{out})` where :math:`H_\text{out}=\text{out\_features}`
|
||||
and all but the last dimension are the same shape as the input.
|
||||
|
||||
Attributes:
|
||||
weight: the learnable weights of the module of shape
|
||||
:math:`(\text{out\_features}, \text{in1\_features}, \text{in2\_features})`.
|
||||
The values are initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where
|
||||
:math:`k = \frac{1}{\text{in1\_features}}`
|
||||
bias: the learnable bias of the module of shape :math:`(\text{out\_features})`.
|
||||
If :attr:`bias` is ``True``, the values are initialized from
|
||||
:math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where
|
||||
:math:`k = \frac{1}{\text{in1\_features}}`
|
||||
|
||||
Examples::
|
||||
|
||||
>>> m = nn.Bilinear(20, 30, 40)
|
||||
>>> input1 = torch.randn(128, 20)
|
||||
>>> input2 = torch.randn(128, 30)
|
||||
>>> output = m(input1, input2)
|
||||
>>> print(output.size())
|
||||
torch.Size([128, 40])
|
||||
"""
|
||||
|
||||
__constants__ = ["in1_features", "in2_features", "out_features"]
|
||||
in1_features: int
|
||||
in2_features: int
|
||||
out_features: int
|
||||
weight: Tensor
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in1_features: int,
|
||||
in2_features: int,
|
||||
out_features: int,
|
||||
bias: bool = True,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
self.in1_features = in1_features
|
||||
self.in2_features = in2_features
|
||||
self.out_features = out_features
|
||||
self.weight = Parameter(
|
||||
torch.empty((out_features, in1_features, in2_features), **factory_kwargs)
|
||||
)
|
||||
|
||||
if bias:
|
||||
self.bias = Parameter(torch.empty(out_features, **factory_kwargs))
|
||||
else:
|
||||
self.register_parameter("bias", None)
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
"""
|
||||
Resets parameters based on their initialization used in ``__init__``.
|
||||
"""
|
||||
if self.in1_features <= 0:
|
||||
raise ValueError(
|
||||
f"in1_features must be > 0, but got (in1_features={self.in1_features})"
|
||||
)
|
||||
bound = 1 / math.sqrt(self.weight.size(1))
|
||||
init.uniform_(self.weight, -bound, bound)
|
||||
if self.bias is not None:
|
||||
init.uniform_(self.bias, -bound, bound)
|
||||
|
||||
def forward(self, input1: Tensor, input2: Tensor) -> Tensor:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
return F.bilinear(input1, input2, self.weight, self.bias)
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
"""
|
||||
Return the extra representation of the module.
|
||||
"""
|
||||
return (
|
||||
f"in1_features={self.in1_features}, in2_features={self.in2_features}, "
|
||||
f"out_features={self.out_features}, bias={self.bias is not None}"
|
||||
)
|
||||
|
||||
|
||||
class LazyLinear(LazyModuleMixin, Linear):
|
||||
r"""A :class:`torch.nn.Linear` module where `in_features` is inferred.
|
||||
|
||||
In this module, the `weight` and `bias` are of :class:`torch.nn.UninitializedParameter`
|
||||
class. They will be initialized after the first call to ``forward`` is done and the
|
||||
module will become a regular :class:`torch.nn.Linear` module. The ``in_features`` argument
|
||||
of the :class:`Linear` is inferred from the ``input.shape[-1]``.
|
||||
|
||||
Check the :class:`torch.nn.modules.lazy.LazyModuleMixin` for further documentation
|
||||
on lazy modules and their limitations.
|
||||
|
||||
Args:
|
||||
out_features: size of each output sample
|
||||
bias: If set to ``False``, the layer will not learn an additive bias.
|
||||
Default: ``True``
|
||||
|
||||
Attributes:
|
||||
weight: the learnable weights of the module of shape
|
||||
:math:`(\text{out\_features}, \text{in\_features})`. The values are
|
||||
initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where
|
||||
:math:`k = \frac{1}{\text{in\_features}}`
|
||||
bias: the learnable bias of the module of shape :math:`(\text{out\_features})`.
|
||||
If :attr:`bias` is ``True``, the values are initialized from
|
||||
:math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
|
||||
:math:`k = \frac{1}{\text{in\_features}}`
|
||||
|
||||
|
||||
"""
|
||||
|
||||
cls_to_become = Linear # type: ignore[assignment]
|
||||
# pyrefly: ignore [bad-override]
|
||||
weight: UninitializedParameter
|
||||
bias: UninitializedParameter # type: ignore[assignment]
|
||||
|
||||
def __init__(
|
||||
self, out_features: int, bias: bool = True, device=None, dtype=None
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
# bias is hardcoded to False to avoid creating tensor
|
||||
# that will soon be overwritten.
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
super().__init__(0, 0, False)
|
||||
# pyrefly: ignore [unexpected-keyword]
|
||||
self.weight = UninitializedParameter(**factory_kwargs)
|
||||
self.out_features = out_features
|
||||
if bias:
|
||||
# pyrefly: ignore [unexpected-keyword]
|
||||
self.bias = UninitializedParameter(**factory_kwargs)
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
"""
|
||||
Resets parameters based on their initialization used in ``__init__``.
|
||||
"""
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
if not self.has_uninitialized_params() and self.in_features != 0:
|
||||
super().reset_parameters()
|
||||
|
||||
def initialize_parameters(self, input) -> None: # type: ignore[override]
|
||||
"""
|
||||
Infers ``in_features`` based on ``input`` and initializes parameters.
|
||||
"""
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
if self.has_uninitialized_params():
|
||||
with torch.no_grad():
|
||||
self.in_features = input.shape[-1]
|
||||
self.weight.materialize((self.out_features, self.in_features))
|
||||
if self.bias is not None:
|
||||
self.bias.materialize((self.out_features,))
|
||||
self.reset_parameters()
|
||||
if self.in_features == 0:
|
||||
if input.shape[-1] != self.weight.shape[-1]:
|
||||
raise AssertionError(
|
||||
f"The in_features inferred from input: {input.shape[-1]} "
|
||||
f"is not equal to in_features from self.weight: "
|
||||
f"{self.weight.shape[-1]}"
|
||||
)
|
||||
self.in_features = input.shape[-1]
|
||||
|
||||
|
||||
# TODO: PartialLinear - maybe in sparse?
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,441 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import numbers
|
||||
|
||||
import torch
|
||||
from torch import Size, Tensor
|
||||
from torch.nn import functional as F, init
|
||||
from torch.nn.parameter import Parameter
|
||||
|
||||
from ._functions import CrossMapLRN2d as _cross_map_lrn2d
|
||||
from .module import Module
|
||||
|
||||
|
||||
__all__ = ["LocalResponseNorm", "CrossMapLRN2d", "LayerNorm", "GroupNorm", "RMSNorm"]
|
||||
|
||||
|
||||
class LocalResponseNorm(Module):
|
||||
r"""Applies local response normalization over an input signal.
|
||||
|
||||
The input signal is composed of several input planes, where channels occupy the second dimension.
|
||||
Applies normalization across channels.
|
||||
|
||||
.. math::
|
||||
b_{c} = a_{c}\left(k + \frac{\alpha}{n}
|
||||
\sum_{c'=\max(0, c-n/2)}^{\min(N-1,c+n/2)}a_{c'}^2\right)^{-\beta}
|
||||
|
||||
Args:
|
||||
size: amount of neighbouring channels used for normalization
|
||||
alpha: multiplicative factor. Default: 0.0001
|
||||
beta: exponent. Default: 0.75
|
||||
k: additive factor. Default: 1
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, *)`
|
||||
- Output: :math:`(N, C, *)` (same shape as input)
|
||||
|
||||
Examples::
|
||||
|
||||
>>> lrn = nn.LocalResponseNorm(2)
|
||||
>>> signal_2d = torch.randn(32, 5, 24, 24)
|
||||
>>> signal_4d = torch.randn(16, 5, 7, 7, 7, 7)
|
||||
>>> output_2d = lrn(signal_2d)
|
||||
>>> output_4d = lrn(signal_4d)
|
||||
|
||||
"""
|
||||
|
||||
__constants__ = ["size", "alpha", "beta", "k"]
|
||||
size: int
|
||||
alpha: float
|
||||
beta: float
|
||||
k: float
|
||||
|
||||
def __init__(
|
||||
self, size: int, alpha: float = 1e-4, beta: float = 0.75, k: float = 1.0
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.size = size
|
||||
self.alpha = alpha
|
||||
self.beta = beta
|
||||
self.k = k
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
return F.local_response_norm(input, self.size, self.alpha, self.beta, self.k)
|
||||
|
||||
def extra_repr(self):
|
||||
"""
|
||||
Return the extra representation of the module.
|
||||
"""
|
||||
return "{size}, alpha={alpha}, beta={beta}, k={k}".format(**self.__dict__)
|
||||
|
||||
|
||||
class CrossMapLRN2d(Module):
|
||||
size: int
|
||||
alpha: float
|
||||
beta: float
|
||||
k: float
|
||||
|
||||
def __init__(
|
||||
self, size: int, alpha: float = 1e-4, beta: float = 0.75, k: float = 1
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.size = size
|
||||
self.alpha = alpha
|
||||
self.beta = beta
|
||||
self.k = k
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
return _cross_map_lrn2d.apply(input, self.size, self.alpha, self.beta, self.k)
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
"""
|
||||
Return the extra representation of the module.
|
||||
"""
|
||||
return "{size}, alpha={alpha}, beta={beta}, k={k}".format(**self.__dict__)
|
||||
|
||||
|
||||
_shape_t = int | list[int] | Size
|
||||
|
||||
|
||||
class LayerNorm(Module):
|
||||
r"""Applies Layer Normalization over a mini-batch of inputs.
|
||||
|
||||
This layer implements the operation as described in
|
||||
the paper `Layer Normalization <https://arxiv.org/abs/1607.06450>`__
|
||||
|
||||
.. math::
|
||||
y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta
|
||||
|
||||
The mean and standard-deviation are calculated over the last `D` dimensions, where `D`
|
||||
is the dimension of :attr:`normalized_shape`. For example, if :attr:`normalized_shape`
|
||||
is ``(3, 5)`` (a 2-dimensional shape), the mean and standard-deviation are computed over
|
||||
the last 2 dimensions of the input (i.e. ``input.mean((-2, -1))``).
|
||||
:math:`\gamma` and :math:`\beta` are learnable affine transform parameters of
|
||||
:attr:`normalized_shape` if :attr:`elementwise_affine` is ``True``.
|
||||
The variance is calculated via the biased estimator, equivalent to
|
||||
`torch.var(input, correction=0)`.
|
||||
|
||||
.. note::
|
||||
Unlike Batch Normalization and Instance Normalization, which applies
|
||||
scalar scale and bias for each entire channel/plane with the
|
||||
:attr:`affine` option, Layer Normalization applies per-element scale and
|
||||
bias with :attr:`elementwise_affine`.
|
||||
|
||||
This layer uses statistics computed from input data in both training and
|
||||
evaluation modes.
|
||||
|
||||
Args:
|
||||
normalized_shape (int or list or torch.Size): input shape from an expected input
|
||||
of size
|
||||
|
||||
.. math::
|
||||
[* \times \text{normalized\_shape}[0] \times \text{normalized\_shape}[1]
|
||||
\times \ldots \times \text{normalized\_shape}[-1]]
|
||||
|
||||
If a single integer is used, it is treated as a singleton list, and this module will
|
||||
normalize over the last dimension which is expected to be of that specific size.
|
||||
eps: a value added to the denominator for numerical stability. Default: 1e-5
|
||||
elementwise_affine: a boolean value that when set to ``True``, this module
|
||||
has learnable per-element affine parameters initialized to ones (for weights)
|
||||
and zeros (for biases). Default: ``True``
|
||||
bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
|
||||
:attr:`elementwise_affine` is ``True``). Default: ``True``
|
||||
|
||||
Attributes:
|
||||
weight: the learnable weights of the module of shape
|
||||
:math:`\text{normalized\_shape}` when :attr:`elementwise_affine` is set to ``True``.
|
||||
The values are initialized to 1.
|
||||
bias: the learnable bias of the module of shape
|
||||
:math:`\text{normalized\_shape}` when :attr:`elementwise_affine` is set to ``True``.
|
||||
The values are initialized to 0.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, *)`
|
||||
- Output: :math:`(N, *)` (same shape as input)
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # NLP Example
|
||||
>>> batch, sentence_length, embedding_dim = 20, 5, 10
|
||||
>>> embedding = torch.randn(batch, sentence_length, embedding_dim)
|
||||
>>> layer_norm = nn.LayerNorm(embedding_dim)
|
||||
>>> # Activate module
|
||||
>>> layer_norm(embedding)
|
||||
>>>
|
||||
>>> # Image Example
|
||||
>>> N, C, H, W = 20, 5, 10, 10
|
||||
>>> input = torch.randn(N, C, H, W)
|
||||
>>> # Normalize over the last three dimensions (i.e. the channel and spatial dimensions)
|
||||
>>> # as shown in the image below
|
||||
>>> layer_norm = nn.LayerNorm([C, H, W])
|
||||
>>> output = layer_norm(input)
|
||||
|
||||
.. image:: ../_static/img/nn/layer_norm.jpg
|
||||
:scale: 50 %
|
||||
|
||||
"""
|
||||
|
||||
__constants__ = ["normalized_shape", "eps", "elementwise_affine"]
|
||||
normalized_shape: tuple[int, ...]
|
||||
eps: float
|
||||
elementwise_affine: bool
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
normalized_shape: _shape_t,
|
||||
eps: float = 1e-5,
|
||||
elementwise_affine: bool = True,
|
||||
bias: bool = True,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
if isinstance(normalized_shape, numbers.Integral):
|
||||
# mypy error: incompatible types in assignment
|
||||
normalized_shape = (normalized_shape,) # type: ignore[assignment]
|
||||
self.normalized_shape = tuple(normalized_shape) # type: ignore[arg-type]
|
||||
self.eps = eps
|
||||
self.elementwise_affine = elementwise_affine
|
||||
if self.elementwise_affine:
|
||||
self.weight = Parameter(
|
||||
torch.empty(self.normalized_shape, **factory_kwargs)
|
||||
)
|
||||
if bias:
|
||||
self.bias = Parameter(
|
||||
torch.empty(self.normalized_shape, **factory_kwargs)
|
||||
)
|
||||
else:
|
||||
self.register_parameter("bias", None)
|
||||
else:
|
||||
self.register_parameter("weight", None)
|
||||
self.register_parameter("bias", None)
|
||||
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
if self.elementwise_affine:
|
||||
init.ones_(self.weight)
|
||||
if self.bias is not None:
|
||||
init.zeros_(self.bias)
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
return F.layer_norm(
|
||||
input, self.normalized_shape, self.weight, self.bias, self.eps
|
||||
)
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
return (
|
||||
"{normalized_shape}, eps={eps}, elementwise_affine={elementwise_affine}, "
|
||||
"bias={use_bias}".format(**self.__dict__, use_bias=self.bias is not None)
|
||||
)
|
||||
|
||||
|
||||
class GroupNorm(Module):
|
||||
r"""Applies Group Normalization over a mini-batch of inputs.
|
||||
|
||||
This layer implements the operation as described in
|
||||
the paper `Group Normalization <https://arxiv.org/abs/1803.08494>`__
|
||||
|
||||
.. math::
|
||||
y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta
|
||||
|
||||
The input channels are separated into :attr:`num_groups` groups, each containing
|
||||
``num_channels / num_groups`` channels. :attr:`num_channels` must be divisible by
|
||||
:attr:`num_groups`. The mean and standard-deviation are calculated
|
||||
separately over each group. :math:`\gamma` and :math:`\beta` are learnable
|
||||
per-channel affine transform parameter vectors of size :attr:`num_channels` if
|
||||
:attr:`affine` is ``True``.
|
||||
The variance is calculated via the biased estimator, equivalent to
|
||||
`torch.var(input, correction=0)`.
|
||||
|
||||
This layer uses statistics computed from input data in both training and
|
||||
evaluation modes.
|
||||
|
||||
Args:
|
||||
num_groups (int): number of groups to separate the channels into
|
||||
num_channels (int): number of channels expected in input
|
||||
eps: a value added to the denominator for numerical stability. Default: 1e-5
|
||||
affine: a boolean value that when set to ``True``, this module
|
||||
has learnable per-channel affine parameters initialized to ones (for weights)
|
||||
and zeros (for biases). Default: ``True``
|
||||
bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
|
||||
:attr:`affine` is ``True``). Default: ``True``
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, *)` where :math:`C=\text{num\_channels}`
|
||||
- Output: :math:`(N, C, *)` (same shape as input)
|
||||
|
||||
Examples::
|
||||
|
||||
>>> input = torch.randn(20, 6, 10, 10)
|
||||
>>> # Separate 6 channels into 3 groups
|
||||
>>> m = nn.GroupNorm(3, 6)
|
||||
>>> # Separate 6 channels into 6 groups (equivalent with InstanceNorm)
|
||||
>>> m = nn.GroupNorm(6, 6)
|
||||
>>> # Put all 6 channels into a single group (equivalent with LayerNorm)
|
||||
>>> m = nn.GroupNorm(1, 6)
|
||||
>>> # Activating the module
|
||||
>>> output = m(input)
|
||||
"""
|
||||
|
||||
__constants__ = ["num_groups", "num_channels", "eps", "affine"]
|
||||
num_groups: int
|
||||
num_channels: int
|
||||
eps: float
|
||||
affine: bool
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_groups: int,
|
||||
num_channels: int,
|
||||
eps: float = 1e-5,
|
||||
affine: bool = True,
|
||||
device=None,
|
||||
dtype=None,
|
||||
*,
|
||||
bias: bool = True,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
if num_channels % num_groups != 0:
|
||||
raise ValueError(
|
||||
f"num_channels ({num_channels}) must be divisible by num_groups ({num_groups})"
|
||||
)
|
||||
|
||||
self.num_groups = num_groups
|
||||
self.num_channels = num_channels
|
||||
self.eps = eps
|
||||
self.affine = affine
|
||||
if self.affine:
|
||||
self.weight = Parameter(torch.empty(num_channels, **factory_kwargs))
|
||||
if bias:
|
||||
self.bias = Parameter(torch.empty(num_channels, **factory_kwargs))
|
||||
else:
|
||||
self.register_parameter("bias", None)
|
||||
else:
|
||||
self.register_parameter("weight", None)
|
||||
self.register_parameter("bias", None)
|
||||
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
if self.affine:
|
||||
init.ones_(self.weight)
|
||||
if self.bias is not None:
|
||||
init.zeros_(self.bias)
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
return F.group_norm(input, self.num_groups, self.weight, self.bias, self.eps)
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
return (
|
||||
"{num_groups}, {num_channels}, eps={eps}, affine={affine}, "
|
||||
"bias={use_bias}".format(**self.__dict__, use_bias=self.bias is not None)
|
||||
)
|
||||
|
||||
|
||||
class RMSNorm(Module):
|
||||
r"""Applies Root Mean Square Layer Normalization over a mini-batch of inputs.
|
||||
|
||||
This layer implements the operation as described in
|
||||
the paper `Root Mean Square Layer Normalization <https://arxiv.org/pdf/1910.07467.pdf>`__
|
||||
|
||||
.. math::
|
||||
y_i = \frac{x_i}{\mathrm{RMS}(x)} * \gamma_i, \quad
|
||||
\text{where} \quad \text{RMS}(x) = \sqrt{\epsilon + \frac{1}{n} \sum_{i=1}^{n} x_i^2}
|
||||
|
||||
The RMS is taken over the last ``D`` dimensions, where ``D``
|
||||
is the dimension of :attr:`normalized_shape`. For example, if :attr:`normalized_shape`
|
||||
is ``(3, 5)`` (a 2-dimensional shape), the RMS is computed over
|
||||
the last 2 dimensions of the input.
|
||||
|
||||
Args:
|
||||
normalized_shape (int or list or torch.Size): input shape from an expected input
|
||||
of size
|
||||
|
||||
.. math::
|
||||
[* \times \text{normalized\_shape}[0] \times \text{normalized\_shape}[1]
|
||||
\times \ldots \times \text{normalized\_shape}[-1]]
|
||||
|
||||
If a single integer is used, it is treated as a singleton list, and this module will
|
||||
normalize over the last dimension which is expected to be of that specific size.
|
||||
eps (float, optional): a value added to the denominator for numerical stability.
|
||||
If not specified, uses the machine epsilon of the computation (opmath) type:
|
||||
fp16/bf16 and fp32 inputs use ``torch.finfo(torch.float32).eps``, while fp64
|
||||
inputs use ``torch.finfo(torch.float64).eps``. Default: ``None``
|
||||
elementwise_affine: a boolean value that when set to ``True``, this module
|
||||
has learnable per-element affine parameters initialized to ones (for weights). Default: ``True``.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, *)`
|
||||
- Output: :math:`(N, *)` (same shape as input)
|
||||
|
||||
Examples::
|
||||
|
||||
>>> rms_norm = nn.RMSNorm([2, 3])
|
||||
>>> input = torch.randn(2, 2, 3)
|
||||
>>> rms_norm(input)
|
||||
|
||||
"""
|
||||
|
||||
__constants__ = ["normalized_shape", "eps", "elementwise_affine"]
|
||||
normalized_shape: tuple[int, ...]
|
||||
eps: float | None
|
||||
elementwise_affine: bool
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
normalized_shape: _shape_t,
|
||||
eps: float | None = None,
|
||||
elementwise_affine: bool = True,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
if isinstance(normalized_shape, numbers.Integral):
|
||||
# mypy error: incompatible types in assignment
|
||||
normalized_shape = (normalized_shape,) # type: ignore[assignment]
|
||||
self.normalized_shape = tuple(normalized_shape) # type: ignore[arg-type]
|
||||
self.eps = eps
|
||||
self.elementwise_affine = elementwise_affine
|
||||
if self.elementwise_affine:
|
||||
self.weight = Parameter(
|
||||
torch.empty(self.normalized_shape, **factory_kwargs)
|
||||
)
|
||||
else:
|
||||
self.register_parameter("weight", None)
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
"""
|
||||
Resets parameters based on their initialization used in __init__.
|
||||
"""
|
||||
if self.elementwise_affine:
|
||||
init.ones_(self.weight)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
return F.rms_norm(x, self.normalized_shape, self.weight, self.eps)
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
"""
|
||||
Return the extra representation of the module.
|
||||
"""
|
||||
return (
|
||||
"{normalized_shape}, eps={eps}, "
|
||||
"elementwise_affine={elementwise_affine}".format(**self.__dict__)
|
||||
)
|
||||
|
||||
|
||||
# TODO: ContrastiveNorm2d
|
||||
# TODO: DivisiveNorm2d
|
||||
# TODO: SubtractiveNorm2d
|
||||
@@ -0,0 +1,842 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from collections.abc import Sequence
|
||||
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
from torch.nn.common_types import _size_2_t, _size_4_t, _size_6_t
|
||||
|
||||
from .module import Module
|
||||
from .utils import _ntuple, _pair, _quadruple
|
||||
|
||||
|
||||
# TODO: grad_output size asserts in THNN
|
||||
|
||||
__all__ = [
|
||||
"CircularPad1d",
|
||||
"CircularPad2d",
|
||||
"CircularPad3d",
|
||||
"ConstantPad1d",
|
||||
"ConstantPad2d",
|
||||
"ConstantPad3d",
|
||||
"ReflectionPad1d",
|
||||
"ReflectionPad2d",
|
||||
"ReflectionPad3d",
|
||||
"ReplicationPad1d",
|
||||
"ReplicationPad2d",
|
||||
"ReplicationPad3d",
|
||||
"ZeroPad1d",
|
||||
"ZeroPad2d",
|
||||
"ZeroPad3d",
|
||||
]
|
||||
|
||||
|
||||
class _CircularPadNd(Module):
|
||||
__constants__ = ["padding"]
|
||||
padding: Sequence[int]
|
||||
|
||||
def _check_input_dim(self, input):
|
||||
raise NotImplementedError
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
self._check_input_dim(input)
|
||||
return F.pad(input, self.padding, "circular")
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
return f"{self.padding}"
|
||||
|
||||
|
||||
class CircularPad1d(_CircularPadNd):
|
||||
r"""Pads the input tensor using circular padding of the input boundary.
|
||||
|
||||
Tensor values at the beginning of the dimension are used to pad the end,
|
||||
and values at the end are used to pad the beginning. If negative padding is
|
||||
applied then the ends of the tensor get removed.
|
||||
|
||||
For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
|
||||
|
||||
Args:
|
||||
padding (int, tuple): the size of the padding. If is `int`, uses the same
|
||||
padding in all boundaries. If a 2-`tuple`, uses
|
||||
(:math:`\text{padding\_left}`, :math:`\text{padding\_right}`)
|
||||
Note that padding size should be less than or equal to the corresponding input dimension.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(C, W_{in})` or :math:`(N, C, W_{in})`.
|
||||
- Output: :math:`(C, W_{out})` or :math:`(N, C, W_{out})`, where
|
||||
|
||||
:math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +IGNORE_WANT("not sure why xdoctest is choking on this")
|
||||
>>> m = nn.CircularPad1d(2)
|
||||
>>> input = torch.arange(8, dtype=torch.float).reshape(1, 2, 4)
|
||||
>>> input
|
||||
tensor([[[0., 1., 2., 3.],
|
||||
[4., 5., 6., 7.]]])
|
||||
>>> m(input)
|
||||
tensor([[[2., 3., 0., 1., 2., 3., 0., 1.],
|
||||
[6., 7., 4., 5., 6., 7., 4., 5.]]])
|
||||
>>> # using different paddings for different sides
|
||||
>>> m = nn.CircularPad1d((3, 1))
|
||||
>>> m(input)
|
||||
tensor([[[1., 2., 3., 0., 1., 2., 3., 0.],
|
||||
[5., 6., 7., 4., 5., 6., 7., 4.]]])
|
||||
"""
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
padding: tuple[int, int]
|
||||
|
||||
def __init__(self, padding: _size_2_t) -> None:
|
||||
super().__init__()
|
||||
self.padding = _pair(padding)
|
||||
|
||||
def _check_input_dim(self, input) -> None:
|
||||
if input.dim() != 2 and input.dim() != 3:
|
||||
raise ValueError(f"expected 2D or 3D input (got {input.dim()}D input)")
|
||||
|
||||
|
||||
class CircularPad2d(_CircularPadNd):
|
||||
r"""Pads the input tensor using circular padding of the input boundary.
|
||||
|
||||
Tensor values at the beginning of the dimension are used to pad the end,
|
||||
and values at the end are used to pad the beginning. If negative padding is
|
||||
applied then the ends of the tensor get removed.
|
||||
|
||||
For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
|
||||
|
||||
Args:
|
||||
padding (int, tuple): the size of the padding. If is `int`, uses the same
|
||||
padding in all boundaries. If a 4-`tuple`, uses (:math:`\text{padding\_left}`,
|
||||
:math:`\text{padding\_right}`, :math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`)
|
||||
Note that padding size should be less than or equal to the corresponding input dimension.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.
|
||||
- Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where
|
||||
|
||||
:math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}`
|
||||
|
||||
:math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
|
||||
|
||||
Examples::
|
||||
|
||||
>>> m = nn.CircularPad2d(2)
|
||||
>>> input = torch.arange(9, dtype=torch.float).reshape(1, 1, 3, 3)
|
||||
>>> input
|
||||
tensor([[[[0., 1., 2.],
|
||||
[3., 4., 5.],
|
||||
[6., 7., 8.]]]])
|
||||
>>> m(input)
|
||||
tensor([[[[4., 5., 3., 4., 5., 3., 4.],
|
||||
[7., 8., 6., 7., 8., 6., 7.],
|
||||
[1., 2., 0., 1., 2., 0., 1.],
|
||||
[4., 5., 3., 4., 5., 3., 4.],
|
||||
[7., 8., 6., 7., 8., 6., 7.],
|
||||
[1., 2., 0., 1., 2., 0., 1.],
|
||||
[4., 5., 3., 4., 5., 3., 4.]]]])
|
||||
>>> # using different paddings for different sides
|
||||
>>> m = nn.CircularPad2d((1, 1, 2, 0))
|
||||
>>> m(input)
|
||||
tensor([[[[5., 3., 4., 5., 3.],
|
||||
[8., 6., 7., 8., 6.],
|
||||
[2., 0., 1., 2., 0.],
|
||||
[5., 3., 4., 5., 3.],
|
||||
[8., 6., 7., 8., 6.]]]])
|
||||
"""
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
padding: tuple[int, int, int, int]
|
||||
|
||||
def __init__(self, padding: _size_4_t) -> None:
|
||||
super().__init__()
|
||||
self.padding = _quadruple(padding)
|
||||
|
||||
def _check_input_dim(self, input) -> None:
|
||||
if input.dim() != 3 and input.dim() != 4:
|
||||
raise ValueError(f"expected 3D or 4D input (got {input.dim()}D input)")
|
||||
|
||||
|
||||
class CircularPad3d(_CircularPadNd):
|
||||
r"""Pads the input tensor using circular padding of the input boundary.
|
||||
|
||||
Tensor values at the beginning of the dimension are used to pad the end,
|
||||
and values at the end are used to pad the beginning. If negative padding is
|
||||
applied then the ends of the tensor get removed.
|
||||
|
||||
For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
|
||||
|
||||
Args:
|
||||
padding (int, tuple): the size of the padding. If is `int`, uses the same
|
||||
padding in all boundaries. If a 6-`tuple`, uses
|
||||
(:math:`\text{padding\_left}`, :math:`\text{padding\_right}`,
|
||||
:math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`,
|
||||
:math:`\text{padding\_front}`, :math:`\text{padding\_back}`)
|
||||
Note that padding size should be less than or equal to the corresponding input dimension.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.
|
||||
- Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or :math:`(C, D_{out}, H_{out}, W_{out})`,
|
||||
where
|
||||
|
||||
:math:`D_{out} = D_{in} + \text{padding\_front} + \text{padding\_back}`
|
||||
|
||||
:math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}`
|
||||
|
||||
:math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +IGNORE_WANT("non-deterministic")
|
||||
>>> m = nn.CircularPad3d(3)
|
||||
>>> input = torch.randn(16, 3, 8, 320, 480)
|
||||
>>> output = m(input)
|
||||
>>> # using different paddings for different sides
|
||||
>>> m = nn.CircularPad3d((3, 3, 6, 6, 1, 1))
|
||||
>>> output = m(input)
|
||||
"""
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
padding: tuple[int, int, int, int, int, int]
|
||||
|
||||
def __init__(self, padding: _size_6_t) -> None:
|
||||
super().__init__()
|
||||
self.padding = _ntuple(6)(padding)
|
||||
|
||||
def _check_input_dim(self, input) -> None:
|
||||
if input.dim() != 4 and input.dim() != 5:
|
||||
raise ValueError(f"expected 4D or 5D input (got {input.dim()}D input)")
|
||||
|
||||
|
||||
class _ConstantPadNd(Module):
|
||||
__constants__ = ["padding", "value"]
|
||||
value: float
|
||||
padding: Sequence[int]
|
||||
|
||||
def __init__(self, value: float) -> None:
|
||||
super().__init__()
|
||||
self.value = value
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
return F.pad(input, self.padding, "constant", self.value)
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
return f"padding={self.padding}, value={self.value}"
|
||||
|
||||
|
||||
class ConstantPad1d(_ConstantPadNd):
|
||||
r"""Pads the input tensor boundaries with a constant value.
|
||||
|
||||
For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
|
||||
|
||||
Args:
|
||||
padding (int, tuple): the size of the padding. If is `int`, uses the same
|
||||
padding in both boundaries. If a 2-`tuple`, uses
|
||||
(:math:`\text{padding\_left}`, :math:`\text{padding\_right}`)
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(C, W_{in})` or :math:`(N, C, W_{in})`.
|
||||
- Output: :math:`(C, W_{out})` or :math:`(N, C, W_{out})`, where
|
||||
|
||||
:math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +IGNORE_WANT("non-deterministic")
|
||||
>>> m = nn.ConstantPad1d(2, 3.5)
|
||||
>>> input = torch.randn(1, 2, 4)
|
||||
>>> input
|
||||
tensor([[[-1.0491, -0.7152, -0.0749, 0.8530],
|
||||
[-1.3287, 1.8966, 0.1466, -0.2771]]])
|
||||
>>> m(input)
|
||||
tensor([[[ 3.5000, 3.5000, -1.0491, -0.7152, -0.0749, 0.8530, 3.5000,
|
||||
3.5000],
|
||||
[ 3.5000, 3.5000, -1.3287, 1.8966, 0.1466, -0.2771, 3.5000,
|
||||
3.5000]]])
|
||||
>>> m = nn.ConstantPad1d(2, 3.5)
|
||||
>>> input = torch.randn(1, 2, 3)
|
||||
>>> input
|
||||
tensor([[[ 1.6616, 1.4523, -1.1255],
|
||||
[-3.6372, 0.1182, -1.8652]]])
|
||||
>>> m(input)
|
||||
tensor([[[ 3.5000, 3.5000, 1.6616, 1.4523, -1.1255, 3.5000, 3.5000],
|
||||
[ 3.5000, 3.5000, -3.6372, 0.1182, -1.8652, 3.5000, 3.5000]]])
|
||||
>>> # using different paddings for different sides
|
||||
>>> m = nn.ConstantPad1d((3, 1), 3.5)
|
||||
>>> m(input)
|
||||
tensor([[[ 3.5000, 3.5000, 3.5000, 1.6616, 1.4523, -1.1255, 3.5000],
|
||||
[ 3.5000, 3.5000, 3.5000, -3.6372, 0.1182, -1.8652, 3.5000]]])
|
||||
"""
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
padding: tuple[int, int]
|
||||
|
||||
def __init__(self, padding: _size_2_t, value: float) -> None:
|
||||
super().__init__(value)
|
||||
self.padding = _pair(padding)
|
||||
|
||||
|
||||
class ConstantPad2d(_ConstantPadNd):
|
||||
r"""Pads the input tensor boundaries with a constant value.
|
||||
|
||||
For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
|
||||
|
||||
Args:
|
||||
padding (int, tuple): the size of the padding. If is `int`, uses the same
|
||||
padding in all boundaries. If a 4-`tuple`, uses (:math:`\text{padding\_left}`,
|
||||
:math:`\text{padding\_right}`, :math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`)
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.
|
||||
- Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where
|
||||
|
||||
:math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}`
|
||||
|
||||
:math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +IGNORE_WANT("non-deterministic")
|
||||
>>> m = nn.ConstantPad2d(2, 3.5)
|
||||
>>> input = torch.randn(1, 2, 2)
|
||||
>>> input
|
||||
tensor([[[ 1.6585, 0.4320],
|
||||
[-0.8701, -0.4649]]])
|
||||
>>> m(input)
|
||||
tensor([[[ 3.5000, 3.5000, 3.5000, 3.5000, 3.5000, 3.5000],
|
||||
[ 3.5000, 3.5000, 3.5000, 3.5000, 3.5000, 3.5000],
|
||||
[ 3.5000, 3.5000, 1.6585, 0.4320, 3.5000, 3.5000],
|
||||
[ 3.5000, 3.5000, -0.8701, -0.4649, 3.5000, 3.5000],
|
||||
[ 3.5000, 3.5000, 3.5000, 3.5000, 3.5000, 3.5000],
|
||||
[ 3.5000, 3.5000, 3.5000, 3.5000, 3.5000, 3.5000]]])
|
||||
>>> # using different paddings for different sides
|
||||
>>> m = nn.ConstantPad2d((3, 0, 2, 1), 3.5)
|
||||
>>> m(input)
|
||||
tensor([[[ 3.5000, 3.5000, 3.5000, 3.5000, 3.5000],
|
||||
[ 3.5000, 3.5000, 3.5000, 3.5000, 3.5000],
|
||||
[ 3.5000, 3.5000, 3.5000, 1.6585, 0.4320],
|
||||
[ 3.5000, 3.5000, 3.5000, -0.8701, -0.4649],
|
||||
[ 3.5000, 3.5000, 3.5000, 3.5000, 3.5000]]])
|
||||
"""
|
||||
|
||||
__constants__ = ["padding", "value"]
|
||||
# pyrefly: ignore [bad-override]
|
||||
padding: tuple[int, int, int, int]
|
||||
|
||||
def __init__(self, padding: _size_4_t, value: float) -> None:
|
||||
super().__init__(value)
|
||||
self.padding = _quadruple(padding)
|
||||
|
||||
|
||||
class ConstantPad3d(_ConstantPadNd):
|
||||
r"""Pads the input tensor boundaries with a constant value.
|
||||
|
||||
For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
|
||||
|
||||
Args:
|
||||
padding (int, tuple): the size of the padding. If is `int`, uses the same
|
||||
padding in all boundaries. If a 6-`tuple`, uses
|
||||
(:math:`\text{padding\_left}`, :math:`\text{padding\_right}`,
|
||||
:math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`,
|
||||
:math:`\text{padding\_front}`, :math:`\text{padding\_back}`)
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.
|
||||
- Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or
|
||||
:math:`(C, D_{out}, H_{out}, W_{out})`, where
|
||||
|
||||
:math:`D_{out} = D_{in} + \text{padding\_front} + \text{padding\_back}`
|
||||
|
||||
:math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}`
|
||||
|
||||
:math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
|
||||
|
||||
Examples::
|
||||
|
||||
>>> m = nn.ConstantPad3d(3, 3.5)
|
||||
>>> input = torch.randn(16, 3, 10, 20, 30)
|
||||
>>> output = m(input)
|
||||
>>> # using different paddings for different sides
|
||||
>>> m = nn.ConstantPad3d((3, 3, 6, 6, 0, 1), 3.5)
|
||||
>>> output = m(input)
|
||||
"""
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
padding: tuple[int, int, int, int, int, int]
|
||||
|
||||
def __init__(self, padding: _size_6_t, value: float) -> None:
|
||||
super().__init__(value)
|
||||
self.padding = _ntuple(6)(padding)
|
||||
|
||||
|
||||
class _ReflectionPadNd(Module):
|
||||
__constants__ = ["padding"]
|
||||
padding: Sequence[int]
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
return F.pad(input, self.padding, "reflect")
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
return f"{self.padding}"
|
||||
|
||||
|
||||
class ReflectionPad1d(_ReflectionPadNd):
|
||||
r"""Pads the input tensor using the reflection of the input boundary.
|
||||
|
||||
For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
|
||||
|
||||
Args:
|
||||
padding (int, tuple): the size of the padding. If is `int`, uses the same
|
||||
padding in all boundaries. If a 2-`tuple`, uses
|
||||
(:math:`\text{padding\_left}`, :math:`\text{padding\_right}`)
|
||||
Note that padding size should be less than the corresponding input dimension.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(C, W_{in})` or :math:`(N, C, W_{in})`.
|
||||
- Output: :math:`(C, W_{out})` or :math:`(N, C, W_{out})`, where
|
||||
|
||||
:math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
|
||||
|
||||
Examples::
|
||||
|
||||
>>> m = nn.ReflectionPad1d(2)
|
||||
>>> # xdoctest: +IGNORE_WANT("other tests seem to modify printing styles")
|
||||
>>> input = torch.arange(8, dtype=torch.float).reshape(1, 2, 4)
|
||||
>>> input
|
||||
tensor([[[0., 1., 2., 3.],
|
||||
[4., 5., 6., 7.]]])
|
||||
>>> m(input)
|
||||
tensor([[[2., 1., 0., 1., 2., 3., 2., 1.],
|
||||
[6., 5., 4., 5., 6., 7., 6., 5.]]])
|
||||
>>> # using different paddings for different sides
|
||||
>>> m = nn.ReflectionPad1d((3, 1))
|
||||
>>> m(input)
|
||||
tensor([[[3., 2., 1., 0., 1., 2., 3., 2.],
|
||||
[7., 6., 5., 4., 5., 6., 7., 6.]]])
|
||||
"""
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
padding: tuple[int, int]
|
||||
|
||||
def __init__(self, padding: _size_2_t) -> None:
|
||||
super().__init__()
|
||||
self.padding = _pair(padding)
|
||||
|
||||
|
||||
class ReflectionPad2d(_ReflectionPadNd):
|
||||
r"""Pads the input tensor using the reflection of the input boundary.
|
||||
|
||||
For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
|
||||
|
||||
Args:
|
||||
padding (int, tuple): the size of the padding. If is `int`, uses the same
|
||||
padding in all boundaries. If a 4-`tuple`, uses (:math:`\text{padding\_left}`,
|
||||
:math:`\text{padding\_right}`, :math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`)
|
||||
Note that padding size should be less than the corresponding input dimension.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.
|
||||
- Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})` where
|
||||
|
||||
:math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}`
|
||||
|
||||
:math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +IGNORE_WANT("not sure why xdoctest is choking on this")
|
||||
>>> m = nn.ReflectionPad2d(2)
|
||||
>>> input = torch.arange(9, dtype=torch.float).reshape(1, 1, 3, 3)
|
||||
>>> input
|
||||
tensor([[[[0., 1., 2.],
|
||||
[3., 4., 5.],
|
||||
[6., 7., 8.]]]])
|
||||
>>> m(input)
|
||||
tensor([[[[8., 7., 6., 7., 8., 7., 6.],
|
||||
[5., 4., 3., 4., 5., 4., 3.],
|
||||
[2., 1., 0., 1., 2., 1., 0.],
|
||||
[5., 4., 3., 4., 5., 4., 3.],
|
||||
[8., 7., 6., 7., 8., 7., 6.],
|
||||
[5., 4., 3., 4., 5., 4., 3.],
|
||||
[2., 1., 0., 1., 2., 1., 0.]]]])
|
||||
>>> # using different paddings for different sides
|
||||
>>> m = nn.ReflectionPad2d((1, 1, 2, 0))
|
||||
>>> m(input)
|
||||
tensor([[[[7., 6., 7., 8., 7.],
|
||||
[4., 3., 4., 5., 4.],
|
||||
[1., 0., 1., 2., 1.],
|
||||
[4., 3., 4., 5., 4.],
|
||||
[7., 6., 7., 8., 7.]]]])
|
||||
"""
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
padding: tuple[int, int, int, int]
|
||||
|
||||
def __init__(self, padding: _size_4_t) -> None:
|
||||
super().__init__()
|
||||
self.padding = _quadruple(padding)
|
||||
|
||||
|
||||
class ReflectionPad3d(_ReflectionPadNd):
|
||||
r"""Pads the input tensor using the reflection of the input boundary.
|
||||
|
||||
For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
|
||||
|
||||
Args:
|
||||
padding (int, tuple): the size of the padding. If is `int`, uses the same
|
||||
padding in all boundaries. If a 6-`tuple`, uses
|
||||
(:math:`\text{padding\_left}`, :math:`\text{padding\_right}`,
|
||||
:math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`,
|
||||
:math:`\text{padding\_front}`, :math:`\text{padding\_back}`)
|
||||
Note that padding size should be less than the corresponding input dimension.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.
|
||||
- Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or :math:`(C, D_{out}, H_{out}, W_{out})`,
|
||||
where
|
||||
|
||||
:math:`D_{out} = D_{in} + \text{padding\_front} + \text{padding\_back}`
|
||||
|
||||
:math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}`
|
||||
|
||||
:math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +IGNORE_WANT("not sure why xdoctest is choking on this")
|
||||
>>> m = nn.ReflectionPad3d(1)
|
||||
>>> input = torch.arange(8, dtype=torch.float).reshape(1, 1, 2, 2, 2)
|
||||
>>> m(input)
|
||||
tensor([[[[[7., 6., 7., 6.],
|
||||
[5., 4., 5., 4.],
|
||||
[7., 6., 7., 6.],
|
||||
[5., 4., 5., 4.]],
|
||||
[[3., 2., 3., 2.],
|
||||
[1., 0., 1., 0.],
|
||||
[3., 2., 3., 2.],
|
||||
[1., 0., 1., 0.]],
|
||||
[[7., 6., 7., 6.],
|
||||
[5., 4., 5., 4.],
|
||||
[7., 6., 7., 6.],
|
||||
[5., 4., 5., 4.]],
|
||||
[[3., 2., 3., 2.],
|
||||
[1., 0., 1., 0.],
|
||||
[3., 2., 3., 2.],
|
||||
[1., 0., 1., 0.]]]]])
|
||||
"""
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
padding: tuple[int, int, int, int, int, int]
|
||||
|
||||
def __init__(self, padding: _size_6_t) -> None:
|
||||
super().__init__()
|
||||
self.padding = _ntuple(6)(padding)
|
||||
|
||||
|
||||
class _ReplicationPadNd(Module):
|
||||
__constants__ = ["padding"]
|
||||
padding: Sequence[int]
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
return F.pad(input, self.padding, "replicate")
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
return f"{self.padding}"
|
||||
|
||||
|
||||
class ReplicationPad1d(_ReplicationPadNd):
|
||||
r"""Pads the input tensor using replication of the input boundary.
|
||||
|
||||
For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
|
||||
|
||||
Args:
|
||||
padding (int, tuple): the size of the padding. If is `int`, uses the same
|
||||
padding in all boundaries. If a 2-`tuple`, uses
|
||||
(:math:`\text{padding\_left}`, :math:`\text{padding\_right}`)
|
||||
Note that the output dimensions must remain positive.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(C, W_{in})` or :math:`(N, C, W_{in})`.
|
||||
- Output: :math:`(C, W_{out})` or :math:`(N, C, W_{out})`, where
|
||||
|
||||
:math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +IGNORE_WANT("not sure why xdoctest is choking on this")
|
||||
>>> m = nn.ReplicationPad1d(2)
|
||||
>>> input = torch.arange(8, dtype=torch.float).reshape(1, 2, 4)
|
||||
>>> input
|
||||
tensor([[[0., 1., 2., 3.],
|
||||
[4., 5., 6., 7.]]])
|
||||
>>> m(input)
|
||||
tensor([[[0., 0., 0., 1., 2., 3., 3., 3.],
|
||||
[4., 4., 4., 5., 6., 7., 7., 7.]]])
|
||||
>>> # using different paddings for different sides
|
||||
>>> m = nn.ReplicationPad1d((3, 1))
|
||||
>>> m(input)
|
||||
tensor([[[0., 0., 0., 0., 1., 2., 3., 3.],
|
||||
[4., 4., 4., 4., 5., 6., 7., 7.]]])
|
||||
"""
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
padding: tuple[int, int]
|
||||
|
||||
def __init__(self, padding: _size_2_t) -> None:
|
||||
super().__init__()
|
||||
self.padding = _pair(padding)
|
||||
|
||||
|
||||
class ReplicationPad2d(_ReplicationPadNd):
|
||||
r"""Pads the input tensor using replication of the input boundary.
|
||||
|
||||
For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
|
||||
|
||||
Args:
|
||||
padding (int, tuple): the size of the padding. If is `int`, uses the same
|
||||
padding in all boundaries. If a 4-`tuple`, uses (:math:`\text{padding\_left}`,
|
||||
:math:`\text{padding\_right}`, :math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`)
|
||||
Note that the output dimensions must remain positive.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.
|
||||
- Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where
|
||||
|
||||
:math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}`
|
||||
|
||||
:math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
|
||||
|
||||
Examples::
|
||||
|
||||
>>> m = nn.ReplicationPad2d(2)
|
||||
>>> # xdoctest: +IGNORE_WANT("non-deterministic")
|
||||
>>> input = torch.arange(9, dtype=torch.float).reshape(1, 1, 3, 3)
|
||||
>>> input
|
||||
tensor([[[[0., 1., 2.],
|
||||
[3., 4., 5.],
|
||||
[6., 7., 8.]]]])
|
||||
>>> m(input)
|
||||
tensor([[[[0., 0., 0., 1., 2., 2., 2.],
|
||||
[0., 0., 0., 1., 2., 2., 2.],
|
||||
[0., 0., 0., 1., 2., 2., 2.],
|
||||
[3., 3., 3., 4., 5., 5., 5.],
|
||||
[6., 6., 6., 7., 8., 8., 8.],
|
||||
[6., 6., 6., 7., 8., 8., 8.],
|
||||
[6., 6., 6., 7., 8., 8., 8.]]]])
|
||||
>>> # using different paddings for different sides
|
||||
>>> m = nn.ReplicationPad2d((1, 1, 2, 0))
|
||||
>>> m(input)
|
||||
tensor([[[[0., 0., 1., 2., 2.],
|
||||
[0., 0., 1., 2., 2.],
|
||||
[0., 0., 1., 2., 2.],
|
||||
[3., 3., 4., 5., 5.],
|
||||
[6., 6., 7., 8., 8.]]]])
|
||||
"""
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
padding: tuple[int, int, int, int]
|
||||
|
||||
def __init__(self, padding: _size_4_t) -> None:
|
||||
super().__init__()
|
||||
self.padding = _quadruple(padding)
|
||||
|
||||
|
||||
class ReplicationPad3d(_ReplicationPadNd):
|
||||
r"""Pads the input tensor using replication of the input boundary.
|
||||
|
||||
For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
|
||||
|
||||
Args:
|
||||
padding (int, tuple): the size of the padding. If is `int`, uses the same
|
||||
padding in all boundaries. If a 6-`tuple`, uses
|
||||
(:math:`\text{padding\_left}`, :math:`\text{padding\_right}`,
|
||||
:math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`,
|
||||
:math:`\text{padding\_front}`, :math:`\text{padding\_back}`)
|
||||
Note that the output dimensions must remain positive.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.
|
||||
- Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or :math:`(C, D_{out}, H_{out}, W_{out})`,
|
||||
where
|
||||
|
||||
:math:`D_{out} = D_{in} + \text{padding\_front} + \text{padding\_back}`
|
||||
|
||||
:math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}`
|
||||
|
||||
:math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +IGNORE_WANT("non-deterministic")
|
||||
>>> m = nn.ReplicationPad3d(3)
|
||||
>>> input = torch.randn(16, 3, 8, 320, 480)
|
||||
>>> output = m(input)
|
||||
>>> # using different paddings for different sides
|
||||
>>> m = nn.ReplicationPad3d((3, 3, 6, 6, 1, 1))
|
||||
>>> output = m(input)
|
||||
"""
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
padding: tuple[int, int, int, int, int, int]
|
||||
|
||||
def __init__(self, padding: _size_6_t) -> None:
|
||||
super().__init__()
|
||||
self.padding = _ntuple(6)(padding)
|
||||
|
||||
|
||||
class ZeroPad1d(ConstantPad1d):
|
||||
r"""Pads the input tensor boundaries with zero.
|
||||
|
||||
For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
|
||||
|
||||
Args:
|
||||
padding (int, tuple): the size of the padding. If is `int`, uses the same
|
||||
padding in both boundaries. If a 2-`tuple`, uses
|
||||
(:math:`\text{padding\_left}`, :math:`\text{padding\_right}`)
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(C, W_{in})` or :math:`(N, C, W_{in})`.
|
||||
- Output: :math:`(C, W_{out})` or :math:`(N, C, W_{out})`, where
|
||||
|
||||
:math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +IGNORE_WANT("non-deterministic")
|
||||
>>> m = nn.ZeroPad1d(2)
|
||||
>>> input = torch.randn(1, 2, 4)
|
||||
>>> input
|
||||
tensor([[[-1.0491, -0.7152, -0.0749, 0.8530],
|
||||
[-1.3287, 1.8966, 0.1466, -0.2771]]])
|
||||
>>> m(input)
|
||||
tensor([[[ 0.0000, 0.0000, -1.0491, -0.7152, -0.0749, 0.8530, 0.0000,
|
||||
0.0000],
|
||||
[ 0.0000, 0.0000, -1.3287, 1.8966, 0.1466, -0.2771, 0.0000,
|
||||
0.0000]]])
|
||||
>>> m = nn.ZeroPad1d(2)
|
||||
>>> input = torch.randn(1, 2, 3)
|
||||
>>> input
|
||||
tensor([[[ 1.6616, 1.4523, -1.1255],
|
||||
[-3.6372, 0.1182, -1.8652]]])
|
||||
>>> m(input)
|
||||
tensor([[[ 0.0000, 0.0000, 1.6616, 1.4523, -1.1255, 0.0000, 0.0000],
|
||||
[ 0.0000, 0.0000, -3.6372, 0.1182, -1.8652, 0.0000, 0.0000]]])
|
||||
>>> # using different paddings for different sides
|
||||
>>> m = nn.ZeroPad1d((3, 1))
|
||||
>>> m(input)
|
||||
tensor([[[ 0.0000, 0.0000, 0.0000, 1.6616, 1.4523, -1.1255, 0.0000],
|
||||
[ 0.0000, 0.0000, 0.0000, -3.6372, 0.1182, -1.8652, 0.0000]]])
|
||||
"""
|
||||
|
||||
padding: tuple[int, int]
|
||||
|
||||
def __init__(self, padding: _size_2_t) -> None:
|
||||
super().__init__(padding, 0.0)
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
"""
|
||||
Return the extra representation of the module.
|
||||
"""
|
||||
return f"{self.padding}"
|
||||
|
||||
|
||||
class ZeroPad2d(ConstantPad2d):
|
||||
r"""Pads the input tensor boundaries with zero.
|
||||
|
||||
For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
|
||||
|
||||
Args:
|
||||
padding (int, tuple): the size of the padding. If is `int`, uses the same
|
||||
padding in all boundaries. If a 4-`tuple`, uses (:math:`\text{padding\_left}`,
|
||||
:math:`\text{padding\_right}`, :math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`)
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.
|
||||
- Output: :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`, where
|
||||
|
||||
:math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}`
|
||||
|
||||
:math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # xdoctest: +IGNORE_WANT("non-deterministic")
|
||||
>>> m = nn.ZeroPad2d(2)
|
||||
>>> input = torch.randn(1, 1, 3, 3)
|
||||
>>> input
|
||||
tensor([[[[-0.1678, -0.4418, 1.9466],
|
||||
[ 0.9604, -0.4219, -0.5241],
|
||||
[-0.9162, -0.5436, -0.6446]]]])
|
||||
>>> m(input)
|
||||
tensor([[[[ 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
|
||||
[ 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
|
||||
[ 0.0000, 0.0000, -0.1678, -0.4418, 1.9466, 0.0000, 0.0000],
|
||||
[ 0.0000, 0.0000, 0.9604, -0.4219, -0.5241, 0.0000, 0.0000],
|
||||
[ 0.0000, 0.0000, -0.9162, -0.5436, -0.6446, 0.0000, 0.0000],
|
||||
[ 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
|
||||
[ 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000]]]])
|
||||
>>> # using different paddings for different sides
|
||||
>>> m = nn.ZeroPad2d((1, 1, 2, 0))
|
||||
>>> m(input)
|
||||
tensor([[[[ 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
|
||||
[ 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
|
||||
[ 0.0000, -0.1678, -0.4418, 1.9466, 0.0000],
|
||||
[ 0.0000, 0.9604, -0.4219, -0.5241, 0.0000],
|
||||
[ 0.0000, -0.9162, -0.5436, -0.6446, 0.0000]]]])
|
||||
"""
|
||||
|
||||
padding: tuple[int, int, int, int]
|
||||
|
||||
def __init__(self, padding: _size_4_t) -> None:
|
||||
super().__init__(padding, 0.0)
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
"""
|
||||
Return the extra representation of the module.
|
||||
"""
|
||||
return f"{self.padding}"
|
||||
|
||||
|
||||
class ZeroPad3d(ConstantPad3d):
|
||||
r"""Pads the input tensor boundaries with zero.
|
||||
|
||||
For `N`-dimensional padding, use :func:`torch.nn.functional.pad()`.
|
||||
|
||||
Args:
|
||||
padding (int, tuple): the size of the padding. If is `int`, uses the same
|
||||
padding in all boundaries. If a 6-`tuple`, uses
|
||||
(:math:`\text{padding\_left}`, :math:`\text{padding\_right}`,
|
||||
:math:`\text{padding\_top}`, :math:`\text{padding\_bottom}`,
|
||||
:math:`\text{padding\_front}`, :math:`\text{padding\_back}`)
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, D_{in}, H_{in}, W_{in})` or :math:`(C, D_{in}, H_{in}, W_{in})`.
|
||||
- Output: :math:`(N, C, D_{out}, H_{out}, W_{out})` or
|
||||
:math:`(C, D_{out}, H_{out}, W_{out})`, where
|
||||
|
||||
:math:`D_{out} = D_{in} + \text{padding\_front} + \text{padding\_back}`
|
||||
|
||||
:math:`H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}`
|
||||
|
||||
:math:`W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}`
|
||||
|
||||
Examples::
|
||||
|
||||
>>> m = nn.ZeroPad3d(3)
|
||||
>>> input = torch.randn(16, 3, 10, 20, 30)
|
||||
>>> output = m(input)
|
||||
>>> # using different paddings for different sides
|
||||
>>> m = nn.ZeroPad3d((3, 3, 6, 6, 0, 1))
|
||||
>>> output = m(input)
|
||||
"""
|
||||
|
||||
padding: tuple[int, int, int, int, int, int]
|
||||
|
||||
def __init__(self, padding: _size_6_t) -> None:
|
||||
super().__init__(padding, 0.0)
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
"""
|
||||
Return the extra representation of the module.
|
||||
"""
|
||||
return f"{self.padding}"
|
||||
@@ -0,0 +1,127 @@
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
|
||||
from .module import Module
|
||||
|
||||
|
||||
__all__ = ["PixelShuffle", "PixelUnshuffle"]
|
||||
|
||||
|
||||
class PixelShuffle(Module):
|
||||
r"""Rearrange elements in a tensor according to an upscaling factor.
|
||||
|
||||
Rearranges elements in a tensor of shape :math:`(*, C \times r^2, H, W)`
|
||||
to a tensor of shape :math:`(*, C, H \times r, W \times r)`, where r is an upscale factor.
|
||||
|
||||
This is useful for implementing efficient sub-pixel convolution
|
||||
with a stride of :math:`1/r`.
|
||||
|
||||
See the paper:
|
||||
`Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network`_
|
||||
by Shi et al. (2016) for more details.
|
||||
|
||||
Args:
|
||||
upscale_factor (int): factor to increase spatial resolution by
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(*, C_{in}, H_{in}, W_{in})`, where * is zero or more batch dimensions
|
||||
- Output: :math:`(*, C_{out}, H_{out}, W_{out})`, where
|
||||
|
||||
.. math::
|
||||
C_{out} = C_{in} \div \text{upscale\_factor}^2
|
||||
|
||||
.. math::
|
||||
H_{out} = H_{in} \times \text{upscale\_factor}
|
||||
|
||||
.. math::
|
||||
W_{out} = W_{in} \times \text{upscale\_factor}
|
||||
|
||||
Examples::
|
||||
|
||||
>>> pixel_shuffle = nn.PixelShuffle(3)
|
||||
>>> input = torch.randn(1, 9, 4, 4)
|
||||
>>> output = pixel_shuffle(input)
|
||||
>>> print(output.size())
|
||||
torch.Size([1, 1, 12, 12])
|
||||
|
||||
.. _Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network:
|
||||
https://arxiv.org/abs/1609.05158
|
||||
"""
|
||||
|
||||
__constants__ = ["upscale_factor"]
|
||||
upscale_factor: int
|
||||
|
||||
def __init__(self, upscale_factor: int) -> None:
|
||||
super().__init__()
|
||||
self.upscale_factor = upscale_factor
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
return F.pixel_shuffle(input, self.upscale_factor)
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
"""
|
||||
Return the extra representation of the module.
|
||||
"""
|
||||
return f"upscale_factor={self.upscale_factor}"
|
||||
|
||||
|
||||
class PixelUnshuffle(Module):
|
||||
r"""Reverse the PixelShuffle operation.
|
||||
|
||||
Reverses the :class:`~torch.nn.PixelShuffle` operation by rearranging elements
|
||||
in a tensor of shape :math:`(*, C, H \times r, W \times r)` to a tensor of shape
|
||||
:math:`(*, C \times r^2, H, W)`, where r is a downscale factor.
|
||||
|
||||
See the paper:
|
||||
`Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network`_
|
||||
by Shi et al. (2016) for more details.
|
||||
|
||||
Args:
|
||||
downscale_factor (int): factor to decrease spatial resolution by
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(*, C_{in}, H_{in}, W_{in})`, where * is zero or more batch dimensions
|
||||
- Output: :math:`(*, C_{out}, H_{out}, W_{out})`, where
|
||||
|
||||
.. math::
|
||||
C_{out} = C_{in} \times \text{downscale\_factor}^2
|
||||
|
||||
.. math::
|
||||
H_{out} = H_{in} \div \text{downscale\_factor}
|
||||
|
||||
.. math::
|
||||
W_{out} = W_{in} \div \text{downscale\_factor}
|
||||
|
||||
Examples::
|
||||
|
||||
>>> pixel_unshuffle = nn.PixelUnshuffle(3)
|
||||
>>> input = torch.randn(1, 1, 12, 12)
|
||||
>>> output = pixel_unshuffle(input)
|
||||
>>> print(output.size())
|
||||
torch.Size([1, 9, 4, 4])
|
||||
|
||||
.. _Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network:
|
||||
https://arxiv.org/abs/1609.05158
|
||||
"""
|
||||
|
||||
__constants__ = ["downscale_factor"]
|
||||
downscale_factor: int
|
||||
|
||||
def __init__(self, downscale_factor: int) -> None:
|
||||
super().__init__()
|
||||
self.downscale_factor = downscale_factor
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
return F.pixel_unshuffle(input, self.downscale_factor)
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
"""
|
||||
Return the extra representation of the module.
|
||||
"""
|
||||
return f"downscale_factor={self.downscale_factor}"
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,543 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.nn import functional as F, init
|
||||
from torch.nn.parameter import Parameter
|
||||
|
||||
from .module import Module
|
||||
|
||||
|
||||
__all__ = ["Embedding", "EmbeddingBag"]
|
||||
|
||||
|
||||
class Embedding(Module):
|
||||
r"""A simple lookup table that stores embeddings of a fixed dictionary and size.
|
||||
|
||||
This module is often used to store word embeddings and retrieve them using indices.
|
||||
The input to the module is a list of indices, and the output is the corresponding
|
||||
word embeddings.
|
||||
|
||||
Args:
|
||||
num_embeddings (int): size of the dictionary of embeddings
|
||||
embedding_dim (int): the size of each embedding vector
|
||||
padding_idx (int, optional): If specified, the entries at :attr:`padding_idx` do not contribute to the gradient;
|
||||
therefore, the embedding vector at :attr:`padding_idx` is not updated during training,
|
||||
i.e. it remains as a fixed "pad". For a newly constructed Embedding,
|
||||
the embedding vector at :attr:`padding_idx` will default to all zeros,
|
||||
but can be updated to another value to be used as the padding vector.
|
||||
max_norm (float, optional): If given, each embedding vector with norm larger than :attr:`max_norm`
|
||||
is renormalized to have norm :attr:`max_norm`.
|
||||
norm_type (float, optional): The p of the p-norm to compute for the :attr:`max_norm` option. Default ``2``.
|
||||
scale_grad_by_freq (bool, optional): If given, this will scale gradients by the inverse of frequency of
|
||||
the words in the mini-batch. Default ``False``.
|
||||
sparse (bool, optional): If ``True``, gradient w.r.t. :attr:`weight` matrix will be a sparse tensor.
|
||||
See Notes for more details regarding sparse gradients.
|
||||
|
||||
Attributes:
|
||||
weight (Tensor): the learnable weights of the module of shape (num_embeddings, embedding_dim)
|
||||
initialized from :math:`\mathcal{N}(0, 1)`
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(*)`, IntTensor or LongTensor of arbitrary shape containing the indices to extract
|
||||
- Output: :math:`(*, H)`, where `*` is the input shape and :math:`H=\text{embedding\_dim}`
|
||||
|
||||
.. note::
|
||||
Keep in mind that only a limited number of optimizers support
|
||||
sparse gradients: currently it's :class:`optim.SGD` (`CUDA` and `CPU`),
|
||||
:class:`optim.SparseAdam` (`CUDA` and `CPU`) and :class:`optim.Adagrad` (`CPU`)
|
||||
|
||||
.. note::
|
||||
When :attr:`max_norm` is not ``None``, :class:`Embedding`'s forward method will modify the
|
||||
:attr:`weight` tensor in-place. Since tensors needed for gradient computations cannot be
|
||||
modified in-place, performing a differentiable operation on ``Embedding.weight`` before
|
||||
calling :class:`Embedding`'s forward method requires cloning ``Embedding.weight`` when
|
||||
:attr:`max_norm` is not ``None``. For example::
|
||||
|
||||
n, d, m = 3, 5, 7
|
||||
embedding = nn.Embedding(n, d, max_norm=1.0)
|
||||
W = torch.randn((m, d), requires_grad=True)
|
||||
idx = torch.tensor([1, 2])
|
||||
a = (
|
||||
embedding.weight.clone() @ W.t()
|
||||
) # weight must be cloned for this to be differentiable
|
||||
b = embedding(idx) @ W.t() # modifies weight in-place
|
||||
out = a.unsqueeze(0) + b.unsqueeze(1)
|
||||
loss = out.sigmoid().prod()
|
||||
loss.backward()
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # an Embedding module containing 10 tensors of size 3
|
||||
>>> embedding = nn.Embedding(10, 3)
|
||||
>>> # a batch of 2 samples of 4 indices each
|
||||
>>> input = torch.LongTensor([[1, 2, 4, 5], [4, 3, 2, 9]])
|
||||
>>> # xdoctest: +IGNORE_WANT("non-deterministic")
|
||||
>>> embedding(input)
|
||||
tensor([[[-0.0251, -1.6902, 0.7172],
|
||||
[-0.6431, 0.0748, 0.6969],
|
||||
[ 1.4970, 1.3448, -0.9685],
|
||||
[-0.3677, -2.7265, -0.1685]],
|
||||
|
||||
[[ 1.4970, 1.3448, -0.9685],
|
||||
[ 0.4362, -0.4004, 0.9400],
|
||||
[-0.6431, 0.0748, 0.6969],
|
||||
[ 0.9124, -2.3616, 1.1151]]])
|
||||
|
||||
|
||||
>>> # example with padding_idx
|
||||
>>> embedding = nn.Embedding(10, 3, padding_idx=0)
|
||||
>>> input = torch.LongTensor([[0, 2, 0, 5]])
|
||||
>>> embedding(input)
|
||||
tensor([[[ 0.0000, 0.0000, 0.0000],
|
||||
[ 0.1535, -2.0309, 0.9315],
|
||||
[ 0.0000, 0.0000, 0.0000],
|
||||
[-0.1655, 0.9897, 0.0635]]])
|
||||
|
||||
>>> # example of changing `pad` vector
|
||||
>>> padding_idx = 0
|
||||
>>> embedding = nn.Embedding(3, 3, padding_idx=padding_idx)
|
||||
>>> embedding.weight
|
||||
Parameter containing:
|
||||
tensor([[ 0.0000, 0.0000, 0.0000],
|
||||
[-0.7895, -0.7089, -0.0364],
|
||||
[ 0.6778, 0.5803, 0.2678]], requires_grad=True)
|
||||
>>> with torch.no_grad():
|
||||
... embedding.weight[padding_idx] = torch.ones(3)
|
||||
>>> embedding.weight
|
||||
Parameter containing:
|
||||
tensor([[ 1.0000, 1.0000, 1.0000],
|
||||
[-0.7895, -0.7089, -0.0364],
|
||||
[ 0.6778, 0.5803, 0.2678]], requires_grad=True)
|
||||
"""
|
||||
|
||||
__constants__ = [
|
||||
"num_embeddings",
|
||||
"embedding_dim",
|
||||
"padding_idx",
|
||||
"max_norm",
|
||||
"norm_type",
|
||||
"scale_grad_by_freq",
|
||||
"sparse",
|
||||
]
|
||||
|
||||
num_embeddings: int
|
||||
embedding_dim: int
|
||||
padding_idx: int | None
|
||||
max_norm: float | None
|
||||
norm_type: float
|
||||
scale_grad_by_freq: bool
|
||||
weight: Tensor
|
||||
freeze: bool
|
||||
sparse: bool
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_embeddings: int,
|
||||
embedding_dim: int,
|
||||
padding_idx: int | None = None,
|
||||
max_norm: float | None = None,
|
||||
norm_type: float = 2.0,
|
||||
scale_grad_by_freq: bool = False,
|
||||
sparse: bool = False,
|
||||
_weight: Tensor | None = None,
|
||||
_freeze: bool = False,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
self.num_embeddings = num_embeddings
|
||||
self.embedding_dim = embedding_dim
|
||||
if padding_idx is not None:
|
||||
if padding_idx > 0:
|
||||
if padding_idx >= self.num_embeddings:
|
||||
raise AssertionError("Padding_idx must be within num_embeddings")
|
||||
elif padding_idx < 0:
|
||||
if padding_idx < -self.num_embeddings:
|
||||
raise AssertionError("Padding_idx must be within num_embeddings")
|
||||
padding_idx = self.num_embeddings + padding_idx
|
||||
self.padding_idx = padding_idx
|
||||
self.max_norm = max_norm
|
||||
self.norm_type = norm_type
|
||||
self.scale_grad_by_freq = scale_grad_by_freq
|
||||
if _weight is None:
|
||||
self.weight = Parameter(
|
||||
torch.empty((num_embeddings, embedding_dim), **factory_kwargs),
|
||||
requires_grad=not _freeze,
|
||||
)
|
||||
self.reset_parameters()
|
||||
else:
|
||||
if list(_weight.shape) != [num_embeddings, embedding_dim]:
|
||||
raise AssertionError(
|
||||
"Shape of weight does not match num_embeddings and embedding_dim"
|
||||
)
|
||||
self.weight = Parameter(_weight, requires_grad=not _freeze)
|
||||
|
||||
self.sparse = sparse
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
init.normal_(self.weight)
|
||||
self._fill_padding_idx_with_zero()
|
||||
|
||||
def _fill_padding_idx_with_zero(self) -> None:
|
||||
if self.padding_idx is not None:
|
||||
with torch.no_grad():
|
||||
self.weight[self.padding_idx].fill_(0)
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
return F.embedding(
|
||||
input,
|
||||
self.weight,
|
||||
self.padding_idx,
|
||||
self.max_norm,
|
||||
self.norm_type,
|
||||
self.scale_grad_by_freq,
|
||||
self.sparse,
|
||||
)
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
s = "{num_embeddings}, {embedding_dim}"
|
||||
if self.padding_idx is not None:
|
||||
s += ", padding_idx={padding_idx}"
|
||||
if self.max_norm is not None:
|
||||
s += ", max_norm={max_norm}"
|
||||
if self.norm_type != 2:
|
||||
s += ", norm_type={norm_type}"
|
||||
if self.scale_grad_by_freq is not False:
|
||||
s += ", scale_grad_by_freq={scale_grad_by_freq}"
|
||||
if self.sparse is not False:
|
||||
s += ", sparse=True"
|
||||
return s.format(**self.__dict__)
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
cls,
|
||||
embeddings,
|
||||
freeze=True,
|
||||
padding_idx=None,
|
||||
max_norm=None,
|
||||
norm_type=2.0,
|
||||
scale_grad_by_freq=False,
|
||||
sparse=False,
|
||||
):
|
||||
r"""Create Embedding instance from given 2-dimensional FloatTensor.
|
||||
|
||||
Args:
|
||||
embeddings (Tensor): FloatTensor containing weights for the Embedding.
|
||||
First dimension is being passed to Embedding as ``num_embeddings``, second as ``embedding_dim``.
|
||||
freeze (bool, optional): If ``True``, the tensor does not get updated in the learning process.
|
||||
Equivalent to ``embedding.weight.requires_grad = False``. Default: ``True``
|
||||
padding_idx (int, optional): If specified, the entries at :attr:`padding_idx` do not contribute to the gradient;
|
||||
therefore, the embedding vector at :attr:`padding_idx` is not updated during training,
|
||||
i.e. it remains as a fixed "pad".
|
||||
max_norm (float, optional): See module initialization documentation.
|
||||
norm_type (float, optional): See module initialization documentation. Default ``2``.
|
||||
scale_grad_by_freq (bool, optional): See module initialization documentation. Default ``False``.
|
||||
sparse (bool, optional): See module initialization documentation.
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # FloatTensor containing pretrained weights
|
||||
>>> weight = torch.FloatTensor([[1, 2.3, 3], [4, 5.1, 6.3]])
|
||||
>>> embedding = nn.Embedding.from_pretrained(weight)
|
||||
>>> # Get embeddings for index 1
|
||||
>>> input = torch.LongTensor([1])
|
||||
>>> # xdoctest: +IGNORE_WANT("non-deterministic")
|
||||
>>> embedding(input)
|
||||
tensor([[ 4.0000, 5.1000, 6.3000]])
|
||||
"""
|
||||
if embeddings.dim() != 2:
|
||||
raise AssertionError("Embeddings parameter is expected to be 2-dimensional")
|
||||
rows, cols = embeddings.shape
|
||||
embedding = cls(
|
||||
num_embeddings=rows,
|
||||
embedding_dim=cols,
|
||||
_weight=embeddings,
|
||||
_freeze=freeze,
|
||||
padding_idx=padding_idx,
|
||||
max_norm=max_norm,
|
||||
norm_type=norm_type,
|
||||
scale_grad_by_freq=scale_grad_by_freq,
|
||||
sparse=sparse,
|
||||
)
|
||||
return embedding
|
||||
|
||||
|
||||
class EmbeddingBag(Module):
|
||||
r"""Compute sums or means of 'bags' of embeddings, without instantiating the intermediate embeddings.
|
||||
|
||||
For bags of constant length, no :attr:`per_sample_weights`, no indices equal to :attr:`padding_idx`,
|
||||
and with 2D inputs, this class
|
||||
|
||||
* with ``mode="sum"`` is equivalent to :class:`~torch.nn.Embedding` followed by ``torch.sum(dim=1)``,
|
||||
* with ``mode="mean"`` is equivalent to :class:`~torch.nn.Embedding` followed by ``torch.mean(dim=1)``,
|
||||
* with ``mode="max"`` is equivalent to :class:`~torch.nn.Embedding` followed by ``torch.max(dim=1)``.
|
||||
|
||||
However, :class:`~torch.nn.EmbeddingBag` is much more time and memory efficient than using a chain of these
|
||||
operations.
|
||||
|
||||
EmbeddingBag also supports per-sample weights as an argument to the forward
|
||||
pass. This scales the output of the Embedding before performing a weighted
|
||||
reduction as specified by ``mode``. If :attr:`per_sample_weights` is passed, the
|
||||
only supported ``mode`` is ``"sum"``, which computes a weighted sum according to
|
||||
:attr:`per_sample_weights`.
|
||||
|
||||
Args:
|
||||
num_embeddings (int): size of the dictionary of embeddings
|
||||
embedding_dim (int): the size of each embedding vector
|
||||
max_norm (float, optional): If given, each embedding vector with norm larger than :attr:`max_norm`
|
||||
is renormalized to have norm :attr:`max_norm`.
|
||||
norm_type (float, optional): The p of the p-norm to compute for the :attr:`max_norm` option. Default ``2``.
|
||||
scale_grad_by_freq (bool, optional): if given, this will scale gradients by the inverse of frequency of
|
||||
the words in the mini-batch. Default ``False``.
|
||||
Note: this option is not supported when ``mode="max"``.
|
||||
mode (str, optional): ``"sum"``, ``"mean"`` or ``"max"``. Specifies the way to reduce the bag.
|
||||
``"sum"`` computes the weighted sum, taking :attr:`per_sample_weights`
|
||||
into consideration. ``"mean"`` computes the average of the values
|
||||
in the bag, ``"max"`` computes the max value over each bag.
|
||||
Default: ``"mean"``
|
||||
sparse (bool, optional): if ``True``, gradient w.r.t. :attr:`weight` matrix will be a sparse tensor. See
|
||||
Notes for more details regarding sparse gradients. Note: this option is not
|
||||
supported when ``mode="max"``.
|
||||
include_last_offset (bool, optional): if ``True``, the size of offsets is equal to the number of bags + 1.
|
||||
The last element is the size of the input, or the ending index position
|
||||
of the last bag (sequence). This matches the CSR format. Ignored when
|
||||
input is 2D. Default ``False``.
|
||||
padding_idx (int, optional): If specified, the entries at :attr:`padding_idx` do not contribute to the
|
||||
gradient; therefore, the embedding vector at :attr:`padding_idx` is not updated
|
||||
during training, i.e. it remains as a fixed "pad". For a newly constructed
|
||||
EmbeddingBag, the embedding vector at :attr:`padding_idx` will default to all
|
||||
zeros, but can be updated to another value to be used as the padding vector.
|
||||
Note that the embedding vector at :attr:`padding_idx` is excluded from the
|
||||
reduction.
|
||||
|
||||
Attributes:
|
||||
weight (Tensor): the learnable weights of the module of shape `(num_embeddings, embedding_dim)`
|
||||
initialized from :math:`\mathcal{N}(0, 1)`.
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # an EmbeddingBag module containing 10 tensors of size 3
|
||||
>>> embedding_sum = nn.EmbeddingBag(10, 3, mode='sum')
|
||||
>>> # a batch of 2 samples of 4 indices each
|
||||
>>> input = torch.tensor([1, 2, 4, 5, 4, 3, 2, 9], dtype=torch.long)
|
||||
>>> offsets = torch.tensor([0, 4], dtype=torch.long)
|
||||
>>> # xdoctest: +IGNORE_WANT("non-deterministic")
|
||||
>>> embedding_sum(input, offsets)
|
||||
tensor([[-0.8861, -5.4350, -0.0523],
|
||||
[ 1.1306, -2.5798, -1.0044]])
|
||||
|
||||
>>> # Example with padding_idx
|
||||
>>> embedding_sum = nn.EmbeddingBag(10, 3, mode='sum', padding_idx=2)
|
||||
>>> input = torch.tensor([2, 2, 2, 2, 4, 3, 2, 9], dtype=torch.long)
|
||||
>>> offsets = torch.tensor([0, 4], dtype=torch.long)
|
||||
>>> embedding_sum(input, offsets)
|
||||
tensor([[ 0.0000, 0.0000, 0.0000],
|
||||
[-0.7082, 3.2145, -2.6251]])
|
||||
|
||||
>>> # An EmbeddingBag can be loaded from an Embedding like so
|
||||
>>> embedding = nn.Embedding(10, 3, padding_idx=2)
|
||||
>>> embedding_sum = nn.EmbeddingBag.from_pretrained(
|
||||
embedding.weight,
|
||||
padding_idx=embedding.padding_idx,
|
||||
mode='sum')
|
||||
"""
|
||||
|
||||
__constants__ = [
|
||||
"num_embeddings",
|
||||
"embedding_dim",
|
||||
"max_norm",
|
||||
"norm_type",
|
||||
"scale_grad_by_freq",
|
||||
"mode",
|
||||
"sparse",
|
||||
"include_last_offset",
|
||||
"padding_idx",
|
||||
]
|
||||
|
||||
num_embeddings: int
|
||||
embedding_dim: int
|
||||
max_norm: float | None
|
||||
norm_type: float
|
||||
scale_grad_by_freq: bool
|
||||
weight: Tensor
|
||||
mode: str
|
||||
sparse: bool
|
||||
include_last_offset: bool
|
||||
padding_idx: int | None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_embeddings: int,
|
||||
embedding_dim: int,
|
||||
max_norm: float | None = None,
|
||||
norm_type: float = 2.0,
|
||||
scale_grad_by_freq: bool = False,
|
||||
mode: str = "mean",
|
||||
sparse: bool = False,
|
||||
_weight: Tensor | None = None,
|
||||
include_last_offset: bool = False,
|
||||
padding_idx: int | None = None,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
self.num_embeddings = num_embeddings
|
||||
self.embedding_dim = embedding_dim
|
||||
self.max_norm = max_norm
|
||||
self.norm_type = norm_type
|
||||
self.scale_grad_by_freq = scale_grad_by_freq
|
||||
if padding_idx is not None:
|
||||
if padding_idx > 0:
|
||||
if padding_idx >= self.num_embeddings:
|
||||
raise AssertionError("padding_idx must be within num_embeddings")
|
||||
elif padding_idx < 0:
|
||||
if padding_idx < -self.num_embeddings:
|
||||
raise AssertionError("padding_idx must be within num_embeddings")
|
||||
padding_idx = self.num_embeddings + padding_idx
|
||||
self.padding_idx = padding_idx
|
||||
if _weight is None:
|
||||
self.weight = Parameter(
|
||||
torch.empty((num_embeddings, embedding_dim), **factory_kwargs)
|
||||
)
|
||||
self.reset_parameters()
|
||||
else:
|
||||
if list(_weight.shape) != [num_embeddings, embedding_dim]:
|
||||
raise AssertionError(
|
||||
"Shape of weight does not match num_embeddings and embedding_dim"
|
||||
)
|
||||
self.weight = Parameter(_weight)
|
||||
self.mode = mode
|
||||
self.sparse = sparse
|
||||
self.include_last_offset = include_last_offset
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
init.normal_(self.weight)
|
||||
self._fill_padding_idx_with_zero()
|
||||
|
||||
def _fill_padding_idx_with_zero(self) -> None:
|
||||
if self.padding_idx is not None:
|
||||
with torch.no_grad():
|
||||
self.weight[self.padding_idx].fill_(0)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input: Tensor,
|
||||
offsets: Tensor | None = None,
|
||||
per_sample_weights: Tensor | None = None,
|
||||
) -> Tensor:
|
||||
"""Forward pass of EmbeddingBag.
|
||||
|
||||
Args:
|
||||
input (Tensor): Tensor containing bags of indices into the embedding matrix.
|
||||
offsets (Tensor, optional): Only used when :attr:`input` is 1D. :attr:`offsets` determines
|
||||
the starting index position of each bag (sequence) in :attr:`input`.
|
||||
per_sample_weights (Tensor, optional): a tensor of float / double weights, or None
|
||||
to indicate all weights should be taken to be ``1``. If specified, :attr:`per_sample_weights`
|
||||
must have exactly the same shape as input and is treated as having the same
|
||||
:attr:`offsets`, if those are not ``None``. Only supported for ``mode='sum'``.
|
||||
|
||||
Returns:
|
||||
Tensor output shape of `(B, embedding_dim)`.
|
||||
|
||||
.. note::
|
||||
|
||||
A few notes about ``input`` and ``offsets``:
|
||||
|
||||
- :attr:`input` and :attr:`offsets` have to be of the same type, either int or long
|
||||
|
||||
- If :attr:`input` is 2D of shape `(B, N)`, it will be treated as ``B`` bags (sequences)
|
||||
each of fixed length ``N``, and this will return ``B`` values aggregated in a way
|
||||
depending on the :attr:`mode`. :attr:`offsets` is ignored and required to be ``None`` in this case.
|
||||
|
||||
- If :attr:`input` is 1D of shape `(N)`, it will be treated as a concatenation of
|
||||
multiple bags (sequences). :attr:`offsets` is required to be a 1D tensor containing the
|
||||
starting index positions of each bag in :attr:`input`. Therefore, for :attr:`offsets` of shape `(B)`,
|
||||
:attr:`input` will be viewed as having ``B`` bags. Empty bags (i.e., having 0-length) will have
|
||||
returned vectors filled by zeros.
|
||||
"""
|
||||
return F.embedding_bag(
|
||||
input,
|
||||
self.weight,
|
||||
offsets,
|
||||
self.max_norm,
|
||||
self.norm_type,
|
||||
self.scale_grad_by_freq,
|
||||
self.mode,
|
||||
self.sparse,
|
||||
per_sample_weights,
|
||||
self.include_last_offset,
|
||||
self.padding_idx,
|
||||
)
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
s = "{num_embeddings}, {embedding_dim}"
|
||||
if self.max_norm is not None:
|
||||
s += ", max_norm={max_norm}"
|
||||
if self.norm_type != 2:
|
||||
s += ", norm_type={norm_type}"
|
||||
if self.scale_grad_by_freq is not False:
|
||||
s += ", scale_grad_by_freq={scale_grad_by_freq}"
|
||||
s += ", mode={mode}"
|
||||
if self.padding_idx is not None:
|
||||
s += ", padding_idx={padding_idx}"
|
||||
return s.format(**{k: repr(v) for k, v in self.__dict__.items()})
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
cls,
|
||||
embeddings: Tensor,
|
||||
freeze: bool = True,
|
||||
max_norm: float | None = None,
|
||||
norm_type: float = 2.0,
|
||||
scale_grad_by_freq: bool = False,
|
||||
mode: str = "mean",
|
||||
sparse: bool = False,
|
||||
include_last_offset: bool = False,
|
||||
padding_idx: int | None = None,
|
||||
) -> "EmbeddingBag":
|
||||
r"""Create EmbeddingBag instance from given 2-dimensional FloatTensor.
|
||||
|
||||
Args:
|
||||
embeddings (Tensor): FloatTensor containing weights for the EmbeddingBag.
|
||||
First dimension is being passed to EmbeddingBag as 'num_embeddings', second as 'embedding_dim'.
|
||||
freeze (bool, optional): If ``True``, the tensor does not get updated in the learning process.
|
||||
Equivalent to ``embeddingbag.weight.requires_grad = False``. Default: ``True``
|
||||
max_norm (float, optional): See module initialization documentation. Default: ``None``
|
||||
norm_type (float, optional): See module initialization documentation. Default ``2``.
|
||||
scale_grad_by_freq (bool, optional): See module initialization documentation. Default ``False``.
|
||||
mode (str, optional): See module initialization documentation. Default: ``"mean"``
|
||||
sparse (bool, optional): See module initialization documentation. Default: ``False``.
|
||||
include_last_offset (bool, optional): See module initialization documentation. Default: ``False``.
|
||||
padding_idx (int, optional): See module initialization documentation. Default: ``None``.
|
||||
|
||||
Examples::
|
||||
|
||||
>>> # FloatTensor containing pretrained weights
|
||||
>>> weight = torch.FloatTensor([[1, 2.3, 3], [4, 5.1, 6.3]])
|
||||
>>> embeddingbag = nn.EmbeddingBag.from_pretrained(weight)
|
||||
>>> # Get embeddings for index 1
|
||||
>>> input = torch.LongTensor([[1, 0]])
|
||||
>>> # xdoctest: +IGNORE_WANT("non-deterministic")
|
||||
>>> embeddingbag(input)
|
||||
tensor([[ 2.5000, 3.7000, 4.6500]])
|
||||
"""
|
||||
if embeddings.dim() != 2:
|
||||
raise AssertionError("Embeddings parameter is expected to be 2-dimensional")
|
||||
rows, cols = embeddings.shape
|
||||
embeddingbag = cls(
|
||||
num_embeddings=rows,
|
||||
embedding_dim=cols,
|
||||
_weight=embeddings,
|
||||
max_norm=max_norm,
|
||||
norm_type=norm_type,
|
||||
scale_grad_by_freq=scale_grad_by_freq,
|
||||
mode=mode,
|
||||
sparse=sparse,
|
||||
include_last_offset=include_last_offset,
|
||||
padding_idx=padding_idx,
|
||||
)
|
||||
embeddingbag.weight.requires_grad = not freeze
|
||||
return embeddingbag
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,298 @@
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
from torch.nn.common_types import _ratio_2_t, _ratio_any_t, _size_2_t, _size_any_t
|
||||
|
||||
from .module import Module
|
||||
|
||||
|
||||
__all__ = ["Upsample", "UpsamplingNearest2d", "UpsamplingBilinear2d"]
|
||||
|
||||
|
||||
class Upsample(Module):
|
||||
r"""Upsamples a given multi-channel 1D (temporal), 2D (spatial) or 3D (volumetric) data.
|
||||
|
||||
The input data is assumed to be of the form
|
||||
`minibatch x channels x [optional depth] x [optional height] x width`.
|
||||
Hence, for spatial inputs, we expect a 4D Tensor and for volumetric inputs, we expect a 5D Tensor.
|
||||
|
||||
The algorithms available for upsampling are nearest neighbor and linear,
|
||||
bilinear, bicubic and trilinear for 3D, 4D and 5D input Tensor,
|
||||
respectively.
|
||||
|
||||
One can either give a :attr:`scale_factor` or the target output :attr:`size` to
|
||||
calculate the output size. (You cannot give both, as it is ambiguous)
|
||||
|
||||
Args:
|
||||
size (int or Tuple[int] or Tuple[int, int] or Tuple[int, int, int], optional):
|
||||
output spatial sizes
|
||||
scale_factor (float or Tuple[float] or Tuple[float, float] or Tuple[float, float, float], optional):
|
||||
multiplier for spatial size. Has to match input size if it is a tuple.
|
||||
mode (str, optional): the upsampling algorithm: one of ``'nearest'``,
|
||||
``'linear'``, ``'bilinear'``, ``'bicubic'`` and ``'trilinear'``.
|
||||
Default: ``'nearest'``
|
||||
align_corners (bool, optional): if ``True``, the corner pixels of the input
|
||||
and output tensors are aligned, and thus preserving the values at
|
||||
those pixels. This only has effect when :attr:`mode` is
|
||||
``'linear'``, ``'bilinear'``, ``'bicubic'``, or ``'trilinear'``.
|
||||
Default: ``False``
|
||||
recompute_scale_factor (bool, optional): recompute the scale_factor for use in the
|
||||
interpolation calculation. If `recompute_scale_factor` is ``True``, then
|
||||
`scale_factor` must be passed in and `scale_factor` is used to compute the
|
||||
output `size`. The computed output `size` will be used to infer new scales for
|
||||
the interpolation. Note that when `scale_factor` is floating-point, it may differ
|
||||
from the recomputed `scale_factor` due to rounding and precision issues.
|
||||
If `recompute_scale_factor` is ``False``, then `size` or `scale_factor` will
|
||||
be used directly for interpolation.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, W_{in})`, :math:`(N, C, H_{in}, W_{in})` or :math:`(N, C, D_{in}, H_{in}, W_{in})`
|
||||
- Output: :math:`(N, C, W_{out})`, :math:`(N, C, H_{out}, W_{out})`
|
||||
or :math:`(N, C, D_{out}, H_{out}, W_{out})`, where
|
||||
|
||||
.. math::
|
||||
D_{out} = \left\lfloor D_{in} \times \text{scale\_factor} \right\rfloor
|
||||
|
||||
.. math::
|
||||
H_{out} = \left\lfloor H_{in} \times \text{scale\_factor} \right\rfloor
|
||||
|
||||
.. math::
|
||||
W_{out} = \left\lfloor W_{in} \times \text{scale\_factor} \right\rfloor
|
||||
|
||||
.. warning::
|
||||
With ``align_corners = True``, the linearly interpolating modes
|
||||
(`linear`, `bilinear`, `bicubic`, and `trilinear`) don't proportionally
|
||||
align the output and input pixels, and thus the output values can depend
|
||||
on the input size. This was the default behavior for these modes up to
|
||||
version 0.3.1. Since then, the default behavior is
|
||||
``align_corners = False``. See below for concrete examples on how this
|
||||
affects the outputs.
|
||||
|
||||
.. note::
|
||||
If you want downsampling/general resizing, you should use :func:`~nn.functional.interpolate`.
|
||||
|
||||
Examples::
|
||||
|
||||
>>> input = torch.arange(1, 5, dtype=torch.float32).view(1, 1, 2, 2)
|
||||
>>> input
|
||||
tensor([[[[1., 2.],
|
||||
[3., 4.]]]])
|
||||
|
||||
>>> m = nn.Upsample(scale_factor=2, mode='nearest')
|
||||
>>> m(input)
|
||||
tensor([[[[1., 1., 2., 2.],
|
||||
[1., 1., 2., 2.],
|
||||
[3., 3., 4., 4.],
|
||||
[3., 3., 4., 4.]]]])
|
||||
|
||||
>>> # xdoctest: +IGNORE_WANT("other tests seem to modify printing styles")
|
||||
>>> m = nn.Upsample(scale_factor=2, mode='bilinear') # align_corners=False
|
||||
>>> m(input)
|
||||
tensor([[[[1.0000, 1.2500, 1.7500, 2.0000],
|
||||
[1.5000, 1.7500, 2.2500, 2.5000],
|
||||
[2.5000, 2.7500, 3.2500, 3.5000],
|
||||
[3.0000, 3.2500, 3.7500, 4.0000]]]])
|
||||
|
||||
>>> m = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
|
||||
>>> m(input)
|
||||
tensor([[[[1.0000, 1.3333, 1.6667, 2.0000],
|
||||
[1.6667, 2.0000, 2.3333, 2.6667],
|
||||
[2.3333, 2.6667, 3.0000, 3.3333],
|
||||
[3.0000, 3.3333, 3.6667, 4.0000]]]])
|
||||
|
||||
>>> # Try scaling the same data in a larger tensor
|
||||
>>> input_3x3 = torch.zeros(3, 3).view(1, 1, 3, 3)
|
||||
>>> input_3x3[:, :, :2, :2].copy_(input)
|
||||
tensor([[[[1., 2.],
|
||||
[3., 4.]]]])
|
||||
>>> input_3x3
|
||||
tensor([[[[1., 2., 0.],
|
||||
[3., 4., 0.],
|
||||
[0., 0., 0.]]]])
|
||||
|
||||
>>> # xdoctest: +IGNORE_WANT("seems to fail when other tests are run in the same session")
|
||||
>>> m = nn.Upsample(scale_factor=2, mode='bilinear') # align_corners=False
|
||||
>>> # Notice that values in top left corner are the same with the small input (except at boundary)
|
||||
>>> m(input_3x3)
|
||||
tensor([[[[1.0000, 1.2500, 1.7500, 1.5000, 0.5000, 0.0000],
|
||||
[1.5000, 1.7500, 2.2500, 1.8750, 0.6250, 0.0000],
|
||||
[2.5000, 2.7500, 3.2500, 2.6250, 0.8750, 0.0000],
|
||||
[2.2500, 2.4375, 2.8125, 2.2500, 0.7500, 0.0000],
|
||||
[0.7500, 0.8125, 0.9375, 0.7500, 0.2500, 0.0000],
|
||||
[0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000]]]])
|
||||
|
||||
>>> m = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
|
||||
>>> # Notice that values in top left corner are now changed
|
||||
>>> m(input_3x3)
|
||||
tensor([[[[1.0000, 1.4000, 1.8000, 1.6000, 0.8000, 0.0000],
|
||||
[1.8000, 2.2000, 2.6000, 2.2400, 1.1200, 0.0000],
|
||||
[2.6000, 3.0000, 3.4000, 2.8800, 1.4400, 0.0000],
|
||||
[2.4000, 2.7200, 3.0400, 2.5600, 1.2800, 0.0000],
|
||||
[1.2000, 1.3600, 1.5200, 1.2800, 0.6400, 0.0000],
|
||||
[0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000]]]])
|
||||
"""
|
||||
|
||||
__constants__ = [
|
||||
"size",
|
||||
"scale_factor",
|
||||
"mode",
|
||||
"align_corners",
|
||||
"name",
|
||||
"recompute_scale_factor",
|
||||
]
|
||||
name: str
|
||||
size: _size_any_t | None
|
||||
scale_factor: _ratio_any_t | None
|
||||
mode: str
|
||||
align_corners: bool | None
|
||||
recompute_scale_factor: bool | None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
size: _size_any_t | None = None,
|
||||
scale_factor: _ratio_any_t | None = None,
|
||||
mode: str = "nearest",
|
||||
align_corners: bool | None = None,
|
||||
recompute_scale_factor: bool | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.name = type(self).__name__
|
||||
self.size = size
|
||||
if isinstance(scale_factor, tuple):
|
||||
self.scale_factor = tuple(float(factor) for factor in scale_factor)
|
||||
else:
|
||||
self.scale_factor = float(scale_factor) if scale_factor else None
|
||||
self.mode = mode
|
||||
self.align_corners = align_corners
|
||||
self.recompute_scale_factor = recompute_scale_factor
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
"""
|
||||
Runs the forward pass.
|
||||
"""
|
||||
return F.interpolate(
|
||||
input,
|
||||
self.size,
|
||||
self.scale_factor,
|
||||
self.mode,
|
||||
self.align_corners,
|
||||
recompute_scale_factor=self.recompute_scale_factor,
|
||||
)
|
||||
|
||||
def __setstate__(self, state):
|
||||
if "recompute_scale_factor" not in state:
|
||||
state["recompute_scale_factor"] = True
|
||||
|
||||
super().__setstate__(state)
|
||||
|
||||
def extra_repr(self) -> str:
|
||||
"""
|
||||
Return the extra representation of the module.
|
||||
"""
|
||||
if self.scale_factor is not None:
|
||||
info = "scale_factor=" + repr(self.scale_factor)
|
||||
else:
|
||||
info = "size=" + repr(self.size)
|
||||
info += ", mode=" + repr(self.mode)
|
||||
return info
|
||||
|
||||
|
||||
class UpsamplingNearest2d(Upsample):
|
||||
r"""Applies a 2D nearest neighbor upsampling to an input signal composed of several input channels.
|
||||
|
||||
To specify the scale, it takes either the :attr:`size` or the :attr:`scale_factor`
|
||||
as it's constructor argument.
|
||||
|
||||
When :attr:`size` is given, it is the output size of the image `(h, w)`.
|
||||
|
||||
Args:
|
||||
size (int or Tuple[int, int], optional): output spatial sizes
|
||||
scale_factor (float or Tuple[float, float], optional): multiplier for
|
||||
spatial size.
|
||||
|
||||
.. warning::
|
||||
This class is deprecated in favor of :func:`~nn.functional.interpolate`.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, H_{in}, W_{in})`
|
||||
- Output: :math:`(N, C, H_{out}, W_{out})` where
|
||||
|
||||
.. math::
|
||||
H_{out} = \left\lfloor H_{in} \times \text{scale\_factor} \right\rfloor
|
||||
|
||||
.. math::
|
||||
W_{out} = \left\lfloor W_{in} \times \text{scale\_factor} \right\rfloor
|
||||
|
||||
Examples::
|
||||
|
||||
>>> input = torch.arange(1, 5, dtype=torch.float32).view(1, 1, 2, 2)
|
||||
>>> input
|
||||
tensor([[[[1., 2.],
|
||||
[3., 4.]]]])
|
||||
|
||||
>>> m = nn.UpsamplingNearest2d(scale_factor=2)
|
||||
>>> m(input)
|
||||
tensor([[[[1., 1., 2., 2.],
|
||||
[1., 1., 2., 2.],
|
||||
[3., 3., 4., 4.],
|
||||
[3., 3., 4., 4.]]]])
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
size: _size_2_t | None = None,
|
||||
scale_factor: _ratio_2_t | None = None,
|
||||
) -> None:
|
||||
super().__init__(size, scale_factor, mode="nearest")
|
||||
|
||||
|
||||
class UpsamplingBilinear2d(Upsample):
|
||||
r"""Applies a 2D bilinear upsampling to an input signal composed of several input channels.
|
||||
|
||||
To specify the scale, it takes either the :attr:`size` or the :attr:`scale_factor`
|
||||
as it's constructor argument.
|
||||
|
||||
When :attr:`size` is given, it is the output size of the image `(h, w)`.
|
||||
|
||||
Args:
|
||||
size (int or Tuple[int, int], optional): output spatial sizes
|
||||
scale_factor (float or Tuple[float, float], optional): multiplier for
|
||||
spatial size.
|
||||
|
||||
.. warning::
|
||||
This class is deprecated in favor of :func:`~nn.functional.interpolate`. It is
|
||||
equivalent to ``nn.functional.interpolate(..., mode='bilinear', align_corners=True)``.
|
||||
|
||||
Shape:
|
||||
- Input: :math:`(N, C, H_{in}, W_{in})`
|
||||
- Output: :math:`(N, C, H_{out}, W_{out})` where
|
||||
|
||||
.. math::
|
||||
H_{out} = \left\lfloor H_{in} \times \text{scale\_factor} \right\rfloor
|
||||
|
||||
.. math::
|
||||
W_{out} = \left\lfloor W_{in} \times \text{scale\_factor} \right\rfloor
|
||||
|
||||
Examples::
|
||||
|
||||
>>> input = torch.arange(1, 5, dtype=torch.float32).view(1, 1, 2, 2)
|
||||
>>> input
|
||||
tensor([[[[1., 2.],
|
||||
[3., 4.]]]])
|
||||
|
||||
>>> # xdoctest: +IGNORE_WANT("do other tests modify the global state?")
|
||||
>>> m = nn.UpsamplingBilinear2d(scale_factor=2)
|
||||
>>> m(input)
|
||||
tensor([[[[1.0000, 1.3333, 1.6667, 2.0000],
|
||||
[1.6667, 2.0000, 2.3333, 2.6667],
|
||||
[2.3333, 2.6667, 3.0000, 3.3333],
|
||||
[3.0000, 3.3333, 3.6667, 4.0000]]]])
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
size: _size_2_t | None = None,
|
||||
scale_factor: _ratio_2_t | None = None,
|
||||
) -> None:
|
||||
super().__init__(size, scale_factor, mode="bilinear", align_corners=True)
|
||||
@@ -0,0 +1,82 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import collections
|
||||
from itertools import repeat
|
||||
from typing import Any
|
||||
|
||||
|
||||
__all__ = ["consume_prefix_in_state_dict_if_present"]
|
||||
|
||||
|
||||
def _ntuple(n, name="parse"):
|
||||
def parse(x):
|
||||
if isinstance(x, collections.abc.Iterable):
|
||||
return tuple(x)
|
||||
return tuple(repeat(x, n))
|
||||
|
||||
parse.__name__ = name
|
||||
return parse
|
||||
|
||||
|
||||
_single = _ntuple(1, "_single")
|
||||
_pair = _ntuple(2, "_pair")
|
||||
_triple = _ntuple(3, "_triple")
|
||||
_quadruple = _ntuple(4, "_quadruple")
|
||||
|
||||
|
||||
def _reverse_repeat_tuple(t, n):
|
||||
r"""Reverse the order of `t` and repeat each element for `n` times.
|
||||
|
||||
This can be used to translate padding arg used by Conv and Pooling modules
|
||||
to the ones used by `F.pad`.
|
||||
"""
|
||||
return tuple(x for x in reversed(t) for _ in range(n))
|
||||
|
||||
|
||||
def _list_with_default(out_size: list[int], defaults: list[int]) -> list[int]:
|
||||
import torch
|
||||
|
||||
if isinstance(out_size, (int, torch.SymInt)):
|
||||
return out_size
|
||||
if len(defaults) <= len(out_size):
|
||||
raise ValueError(f"Input dimension should be at least {len(out_size) + 1}")
|
||||
return [
|
||||
v if v is not None else d
|
||||
for v, d in zip(out_size, defaults[-len(out_size) :], strict=False)
|
||||
]
|
||||
|
||||
|
||||
def consume_prefix_in_state_dict_if_present(
|
||||
state_dict: dict[str, Any],
|
||||
prefix: str,
|
||||
) -> None:
|
||||
r"""Strip the prefix in state_dict in place, if any.
|
||||
|
||||
.. note::
|
||||
Given a `state_dict` from a DP/DDP model, a local model can load it by applying
|
||||
`consume_prefix_in_state_dict_if_present(state_dict, "module.")` before calling
|
||||
:meth:`torch.nn.Module.load_state_dict`.
|
||||
|
||||
Args:
|
||||
state_dict (OrderedDict): a state-dict to be loaded to the model.
|
||||
prefix (str): prefix.
|
||||
"""
|
||||
keys = list(state_dict.keys())
|
||||
for key in keys:
|
||||
if key.startswith(prefix):
|
||||
newkey = key[len(prefix) :]
|
||||
state_dict[newkey] = state_dict.pop(key)
|
||||
|
||||
# also strip the prefix in metadata if any.
|
||||
if hasattr(state_dict, "_metadata"):
|
||||
keys = list(state_dict._metadata.keys())
|
||||
for key in keys:
|
||||
# for the metadata dict, the key can be:
|
||||
# '': for the DDP module, which we want to remove.
|
||||
# 'module': for the actual model.
|
||||
# 'module.xx.xx': for the rest.
|
||||
if len(key) == 0:
|
||||
continue
|
||||
# handling both, 'module' case and 'module.' cases
|
||||
if key == prefix.replace(".", "") or key.startswith(prefix):
|
||||
newkey = key[len(prefix) :]
|
||||
state_dict._metadata[newkey] = state_dict._metadata.pop(key)
|
||||
@@ -0,0 +1,27 @@
|
||||
from typing_extensions import deprecated
|
||||
|
||||
from torch.nn.parallel.data_parallel import data_parallel, DataParallel
|
||||
from torch.nn.parallel.distributed import DistributedDataParallel
|
||||
from torch.nn.parallel.parallel_apply import parallel_apply
|
||||
from torch.nn.parallel.replicate import replicate
|
||||
from torch.nn.parallel.scatter_gather import gather, scatter
|
||||
|
||||
|
||||
__all__ = [
|
||||
"replicate",
|
||||
"scatter",
|
||||
"parallel_apply",
|
||||
"gather",
|
||||
"data_parallel",
|
||||
"DataParallel",
|
||||
"DistributedDataParallel",
|
||||
]
|
||||
|
||||
|
||||
@deprecated(
|
||||
"`torch.nn.parallel.DistributedDataParallelCPU` is deprecated, "
|
||||
"please use `torch.nn.parallel.DistributedDataParallel` instead.",
|
||||
category=FutureWarning,
|
||||
)
|
||||
class DistributedDataParallelCPU(DistributedDataParallel):
|
||||
pass
|
||||
@@ -0,0 +1,176 @@
|
||||
import warnings
|
||||
from itertools import chain
|
||||
|
||||
import torch
|
||||
from torch._utils import _get_device_index
|
||||
from torch.autograd import Function
|
||||
from torch.nn.parallel import comm
|
||||
|
||||
|
||||
class Broadcast(Function):
|
||||
@staticmethod
|
||||
def forward(ctx, target_gpus, *inputs):
|
||||
if not all(i.device.type != "cpu" for i in inputs):
|
||||
raise AssertionError("Broadcast function not implemented for CPU tensors")
|
||||
target_gpus = [_get_device_index(x, True) for x in target_gpus]
|
||||
ctx.target_gpus = target_gpus
|
||||
if len(inputs) == 0:
|
||||
return ()
|
||||
ctx.num_inputs = len(inputs)
|
||||
ctx.input_device = inputs[0].get_device()
|
||||
|
||||
ctx.complex_mask = [inp.is_complex() for inp in inputs]
|
||||
|
||||
outputs = comm.broadcast_coalesced(inputs, ctx.target_gpus)
|
||||
|
||||
for device_outputs in outputs:
|
||||
for i, is_complex in enumerate(ctx.complex_mask):
|
||||
if is_complex:
|
||||
device_outputs[i] = torch.view_as_complex(device_outputs[i])
|
||||
|
||||
non_differentiables = []
|
||||
for idx, input_requires_grad in enumerate(ctx.needs_input_grad[1:]):
|
||||
if not input_requires_grad:
|
||||
non_differentiables.extend(output[idx] for output in outputs)
|
||||
ctx.mark_non_differentiable(*non_differentiables)
|
||||
return tuple(chain.from_iterable(outputs))
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *grad_outputs):
|
||||
grads = ReduceAddCoalesced.apply(
|
||||
ctx.input_device, ctx.num_inputs, *grad_outputs
|
||||
)
|
||||
|
||||
return (None,) + grads
|
||||
|
||||
|
||||
class ReduceAddCoalesced(Function):
|
||||
@staticmethod
|
||||
def forward(ctx, destination, num_inputs, *grads):
|
||||
ctx.target_gpus = [
|
||||
grads[i].get_device() for i in range(0, len(grads), num_inputs)
|
||||
]
|
||||
|
||||
complex_mask = [grads[i].is_complex() for i in range(num_inputs)]
|
||||
ctx.complex_mask = complex_mask
|
||||
|
||||
grads_converted = tuple(
|
||||
torch.view_as_real(g) if g.is_complex() else g for g in grads
|
||||
)
|
||||
|
||||
grads_ = [
|
||||
grads_converted[i : i + num_inputs]
|
||||
for i in range(0, len(grads_converted), num_inputs)
|
||||
]
|
||||
results = comm.reduce_add_coalesced(grads_, destination)
|
||||
|
||||
results = tuple(
|
||||
torch.view_as_complex(r) if is_complex else r
|
||||
for r, is_complex in zip(results, complex_mask)
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *grad_outputs):
|
||||
return (
|
||||
None,
|
||||
None,
|
||||
) + Broadcast.apply(ctx.target_gpus, *grad_outputs)
|
||||
|
||||
|
||||
class Gather(Function):
|
||||
@staticmethod
|
||||
def forward(ctx, target_device, dim, *inputs):
|
||||
if not all(i.device.type != "cpu" for i in inputs):
|
||||
raise AssertionError("Gather function not implemented for CPU tensors")
|
||||
if target_device == "cpu":
|
||||
ctx.target_device = "cpu"
|
||||
else:
|
||||
target_device = _get_device_index(target_device, True)
|
||||
ctx.target_device = target_device
|
||||
ctx.dim = dim
|
||||
ctx.input_gpus = tuple(i.get_device() for i in inputs)
|
||||
if all(t.dim() == 0 for t in inputs) and dim == 0:
|
||||
inputs = tuple(t.view(1) for t in inputs)
|
||||
warnings.warn(
|
||||
"Was asked to gather along dimension 0, but all "
|
||||
"input tensors were scalars; will instead unsqueeze "
|
||||
"and return a vector.",
|
||||
stacklevel=2,
|
||||
)
|
||||
ctx.unsqueezed_scalar = True
|
||||
else:
|
||||
ctx.unsqueezed_scalar = False
|
||||
ctx.input_sizes = tuple(i.size(ctx.dim) for i in inputs)
|
||||
|
||||
is_complex = len(inputs) > 0 and inputs[0].is_complex()
|
||||
|
||||
output = comm.gather(inputs, ctx.dim, ctx.target_device)
|
||||
|
||||
if is_complex:
|
||||
output = torch.view_as_complex(output)
|
||||
|
||||
return output
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
scattered_grads = Scatter.apply(
|
||||
ctx.input_gpus, ctx.input_sizes, ctx.dim, grad_output
|
||||
)
|
||||
if ctx.unsqueezed_scalar:
|
||||
scattered_grads = tuple(g[0] for g in scattered_grads)
|
||||
return (None, None) + scattered_grads
|
||||
|
||||
|
||||
class Scatter(Function):
|
||||
@staticmethod
|
||||
def forward(ctx, target_gpus, chunk_sizes, dim, input):
|
||||
target_gpus = [_get_device_index(x, True) for x in target_gpus]
|
||||
ctx.dim = dim
|
||||
ctx.input_device = input.get_device() if input.device.type != "cpu" else -1
|
||||
streams = None
|
||||
if torch.accelerator.is_available() and ctx.input_device == -1:
|
||||
# Perform CPU to GPU copies in a background stream
|
||||
streams = [_get_stream(torch.device(device)) for device in target_gpus]
|
||||
|
||||
is_complex = input.is_complex()
|
||||
|
||||
outputs = comm.scatter(input, target_gpus, chunk_sizes, ctx.dim, streams)
|
||||
|
||||
if is_complex:
|
||||
outputs = tuple(torch.view_as_complex(o) for o in outputs)
|
||||
|
||||
# Synchronize with the copy stream
|
||||
if streams is not None:
|
||||
for i, output in enumerate(outputs):
|
||||
with torch.accelerator.device_index(target_gpus[i]):
|
||||
main_stream = torch.accelerator.current_stream()
|
||||
main_stream.wait_stream(streams[i])
|
||||
output.record_stream(main_stream)
|
||||
return outputs
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *grad_output):
|
||||
return None, None, None, Gather.apply(ctx.input_device, ctx.dim, *grad_output)
|
||||
|
||||
|
||||
# background streams used for copying
|
||||
_streams: list[torch.Stream | None] | None = None
|
||||
|
||||
|
||||
def _get_stream(device: torch.device):
|
||||
"""Get a background stream for copying between CPU and target device."""
|
||||
global _streams
|
||||
if device.type == "cpu" or not torch.accelerator.is_available():
|
||||
return None
|
||||
if torch.accelerator.current_accelerator().type != device.type:
|
||||
raise AssertionError(
|
||||
f"Expected current accelerator type {torch.accelerator.current_accelerator().type} "
|
||||
f"to match device type {device.type}"
|
||||
)
|
||||
if _streams is None:
|
||||
_streams = [None] * torch.accelerator.device_count()
|
||||
if _streams[device.index] is None:
|
||||
_streams[device.index] = torch.Stream(device.index)
|
||||
return _streams[device.index]
|
||||
@@ -0,0 +1,264 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
from torch._utils import (
|
||||
_flatten_dense_tensors,
|
||||
_get_device_index,
|
||||
_handle_complex,
|
||||
_reorder_tensors_as,
|
||||
_take_tensors,
|
||||
_unflatten_dense_tensors,
|
||||
)
|
||||
from torch.cuda import nccl
|
||||
|
||||
|
||||
def broadcast(tensor, devices=None, *, out=None):
|
||||
r"""Broadcasts a tensor to specified GPU devices.
|
||||
|
||||
Args:
|
||||
tensor (Tensor): tensor to broadcast. Can be on CPU or GPU.
|
||||
devices (Iterable[torch.device, str or int], optional): an iterable of
|
||||
GPU devices, among which to broadcast.
|
||||
out (Sequence[Tensor], optional, keyword-only): the GPU tensors to
|
||||
store output results.
|
||||
|
||||
.. note::
|
||||
Exactly one of :attr:`devices` and :attr:`out` must be specified.
|
||||
|
||||
Returns:
|
||||
- If :attr:`devices` is specified,
|
||||
a tuple containing copies of :attr:`tensor`, placed on
|
||||
:attr:`devices`.
|
||||
- If :attr:`out` is specified,
|
||||
a tuple containing :attr:`out` tensors, each containing a copy of
|
||||
:attr:`tensor`.
|
||||
"""
|
||||
tensor = _handle_complex(tensor)
|
||||
if not ((devices is None) ^ (out is None)):
|
||||
raise RuntimeError(
|
||||
f"Exactly one of 'devices' and 'out' must be specified, but got devices={devices} and out={out}"
|
||||
)
|
||||
if devices is not None:
|
||||
devices = [_get_device_index(d) for d in devices]
|
||||
return torch._C._broadcast(tensor, devices)
|
||||
else:
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
return torch._C._broadcast_out(tensor, out)
|
||||
|
||||
|
||||
def broadcast_coalesced(tensors, devices, buffer_size=10485760):
|
||||
"""Broadcast a sequence of tensors to the specified GPUs.
|
||||
|
||||
Small tensors are first coalesced into a buffer to reduce the number of synchronizations.
|
||||
|
||||
Args:
|
||||
tensors (sequence): tensors to broadcast. Must be on the same device,
|
||||
either CPU or GPU.
|
||||
devices (Iterable[torch.device, str or int]): an iterable of GPU
|
||||
devices, among which to broadcast.
|
||||
buffer_size (int): maximum size of the buffer used for coalescing
|
||||
|
||||
Returns:
|
||||
A tuple containing copies of :attr:`tensor`, placed on :attr:`devices`.
|
||||
"""
|
||||
devices = [_get_device_index(d) for d in devices]
|
||||
tensors = [_handle_complex(t) for t in tensors]
|
||||
return torch._C._broadcast_coalesced(tensors, devices, buffer_size)
|
||||
|
||||
|
||||
def reduce_add(inputs, destination=None):
|
||||
"""Sum tensors from multiple GPUs.
|
||||
|
||||
All inputs should have matching shapes, dtype, and layout. The output tensor
|
||||
will be of the same shape, dtype, and layout.
|
||||
|
||||
Args:
|
||||
inputs (Iterable[Tensor]): an iterable of tensors to add.
|
||||
destination (int, optional): a device on which the output will be
|
||||
placed (default: current device).
|
||||
|
||||
Returns:
|
||||
A tensor containing an elementwise sum of all inputs, placed on the
|
||||
:attr:`destination` device.
|
||||
"""
|
||||
destination = _get_device_index(destination, optional=True)
|
||||
input_size = inputs[0].size()
|
||||
root_index = None # index of input tensor that already is on the correct device
|
||||
for i, inp in enumerate(inputs):
|
||||
if inp.device.type == "cpu":
|
||||
raise AssertionError(
|
||||
f"reduce_add expects all inputs to be on GPUs, but input {i} is on CPU"
|
||||
)
|
||||
if inp.get_device() == destination:
|
||||
root_index = i
|
||||
if inp.size() != input_size:
|
||||
got = "x".join(str(x) for x in inp.size())
|
||||
expected = "x".join(str(x) for x in input_size)
|
||||
raise ValueError(
|
||||
f"input {i} has invalid size: got {got}, but expected {expected}"
|
||||
)
|
||||
if root_index is None:
|
||||
raise RuntimeError(
|
||||
"reduce_add expects destination to be on the same GPU with one of the tensors"
|
||||
)
|
||||
|
||||
if len(inputs) == 1:
|
||||
return inputs[0]
|
||||
|
||||
if nccl.is_available(inputs):
|
||||
result = torch.empty_like(inputs[root_index])
|
||||
nccl.reduce(inputs, output=result, root=root_index)
|
||||
else:
|
||||
destination_device = torch.device(inputs[root_index].device.type, destination)
|
||||
nonroot = [t for i, t in enumerate(inputs) if i != root_index]
|
||||
# make a new tensor w/o clone
|
||||
result = inputs[root_index] + nonroot[0].to(
|
||||
device=destination_device, non_blocking=True
|
||||
)
|
||||
for other in nonroot[1:]:
|
||||
result.add_(other.to(device=destination_device, non_blocking=True))
|
||||
return result
|
||||
|
||||
|
||||
def reduce_add_coalesced(inputs, destination=None, buffer_size=10485760):
|
||||
"""Sum tensors from multiple GPUs.
|
||||
|
||||
Small tensors are first coalesced into a buffer to reduce the number
|
||||
of synchronizations.
|
||||
|
||||
Args:
|
||||
inputs (Iterable[Iterable[Tensor]]): iterable of iterables that
|
||||
contain tensors from a single device.
|
||||
destination (int, optional): a device on which the output will be
|
||||
placed (default: current device).
|
||||
buffer_size (int): maximum size of the buffer used for coalescing
|
||||
|
||||
Returns:
|
||||
A tuple of tensors containing an elementwise sum of each group of
|
||||
inputs, placed on the ``destination`` device.
|
||||
"""
|
||||
# TODO: When `len(inputs) == 1` and all inputs are on `destination`, just
|
||||
# return `inputs`.
|
||||
dense_tensors: list[list] = [[] for _ in inputs] # shape (num_gpus, num_tensors)
|
||||
output = []
|
||||
ref_order = []
|
||||
# process sparse ones first since they may have different sizes on different gpus
|
||||
for tensor_at_gpus in zip(*inputs, strict=True):
|
||||
if all(t.is_sparse for t in tensor_at_gpus):
|
||||
result = reduce_add(tensor_at_gpus, destination) # this will be sparse too
|
||||
output.append(result)
|
||||
ref_order.append(tensor_at_gpus[0])
|
||||
else:
|
||||
for coll, t in zip(dense_tensors, tensor_at_gpus, strict=True):
|
||||
coll.append(t.to_dense() if t.is_sparse else t)
|
||||
ref_order.append(dense_tensors[0][-1])
|
||||
itrs = [_take_tensors(tensors, buffer_size) for tensors in dense_tensors]
|
||||
# now the dense ones, which have consistent sizes
|
||||
for chunks in zip(*itrs, strict=True):
|
||||
flat_tensors = [
|
||||
_flatten_dense_tensors(chunk) for chunk in chunks
|
||||
] # (num_gpus,)
|
||||
flat_result = reduce_add(flat_tensors, destination)
|
||||
for t in _unflatten_dense_tensors(flat_result, chunks[0]):
|
||||
# The unflattened tensors do not share storage, and we don't expose
|
||||
# base flat tensor anyways, so give them different version counters.
|
||||
# See NOTE [ Version Counter in comm.*_coalesced ]
|
||||
output.append(t.data)
|
||||
return tuple(_reorder_tensors_as(output, ref_order))
|
||||
|
||||
|
||||
def scatter(tensor, devices=None, chunk_sizes=None, dim=0, streams=None, *, out=None):
|
||||
"""Scatters tensor across multiple GPUs.
|
||||
|
||||
Args:
|
||||
tensor (Tensor): tensor to scatter. Can be on CPU or GPU.
|
||||
devices (Iterable[torch.device, str or int], optional): an iterable of
|
||||
GPU devices, among which to scatter.
|
||||
chunk_sizes (Iterable[int], optional): sizes of chunks to be placed on
|
||||
each device. It should match :attr:`devices` in length and sums to
|
||||
``tensor.size(dim)``. If not specified, :attr:`tensor` will be divided
|
||||
into equal chunks.
|
||||
dim (int, optional): A dimension along which to chunk :attr:`tensor`.
|
||||
Default: ``0``.
|
||||
streams (Iterable[torch.cuda.Stream], optional): an iterable of Streams, among
|
||||
which to execute the scatter. If not specified, the default stream will
|
||||
be utilized.
|
||||
out (Sequence[Tensor], optional, keyword-only): the GPU tensors to
|
||||
store output results. Sizes of these tensors must match that of
|
||||
:attr:`tensor`, except for :attr:`dim`, where the total size must
|
||||
sum to ``tensor.size(dim)``.
|
||||
|
||||
.. note::
|
||||
Exactly one of :attr:`devices` and :attr:`out` must be specified. When
|
||||
:attr:`out` is specified, :attr:`chunk_sizes` must not be specified and
|
||||
will be inferred from sizes of :attr:`out`.
|
||||
|
||||
Returns:
|
||||
- If :attr:`devices` is specified,
|
||||
a tuple containing chunks of :attr:`tensor`, placed on
|
||||
:attr:`devices`.
|
||||
- If :attr:`out` is specified,
|
||||
a tuple containing :attr:`out` tensors, each containing a chunk of
|
||||
:attr:`tensor`.
|
||||
"""
|
||||
tensor = _handle_complex(tensor)
|
||||
if out is None:
|
||||
# pyrefly: ignore [not-iterable]
|
||||
devices = [_get_device_index(d) for d in devices]
|
||||
return tuple(torch._C._scatter(tensor, devices, chunk_sizes, dim, streams))
|
||||
else:
|
||||
if devices is not None:
|
||||
raise RuntimeError(
|
||||
f"'devices' must not be specified when 'out' is specified, but got devices={devices}"
|
||||
)
|
||||
if chunk_sizes is not None:
|
||||
raise RuntimeError(
|
||||
f"'chunk_sizes' must not be specified when 'out' is specified, but got chunk_sizes={chunk_sizes}"
|
||||
)
|
||||
return tuple(torch._C._scatter_out(tensor, out, dim, streams))
|
||||
|
||||
|
||||
def gather(tensors, dim=0, destination=None, *, out=None):
|
||||
r"""Gathers tensors from multiple GPU devices.
|
||||
|
||||
Args:
|
||||
tensors (Iterable[Tensor]): an iterable of tensors to gather.
|
||||
Tensor sizes in all dimensions other than :attr:`dim` have to match.
|
||||
dim (int, optional): a dimension along which the tensors will be
|
||||
concatenated. Default: ``0``.
|
||||
destination (torch.device, str, or int, optional): the output device.
|
||||
Can be CPU or CUDA. Default: the current CUDA device.
|
||||
out (Tensor, optional, keyword-only): the tensor to store gather result.
|
||||
Its sizes must match those of :attr:`tensors`, except for :attr:`dim`,
|
||||
where the size must equal ``sum(tensor.size(dim) for tensor in tensors)``.
|
||||
Can be on CPU or CUDA.
|
||||
|
||||
.. note::
|
||||
:attr:`destination` must not be specified when :attr:`out` is specified.
|
||||
|
||||
Returns:
|
||||
- If :attr:`destination` is specified,
|
||||
a tensor located on :attr:`destination` device, that is a result of
|
||||
concatenating :attr:`tensors` along :attr:`dim`.
|
||||
- If :attr:`out` is specified,
|
||||
the :attr:`out` tensor, now containing results of concatenating
|
||||
:attr:`tensors` along :attr:`dim`.
|
||||
"""
|
||||
tensors = [_handle_complex(t) for t in tensors]
|
||||
if out is None:
|
||||
if destination == -1:
|
||||
warnings.warn(
|
||||
"Using -1 to represent CPU tensor is deprecated. Please use a "
|
||||
'device object or string instead, e.g., "cpu".',
|
||||
FutureWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
destination = _get_device_index(destination, allow_cpu=True, optional=True)
|
||||
return torch._C._gather(tensors, dim, destination)
|
||||
else:
|
||||
if destination is not None:
|
||||
raise RuntimeError(
|
||||
f"'destination' must not be specified when 'out' is specified, but got destination={destination}"
|
||||
)
|
||||
return torch._C._gather_out(tensors, out, dim)
|
||||
@@ -0,0 +1,289 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import operator
|
||||
import warnings
|
||||
from collections.abc import Sequence
|
||||
from itertools import chain
|
||||
from typing import Any, Generic, TypeVar
|
||||
|
||||
import torch
|
||||
from torch._utils import (
|
||||
_get_all_device_indices,
|
||||
_get_available_device_type,
|
||||
_get_device_index,
|
||||
_get_devices_properties,
|
||||
)
|
||||
from torch.nn.modules import Module
|
||||
from torch.nn.parallel.parallel_apply import parallel_apply
|
||||
from torch.nn.parallel.replicate import replicate
|
||||
from torch.nn.parallel.scatter_gather import gather, scatter_kwargs
|
||||
|
||||
|
||||
__all__ = ["DataParallel", "data_parallel"]
|
||||
|
||||
|
||||
def _check_balance(device_ids: Sequence[int | torch.device]) -> None:
|
||||
imbalance_warn = """
|
||||
There is an imbalance between your GPUs. You may want to exclude GPU {} which
|
||||
has less than 75% of the memory or cores of GPU {}. You can do so by setting
|
||||
the device_ids argument to DataParallel, or by setting the CUDA_VISIBLE_DEVICES
|
||||
environment variable."""
|
||||
device_ids = [_get_device_index(x, True) for x in device_ids]
|
||||
dev_props = _get_devices_properties(device_ids)
|
||||
|
||||
def warn_imbalance(get_prop) -> bool:
|
||||
values = [get_prop(props) for props in dev_props]
|
||||
min_pos, min_val = min(enumerate(values), key=operator.itemgetter(1))
|
||||
max_pos, max_val = max(enumerate(values), key=operator.itemgetter(1))
|
||||
if min_val / max_val < 0.75:
|
||||
warnings.warn(
|
||||
imbalance_warn.format(device_ids[min_pos], device_ids[max_pos]),
|
||||
stacklevel=2,
|
||||
)
|
||||
return True
|
||||
return False
|
||||
|
||||
if warn_imbalance(lambda props: props.total_memory):
|
||||
return
|
||||
if warn_imbalance(lambda props: props.multi_processor_count):
|
||||
return
|
||||
|
||||
|
||||
T = TypeVar("T", bound=Module)
|
||||
|
||||
|
||||
class DataParallel(Module, Generic[T]):
|
||||
r"""Implements data parallelism at the module level.
|
||||
|
||||
This container parallelizes the application of the given :attr:`module` by
|
||||
splitting the input across the specified devices by chunking in the batch
|
||||
dimension (other objects will be copied once per device). In the forward
|
||||
pass, the module is replicated on each device, and each replica handles a
|
||||
portion of the input. During the backwards pass, gradients from each replica
|
||||
are summed into the original module.
|
||||
|
||||
The batch size should be larger than the number of GPUs used.
|
||||
|
||||
.. warning::
|
||||
It is recommended to use :class:`~torch.nn.parallel.DistributedDataParallel`,
|
||||
instead of this class, to do multi-GPU training, even if there is only a single
|
||||
node. See: :ref:`cuda-nn-ddp-instead` and :ref:`ddp`.
|
||||
|
||||
Arbitrary positional and keyword inputs are allowed to be passed into
|
||||
DataParallel but some types are specially handled. tensors will be
|
||||
**scattered** on dim specified (default 0). tuple, list and dict types will
|
||||
be shallow copied. The other types will be shared among different threads
|
||||
and can be corrupted if written to in the model's forward pass.
|
||||
|
||||
The parallelized :attr:`module` must have its parameters and buffers on
|
||||
``device_ids[0]`` before running this :class:`~torch.nn.DataParallel`
|
||||
module.
|
||||
|
||||
.. warning::
|
||||
In each forward, :attr:`module` is **replicated** on each device, so any
|
||||
updates to the running module in ``forward`` will be lost. For example,
|
||||
if :attr:`module` has a counter attribute that is incremented in each
|
||||
``forward``, it will always stay at the initial value because the update
|
||||
is done on the replicas which are destroyed after ``forward``. However,
|
||||
:class:`~torch.nn.DataParallel` guarantees that the replica on
|
||||
``device[0]`` will have its parameters and buffers sharing storage with
|
||||
the base parallelized :attr:`module`. So **in-place** updates to the
|
||||
parameters or buffers on ``device[0]`` will be recorded. E.g.,
|
||||
:class:`~torch.nn.BatchNorm2d` and :func:`~torch.nn.utils.spectral_norm`
|
||||
rely on this behavior to update the buffers.
|
||||
|
||||
.. warning::
|
||||
Forward and backward hooks defined on :attr:`module` and its submodules
|
||||
will be invoked ``len(device_ids)`` times, each with inputs located on
|
||||
a particular device. Particularly, the hooks are only guaranteed to be
|
||||
executed in correct order with respect to operations on corresponding
|
||||
devices. For example, it is not guaranteed that hooks set via
|
||||
:meth:`~torch.nn.Module.register_forward_pre_hook` be executed before
|
||||
`all` ``len(device_ids)`` :meth:`~torch.nn.Module.forward` calls, but
|
||||
that each such hook be executed before the corresponding
|
||||
:meth:`~torch.nn.Module.forward` call of that device.
|
||||
|
||||
.. warning::
|
||||
When :attr:`module` returns a scalar (i.e., 0-dimensional tensor) in
|
||||
:func:`forward`, this wrapper will return a vector of length equal to
|
||||
number of devices used in data parallelism, containing the result from
|
||||
each device.
|
||||
|
||||
.. note::
|
||||
There is a subtlety in using the
|
||||
``pack sequence -> recurrent network -> unpack sequence`` pattern in a
|
||||
:class:`~torch.nn.Module` wrapped in :class:`~torch.nn.DataParallel`.
|
||||
See :ref:`pack-rnn-unpack-with-data-parallelism` section in FAQ for
|
||||
details.
|
||||
|
||||
|
||||
Args:
|
||||
module (Module): module to be parallelized
|
||||
device_ids (list of int or torch.device): CUDA devices (default: all devices)
|
||||
output_device (int or torch.device): device location of output (default: device_ids[0])
|
||||
|
||||
Attributes:
|
||||
module (Module): the module to be parallelized
|
||||
|
||||
Example::
|
||||
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])
|
||||
>>> output = net(input_var) # input_var can be on any device, including CPU
|
||||
"""
|
||||
|
||||
# TODO: update notes/cuda.rst when this class handles 8+ GPUs well
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
module: T,
|
||||
device_ids: Sequence[int | torch.device] | None = None,
|
||||
output_device: int | torch.device | None = None,
|
||||
dim: int = 0,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
torch._C._log_api_usage_once("torch.nn.parallel.DataParallel")
|
||||
device_type = _get_available_device_type()
|
||||
if device_type is None or device_type == "mps":
|
||||
self.module = module
|
||||
self.device_ids = []
|
||||
return
|
||||
|
||||
if device_ids is None:
|
||||
device_ids = _get_all_device_indices()
|
||||
|
||||
if device_ids is None:
|
||||
raise RuntimeError("no available devices were found")
|
||||
|
||||
if output_device is None:
|
||||
output_device = device_ids[0]
|
||||
|
||||
self.dim = dim
|
||||
self.module = module
|
||||
self.device_ids = [_get_device_index(x, True) for x in device_ids]
|
||||
self.output_device = _get_device_index(output_device, True)
|
||||
self.src_device_obj = torch.device(device_type, self.device_ids[0])
|
||||
|
||||
if device_type == "cuda":
|
||||
_check_balance(self.device_ids)
|
||||
|
||||
if len(self.device_ids) == 1:
|
||||
self.module.to(self.src_device_obj)
|
||||
|
||||
def forward(self, *inputs: Any, **kwargs: Any) -> Any:
|
||||
with torch.autograd.profiler.record_function("DataParallel.forward"):
|
||||
if not self.device_ids:
|
||||
return self.module(*inputs, **kwargs)
|
||||
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
for t in chain(self.module.parameters(), self.module.buffers()):
|
||||
if t.device != self.src_device_obj:
|
||||
raise RuntimeError(
|
||||
"module must have its parameters and buffers "
|
||||
f"on device {self.src_device_obj} (device_ids[0]) but found one of "
|
||||
f"them on device: {t.device}"
|
||||
)
|
||||
|
||||
inputs, module_kwargs = self.scatter(inputs, kwargs, self.device_ids)
|
||||
# for forward function without any inputs, empty list and dict will be created
|
||||
# so the module can be executed on one device which is the first one in device_ids
|
||||
if not inputs and not module_kwargs:
|
||||
inputs = ((),)
|
||||
module_kwargs = ({},)
|
||||
|
||||
if len(self.device_ids) == 1:
|
||||
return self.module(*inputs[0], **module_kwargs[0])
|
||||
replicas = self.replicate(self.module, self.device_ids[: len(inputs)])
|
||||
outputs = self.parallel_apply(replicas, inputs, module_kwargs)
|
||||
return self.gather(outputs, self.output_device)
|
||||
|
||||
def replicate(self, module: T, device_ids: Sequence[int | torch.device]) -> list[T]:
|
||||
return replicate(module, device_ids, not torch.is_grad_enabled())
|
||||
|
||||
def scatter(
|
||||
self,
|
||||
inputs: tuple[Any, ...],
|
||||
kwargs: dict[str, Any] | None,
|
||||
device_ids: Sequence[int | torch.device],
|
||||
) -> Any:
|
||||
return scatter_kwargs(inputs, kwargs, device_ids, dim=self.dim)
|
||||
|
||||
def parallel_apply(
|
||||
self, replicas: Sequence[T], inputs: Sequence[Any], kwargs: Any
|
||||
) -> list[Any]:
|
||||
return parallel_apply(
|
||||
replicas, inputs, kwargs, self.device_ids[: len(replicas)]
|
||||
)
|
||||
|
||||
def gather(self, outputs: Any, output_device: int | torch.device) -> Any:
|
||||
return gather(outputs, output_device, dim=self.dim)
|
||||
|
||||
|
||||
def data_parallel(
|
||||
module: Module,
|
||||
inputs: Any,
|
||||
device_ids: Sequence[int | torch.device] | None = None,
|
||||
output_device: int | torch.device | None = None,
|
||||
dim: int = 0,
|
||||
module_kwargs: Any | None = None,
|
||||
) -> torch.Tensor:
|
||||
r"""Evaluate module(input) in parallel across the GPUs given in device_ids.
|
||||
|
||||
This is the functional version of the DataParallel module.
|
||||
|
||||
Args:
|
||||
module (Module): the module to evaluate in parallel
|
||||
inputs (Tensor): inputs to the module
|
||||
device_ids (list of int or torch.device): GPU ids on which to replicate module
|
||||
output_device (list of int or torch.device): GPU location of the output Use -1 to indicate the CPU.
|
||||
(default: device_ids[0])
|
||||
Returns:
|
||||
a Tensor containing the result of module(input) located on
|
||||
output_device
|
||||
"""
|
||||
if not isinstance(inputs, tuple):
|
||||
inputs = (inputs,) if inputs is not None else ()
|
||||
|
||||
device_type = _get_available_device_type()
|
||||
|
||||
if device_type is None:
|
||||
raise RuntimeError("device type could not be determined")
|
||||
|
||||
if device_ids is None:
|
||||
device_ids = _get_all_device_indices()
|
||||
|
||||
if device_ids is None:
|
||||
raise RuntimeError("no available devices were found")
|
||||
|
||||
if output_device is None:
|
||||
output_device = device_ids[0]
|
||||
|
||||
device_ids = [_get_device_index(x, True) for x in device_ids]
|
||||
output_device = _get_device_index(output_device, True)
|
||||
# pyrefly: ignore [bad-argument-type, no-matching-overload]
|
||||
src_device_obj = torch.device(device_type, device_ids[0])
|
||||
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
for t in chain(module.parameters(), module.buffers()):
|
||||
if t.device != src_device_obj:
|
||||
raise RuntimeError(
|
||||
"module must have its parameters and buffers "
|
||||
f"on device {src_device_obj} (device_ids[0]) but found one of "
|
||||
f"them on device: {t.device}"
|
||||
)
|
||||
|
||||
inputs, module_kwargs = scatter_kwargs(inputs, module_kwargs, device_ids, dim)
|
||||
# for module without any inputs, empty list and dict will be created
|
||||
# so the module can be executed on one device which is the first one in device_ids
|
||||
if not inputs and not module_kwargs:
|
||||
inputs = ((),)
|
||||
module_kwargs = ({},)
|
||||
|
||||
if module_kwargs is None:
|
||||
raise AssertionError("module_kwargs should not be None after scatter_kwargs")
|
||||
|
||||
if len(device_ids) == 1:
|
||||
return module(*inputs[0], **module_kwargs[0])
|
||||
used_device_ids = device_ids[: len(inputs)]
|
||||
replicas = replicate(module, used_device_ids)
|
||||
outputs = parallel_apply(replicas, inputs, module_kwargs, used_device_ids)
|
||||
return gather(outputs, output_device, dim)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,146 @@
|
||||
import threading
|
||||
from collections.abc import Sequence
|
||||
from typing import Any, cast
|
||||
|
||||
import torch
|
||||
from torch._utils import ExceptionWrapper
|
||||
from torch.cuda._utils import _get_device_index
|
||||
from torch.nn.modules import Module
|
||||
|
||||
|
||||
__all__ = ["get_a_var", "parallel_apply"]
|
||||
|
||||
|
||||
def get_a_var(
|
||||
obj: torch.Tensor | list[Any] | tuple[Any, ...] | dict[Any, Any],
|
||||
) -> torch.Tensor | None:
|
||||
if isinstance(obj, torch.Tensor):
|
||||
return obj
|
||||
|
||||
if isinstance(obj, (list, tuple)):
|
||||
for result in map(get_a_var, obj):
|
||||
if isinstance(result, torch.Tensor):
|
||||
return result
|
||||
if isinstance(obj, dict):
|
||||
for result in map(get_a_var, obj.items()):
|
||||
if isinstance(result, torch.Tensor):
|
||||
return result
|
||||
return None
|
||||
|
||||
|
||||
def parallel_apply(
|
||||
modules: Sequence[Module],
|
||||
inputs: Sequence[Any],
|
||||
kwargs_tup: Sequence[dict[str, Any]] | None = None,
|
||||
devices: Sequence[int | torch.device | None] | None = None,
|
||||
) -> list[Any]:
|
||||
r"""Apply each `module` in :attr:`modules` in parallel on each of :attr:`devices`.
|
||||
|
||||
Args:
|
||||
modules (Module): modules to be parallelized
|
||||
inputs (tensor): inputs to the modules
|
||||
devices (list of int or torch.device): CUDA devices
|
||||
|
||||
:attr:`modules`, :attr:`inputs`, :attr:`kwargs_tup` (if given), and
|
||||
:attr:`devices` (if given) should all have same length. Moreover, each
|
||||
element of :attr:`inputs` can either be a single object as the only argument
|
||||
to a module, or a collection of positional arguments.
|
||||
"""
|
||||
if len(modules) != len(inputs):
|
||||
raise AssertionError(
|
||||
f"The number of modules {len(modules)} is not equal to "
|
||||
f"the number of inputs {len(inputs)}"
|
||||
)
|
||||
if kwargs_tup is not None:
|
||||
if len(modules) != len(kwargs_tup):
|
||||
raise AssertionError(
|
||||
f"The number of modules {len(modules)} is not equal to "
|
||||
f"the number of kwargs_tup {len(kwargs_tup)}"
|
||||
)
|
||||
else:
|
||||
kwargs_tup = (cast(dict[str, Any], {}),) * len(modules)
|
||||
if devices is not None:
|
||||
if len(modules) != len(devices):
|
||||
raise AssertionError(
|
||||
f"The number of modules {len(modules)} is not equal to "
|
||||
f"the number of devices {len(devices)}"
|
||||
)
|
||||
else:
|
||||
devices = [None] * len(modules)
|
||||
devices = [_get_device_index(x, True) for x in devices]
|
||||
streams = [torch.accelerator.current_stream(x) for x in devices]
|
||||
if not torch.accelerator.is_available():
|
||||
raise AssertionError("No available accelerator found.")
|
||||
device_type = torch.accelerator.current_accelerator().type # type: ignore[union-attr]
|
||||
lock = threading.Lock()
|
||||
results = {}
|
||||
grad_enabled, autocast_enabled = (
|
||||
torch.is_grad_enabled(),
|
||||
torch.is_autocast_enabled(),
|
||||
)
|
||||
|
||||
def _worker(
|
||||
i: int,
|
||||
module: Module,
|
||||
input: Any,
|
||||
kwargs: dict[str, Any],
|
||||
device: int | torch.device | None = None,
|
||||
stream: torch.Stream | None = None,
|
||||
) -> None:
|
||||
torch.set_grad_enabled(grad_enabled)
|
||||
if device is None:
|
||||
t = get_a_var(input)
|
||||
if t is None:
|
||||
with lock:
|
||||
results[i] = ExceptionWrapper(
|
||||
where=f"in replica {i}, no device was provided and no tensor input was found; "
|
||||
"device cannot be resolved"
|
||||
)
|
||||
return
|
||||
device = t.get_device()
|
||||
if isinstance(device, torch.device):
|
||||
device = device.index
|
||||
if stream is None:
|
||||
stream = torch.accelerator.current_stream(device)
|
||||
try:
|
||||
with (
|
||||
torch.accelerator.device_index(device),
|
||||
stream,
|
||||
torch.amp.autocast(device_type, enabled=autocast_enabled),
|
||||
):
|
||||
# this also avoids accidental slicing of `input` if it is a Tensor
|
||||
if not isinstance(input, (list, tuple)):
|
||||
input = (input,)
|
||||
output = module(*input, **kwargs)
|
||||
with lock:
|
||||
results[i] = output
|
||||
except Exception:
|
||||
with lock:
|
||||
results[i] = ExceptionWrapper(
|
||||
where=f"in replica {i} on device {device}"
|
||||
)
|
||||
|
||||
if len(modules) > 1:
|
||||
threads = [
|
||||
threading.Thread(
|
||||
target=_worker, args=(i, module, input, kwargs, device, stream)
|
||||
)
|
||||
for i, (module, input, kwargs, device, stream) in enumerate(
|
||||
zip(modules, inputs, kwargs_tup, devices, streams, strict=True)
|
||||
)
|
||||
]
|
||||
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
else:
|
||||
_worker(0, modules[0], inputs[0], kwargs_tup[0], devices[0], streams[0])
|
||||
|
||||
outputs = []
|
||||
for i in range(len(inputs)):
|
||||
output = results[i]
|
||||
if isinstance(output, ExceptionWrapper):
|
||||
output.reraise()
|
||||
outputs.append(output)
|
||||
return outputs
|
||||
@@ -0,0 +1,214 @@
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Iterator, Sequence
|
||||
from typing import cast, TYPE_CHECKING, TypeVar
|
||||
from typing_extensions import TypeIs
|
||||
|
||||
import torch
|
||||
from torch._utils import _get_device_index
|
||||
from torch.nn.modules import Module
|
||||
from torch.nn.parallel import comm
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch._C import ScriptMethod
|
||||
from torch.jit import ScriptModule
|
||||
from torch.jit._state import EnabledProxy
|
||||
|
||||
|
||||
__all__ = ["replicate"]
|
||||
|
||||
|
||||
def _is_script_module(module: Module) -> TypeIs["ScriptModule"]:
|
||||
import torch.jit
|
||||
|
||||
return isinstance(module, torch.jit.ScriptModule)
|
||||
|
||||
|
||||
def _is_script_method(module: object) -> TypeIs["ScriptMethod"]:
|
||||
import torch.jit
|
||||
|
||||
return isinstance(module, torch._C.ScriptMethod)
|
||||
|
||||
|
||||
def _init_script_module() -> "ScriptModule":
|
||||
import torch.jit
|
||||
|
||||
return torch.jit.ScriptModule()
|
||||
|
||||
|
||||
def _is_jit_enabled() -> "EnabledProxy":
|
||||
import torch.jit._state
|
||||
|
||||
return torch.jit._state._enabled
|
||||
|
||||
|
||||
# Check if we can safely replicate the module.
|
||||
# there are two types of module:
|
||||
# 1. python modules
|
||||
# 2. ScriptModule
|
||||
#
|
||||
# currently a module cannot be replicated properly if the descendants of
|
||||
# any ScriptModule contains python module (type 1 above)
|
||||
def _replicatable_module(module: Module, memo: set[Module] | None = None) -> bool:
|
||||
# module.modules() contains module itself as the first element
|
||||
def descendant_modules(module: Module) -> Iterator[Module]:
|
||||
gen = module.modules()
|
||||
next(gen)
|
||||
return gen
|
||||
|
||||
if not _is_jit_enabled():
|
||||
return True
|
||||
if memo is None:
|
||||
memo = set()
|
||||
|
||||
# memoize visited modules
|
||||
memo.add(module)
|
||||
if _is_script_module(module):
|
||||
memo.update(descendant_modules(module))
|
||||
return all(
|
||||
_is_script_module(descendant) for descendant in descendant_modules(module)
|
||||
)
|
||||
|
||||
for child in module.children():
|
||||
# since any unreplicatable module will cause the check to return
|
||||
# False early, visited modules here can be safely ignored.
|
||||
if child in memo:
|
||||
continue
|
||||
if not _replicatable_module(child, memo):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def _broadcast_coalesced_reshape(
|
||||
tensors: Sequence[torch.Tensor],
|
||||
devices: Sequence[int | torch.device],
|
||||
detach: bool = False,
|
||||
) -> list[list[torch.Tensor]]:
|
||||
from torch.nn.parallel._functions import Broadcast
|
||||
|
||||
if len(tensors) == 0:
|
||||
return []
|
||||
|
||||
if detach:
|
||||
complex_mask = [
|
||||
not isinstance(t, torch.nn.UninitializedParameter) and t.is_complex()
|
||||
for t in tensors
|
||||
]
|
||||
|
||||
outputs = comm.broadcast_coalesced(tensors, devices)
|
||||
|
||||
for device_outputs in outputs:
|
||||
for i, is_complex in enumerate(complex_mask):
|
||||
if is_complex:
|
||||
device_outputs[i] = torch.view_as_complex(device_outputs[i])
|
||||
|
||||
return outputs
|
||||
else:
|
||||
tensor_copies = Broadcast.apply(devices, *tensors)
|
||||
return [
|
||||
list(tensor_copies[i : i + len(tensors)])
|
||||
for i in range(0, len(tensor_copies), len(tensors))
|
||||
]
|
||||
|
||||
|
||||
T = TypeVar("T", bound=Module)
|
||||
|
||||
|
||||
def replicate(
|
||||
network: T,
|
||||
devices: Sequence[int | torch.device],
|
||||
detach: bool = False,
|
||||
) -> list[T]:
|
||||
if not _replicatable_module(network):
|
||||
raise RuntimeError(
|
||||
"Cannot replicate network where python modules are children of ScriptModule"
|
||||
)
|
||||
|
||||
if not devices:
|
||||
return []
|
||||
|
||||
devices = [_get_device_index(x, True) for x in devices]
|
||||
num_replicas = len(devices)
|
||||
|
||||
params = list(network.parameters())
|
||||
param_indices = {param: idx for idx, param in enumerate(params)}
|
||||
param_copies = _broadcast_coalesced_reshape(params, devices, detach)
|
||||
|
||||
buffers = list(network.buffers())
|
||||
buffers_rg: list[torch.Tensor] = []
|
||||
buffers_not_rg: list[torch.Tensor] = []
|
||||
for buf in buffers:
|
||||
if buf.requires_grad and not detach:
|
||||
buffers_rg.append(buf)
|
||||
else:
|
||||
buffers_not_rg.append(buf)
|
||||
|
||||
buffer_indices_rg = {buf: idx for idx, buf in enumerate(buffers_rg)}
|
||||
buffer_indices_not_rg = {buf: idx for idx, buf in enumerate(buffers_not_rg)}
|
||||
|
||||
buffer_copies_rg = _broadcast_coalesced_reshape(buffers_rg, devices, detach=detach)
|
||||
buffer_copies_not_rg = _broadcast_coalesced_reshape(
|
||||
buffers_not_rg, devices, detach=True
|
||||
)
|
||||
|
||||
modules = list(network.modules())
|
||||
module_copies: list[list[Module]] = [[] for _ in devices]
|
||||
module_indices: dict[Module, int] = {}
|
||||
|
||||
for i, module in enumerate(modules):
|
||||
module_indices[module] = i
|
||||
for j in range(num_replicas):
|
||||
replica = module._replicate_for_data_parallel()
|
||||
# This is a temporary fix for DDP. DDP needs to access the
|
||||
# replicated model parameters. It used to do so through
|
||||
# `mode.parameters()`. The fix added in #33907 for DP stops the
|
||||
# `parameters()` API from exposing the replicated parameters.
|
||||
# Hence, we add a `_former_parameters` dict here to support DDP.
|
||||
replica._former_parameters = OrderedDict()
|
||||
|
||||
module_copies[j].append(replica)
|
||||
|
||||
for i, module in enumerate(modules):
|
||||
for key, child in module._modules.items():
|
||||
if child is None:
|
||||
for j in range(num_replicas):
|
||||
replica = module_copies[j][i]
|
||||
replica._modules[key] = None
|
||||
else:
|
||||
module_idx = module_indices[child]
|
||||
for j in range(num_replicas):
|
||||
replica = module_copies[j][i]
|
||||
setattr(replica, key, module_copies[j][module_idx])
|
||||
for key, param in module._parameters.items():
|
||||
if param is None:
|
||||
for j in range(num_replicas):
|
||||
replica = module_copies[j][i]
|
||||
replica._parameters[key] = None
|
||||
else:
|
||||
param_idx = param_indices[param]
|
||||
for j in range(num_replicas):
|
||||
replica = module_copies[j][i]
|
||||
param_copy = param_copies[j][param_idx]
|
||||
# parameters in replicas are no longer leaves,
|
||||
# so setattr them as non-parameter attributes
|
||||
setattr(replica, key, param_copy)
|
||||
# expose the parameter for DDP
|
||||
replica._former_parameters[key] = param_copy # type: ignore[operator, index]
|
||||
for key, buf in module._buffers.items(): # type: ignore[assignment]
|
||||
if buf is None:
|
||||
for j in range(num_replicas):
|
||||
replica = module_copies[j][i]
|
||||
replica._buffers[key] = None
|
||||
else:
|
||||
if buf.requires_grad and not detach:
|
||||
buffer_copies = buffer_copies_rg
|
||||
buffer_idx = buffer_indices_rg[buf]
|
||||
else:
|
||||
buffer_copies = buffer_copies_not_rg
|
||||
buffer_idx = buffer_indices_not_rg[buf]
|
||||
for j in range(num_replicas):
|
||||
replica = module_copies[j][i]
|
||||
setattr(replica, key, buffer_copies[j][buffer_idx])
|
||||
|
||||
return [cast(T, module_copies[j][0]) for j in range(num_replicas)]
|
||||
@@ -0,0 +1,150 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from collections.abc import Sequence
|
||||
from typing import Any, overload, TypeVar
|
||||
from typing_extensions import deprecated
|
||||
|
||||
import torch
|
||||
from torch.nn.parallel._functions import Gather, Scatter
|
||||
|
||||
|
||||
__all__ = ["scatter", "scatter_kwargs", "gather"]
|
||||
|
||||
|
||||
@deprecated(
|
||||
"`is_namedtuple` is deprecated, please use the python checks instead",
|
||||
category=FutureWarning,
|
||||
)
|
||||
def is_namedtuple(obj: Any) -> bool:
|
||||
# Check if type was created from collections.namedtuple or a typing.NamedTuple.
|
||||
return _is_namedtuple(obj)
|
||||
|
||||
|
||||
def _is_namedtuple(obj: Any) -> bool:
|
||||
# Check if type was created from collections.namedtuple or a typing.NamedTuple.
|
||||
return (
|
||||
isinstance(obj, tuple) and hasattr(obj, "_asdict") and hasattr(obj, "_fields")
|
||||
)
|
||||
|
||||
|
||||
T = TypeVar("T", dict, list, tuple)
|
||||
|
||||
|
||||
# For some reason, 'scatter' returns a tuple when given a single Tensor input but a list otherwise.
|
||||
@overload
|
||||
def scatter(
|
||||
inputs: torch.Tensor,
|
||||
target_gpus: Sequence[int | torch.device],
|
||||
dim: int = ...,
|
||||
) -> tuple[torch.Tensor, ...]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def scatter(
|
||||
inputs: T,
|
||||
target_gpus: Sequence[int | torch.device],
|
||||
dim: int = ...,
|
||||
) -> list[T]: ...
|
||||
|
||||
|
||||
def scatter(inputs, target_gpus, dim=0):
|
||||
r"""Slice tensors into approximately equal chunks and distributes them across given GPUs.
|
||||
|
||||
Duplicates references to objects that are not tensors.
|
||||
"""
|
||||
|
||||
def scatter_map(obj):
|
||||
if isinstance(obj, torch.Tensor):
|
||||
return Scatter.apply(target_gpus, None, dim, obj)
|
||||
if _is_namedtuple(obj):
|
||||
return [
|
||||
type(obj)(*args)
|
||||
# pyrefly: ignore [bad-argument-type, no-matching-overload]
|
||||
for args in zip(*map(scatter_map, obj), strict=False)
|
||||
]
|
||||
if isinstance(obj, tuple) and len(obj) > 0:
|
||||
# pyrefly: ignore [bad-argument-type, no-matching-overload]
|
||||
return list(zip(*map(scatter_map, obj), strict=False))
|
||||
if isinstance(obj, list) and len(obj) > 0:
|
||||
# pyrefly: ignore [bad-argument-type, no-matching-overload]
|
||||
return [list(i) for i in zip(*map(scatter_map, obj), strict=False)]
|
||||
if isinstance(obj, dict) and len(obj) > 0:
|
||||
return [
|
||||
type(obj)(i)
|
||||
# pyrefly: ignore [bad-argument-type, no-matching-overload]
|
||||
for i in zip(*map(scatter_map, obj.items()), strict=False)
|
||||
]
|
||||
return [obj for _ in target_gpus]
|
||||
|
||||
# After scatter_map is called, a scatter_map cell will exist. This cell
|
||||
# has a reference to the actual function scatter_map, which has references
|
||||
# to a closure that has a reference to the scatter_map cell (because the
|
||||
# fn is recursive). To avoid this reference cycle, we set the function to
|
||||
# None, clearing the cell
|
||||
try:
|
||||
res = scatter_map(inputs)
|
||||
finally:
|
||||
scatter_map = None # type: ignore[assignment]
|
||||
return res
|
||||
|
||||
|
||||
def scatter_kwargs(
|
||||
inputs: tuple[Any, ...],
|
||||
kwargs: dict[str, Any] | None,
|
||||
target_gpus: Sequence[int | torch.device],
|
||||
dim: int = 0,
|
||||
) -> tuple[tuple[Any, ...], tuple[dict[str, Any], ...]]:
|
||||
r"""Scatter with support for kwargs dictionary."""
|
||||
scattered_inputs = scatter(inputs, target_gpus, dim) if inputs else []
|
||||
scattered_kwargs = scatter(kwargs, target_gpus, dim) if kwargs else []
|
||||
if len(scattered_inputs) < len(scattered_kwargs):
|
||||
scattered_inputs.extend(
|
||||
() for _ in range(len(scattered_kwargs) - len(scattered_inputs))
|
||||
)
|
||||
elif len(scattered_kwargs) < len(inputs):
|
||||
scattered_kwargs.extend(
|
||||
{} for _ in range(len(scattered_inputs) - len(scattered_kwargs))
|
||||
)
|
||||
return tuple(scattered_inputs), tuple(scattered_kwargs)
|
||||
|
||||
|
||||
def gather(outputs: Any, target_device: int | torch.device, dim: int = 0) -> Any:
|
||||
r"""Gather tensors from different GPUs on a specified device.
|
||||
|
||||
This function is useful for gathering the results of a distributed computation.
|
||||
It takes a sequence of objects, one for each GPU, and returns a single object
|
||||
on the specified device.
|
||||
|
||||
Args:
|
||||
outputs (Any): A sequence of objects (potentially tensors) to gather.
|
||||
target_device (Union[int, torch.device]): The device to gather the tensors to.
|
||||
Use 'cpu' for CPU to avoid a deprecation warning.
|
||||
dim (int, optional): The dimension along which to gather. Default: 0.
|
||||
|
||||
Returns:
|
||||
Any: A gathered object (potentially tensor) on the specified device.
|
||||
"""
|
||||
|
||||
def gather_map(outputs):
|
||||
out = outputs[0]
|
||||
if isinstance(out, torch.Tensor):
|
||||
return Gather.apply(target_device, dim, *outputs)
|
||||
if out is None:
|
||||
return None
|
||||
if isinstance(out, dict):
|
||||
if not all(len(out) == len(d) for d in outputs):
|
||||
raise ValueError("All dicts must have the same number of keys")
|
||||
# pyrefly: ignore [not-callable]
|
||||
return type(out)((k, gather_map([d[k] for d in outputs])) for k in out)
|
||||
if _is_namedtuple(out):
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
return type(out)._make(map(gather_map, zip(*outputs, strict=True)))
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
return type(out)(map(gather_map, zip(*outputs, strict=True)))
|
||||
|
||||
# Recursive function calls like this create reference cycles.
|
||||
# Setting the function to None clears the refcycle.
|
||||
try:
|
||||
res = gather_map(outputs)
|
||||
finally:
|
||||
gather_map = None # type: ignore[assignment]
|
||||
return res
|
||||
@@ -0,0 +1,310 @@
|
||||
from collections import OrderedDict
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch._C import _disabled_torch_function_impl
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Parameter",
|
||||
"UninitializedParameter",
|
||||
"is_lazy",
|
||||
"Buffer",
|
||||
"UninitializedBuffer",
|
||||
"UninitializedTensorMixin",
|
||||
]
|
||||
|
||||
|
||||
# Metaclass to combine _TensorMeta and the instance check override for Parameter.
|
||||
class _ParameterMeta(torch._C._TensorMeta):
|
||||
# Make `isinstance(t, Parameter)` return True for custom tensor instances that have the _is_param flag.
|
||||
def __instancecheck__(self, instance) -> bool:
|
||||
if self is Parameter:
|
||||
if isinstance(instance, torch.Tensor) and getattr(
|
||||
instance, "_is_param", False
|
||||
):
|
||||
return True
|
||||
return super().__instancecheck__(instance)
|
||||
|
||||
|
||||
class Parameter(torch.Tensor, metaclass=_ParameterMeta):
|
||||
r"""A kind of Tensor that is to be considered a module parameter.
|
||||
|
||||
Parameters are :class:`~torch.Tensor` subclasses, that have a
|
||||
very special property when used with :class:`Module` s - when they're
|
||||
assigned as Module attributes they are automatically added to the list of
|
||||
its parameters, and will appear e.g. in :meth:`~Module.parameters` iterator.
|
||||
Assigning a Tensor doesn't have such effect. This is because one might
|
||||
want to cache some temporary state, like last hidden state of the RNN, in
|
||||
the model. If there was no such class as :class:`Parameter`, these
|
||||
temporaries would get registered too.
|
||||
|
||||
Args:
|
||||
data (Tensor): parameter tensor.
|
||||
requires_grad (bool, optional): if the parameter requires gradient. Note that
|
||||
the torch.no_grad() context does NOT affect the default behavior of
|
||||
Parameter creation--the Parameter will still have `requires_grad=True` in
|
||||
:class:`~no_grad` mode. See :ref:`locally-disable-grad-doc` for more
|
||||
details. Default: `True`
|
||||
"""
|
||||
|
||||
def __new__(cls, data=None, requires_grad=True):
|
||||
if data is None:
|
||||
data = torch.empty(0)
|
||||
if type(data) is torch.Tensor or type(data) is Parameter:
|
||||
# For ease of BC maintenance, keep this path for standard Tensor.
|
||||
# Eventually (tm), we should change the behavior for standard Tensor to match.
|
||||
return torch.Tensor._make_subclass(cls, data, requires_grad)
|
||||
|
||||
# Path for custom tensors: set a flag on the instance to indicate parameter-ness.
|
||||
t = data.detach().requires_grad_(requires_grad)
|
||||
if type(t) is not type(data):
|
||||
raise RuntimeError(
|
||||
f"Creating a Parameter from an instance of type {type(data).__name__} "
|
||||
"requires that detach() returns an instance of the same type, but return "
|
||||
f"type {type(t).__name__} was found instead. To use the type as a "
|
||||
"Parameter, please correct the detach() semantics defined by "
|
||||
"its __torch_dispatch__() implementation."
|
||||
)
|
||||
t._is_param = True
|
||||
return t
|
||||
|
||||
# Note: the 3 methods below only apply to standard Tensor. Parameters of custom tensor types
|
||||
# are still considered that custom tensor type and these methods will not be called for them.
|
||||
def __deepcopy__(self, memo):
|
||||
if id(self) in memo:
|
||||
return memo[id(self)]
|
||||
else:
|
||||
result = type(self)(
|
||||
self.data.clone(memory_format=torch.preserve_format), self.requires_grad
|
||||
)
|
||||
memo[id(self)] = result
|
||||
return result
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
def __repr__(self) -> str:
|
||||
return "Parameter containing:\n" + super().__repr__()
|
||||
|
||||
def __reduce_ex__(self, proto):
|
||||
state = torch._utils._get_obj_state(self)
|
||||
|
||||
# See Note [Don't serialize hooks]
|
||||
hooks = OrderedDict()
|
||||
if not state:
|
||||
return (
|
||||
torch._utils._rebuild_parameter,
|
||||
(self.data, self.requires_grad, hooks),
|
||||
)
|
||||
|
||||
return (
|
||||
torch._utils._rebuild_parameter_with_state,
|
||||
(self.data, self.requires_grad, hooks, state),
|
||||
)
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
__torch_function__ = _disabled_torch_function_impl
|
||||
|
||||
|
||||
class UninitializedTensorMixin:
|
||||
_allowed_methods = [
|
||||
torch.Tensor.__hash__,
|
||||
torch.Tensor.size,
|
||||
torch.Tensor.copy_,
|
||||
torch.Tensor.is_complex,
|
||||
torch.Tensor.is_floating_point,
|
||||
torch.Tensor.half,
|
||||
torch.Tensor.float,
|
||||
torch.Tensor.double,
|
||||
torch.Tensor.char,
|
||||
torch.Tensor.short,
|
||||
torch.Tensor.int,
|
||||
torch.Tensor.long,
|
||||
torch.Tensor.cuda,
|
||||
torch.Tensor.cpu,
|
||||
torch.Tensor.to,
|
||||
torch.Tensor.get_device,
|
||||
torch._has_compatible_shallow_copy_type,
|
||||
]
|
||||
|
||||
def materialize(self, shape, device=None, dtype=None) -> None:
|
||||
r"""Create a Parameter or Tensor with the same properties of the uninitialized one.
|
||||
|
||||
Given a shape, it materializes a parameter in the same device
|
||||
and with the same `dtype` as the current one or the specified ones in the
|
||||
arguments.
|
||||
|
||||
Args:
|
||||
shape : (tuple): the shape for the materialized tensor.
|
||||
device (:class:`torch.device`): the desired device of the parameters
|
||||
and buffers in this module. Optional.
|
||||
dtype (:class:`torch.dtype`): the desired floating point type of
|
||||
the floating point parameters and buffers in this module. Optional.
|
||||
"""
|
||||
if device is None:
|
||||
device = self.data.device
|
||||
if dtype is None:
|
||||
dtype = self.data.dtype
|
||||
self.data = torch.empty(shape, device=device, dtype=dtype)
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
self.__class__ = self.cls_to_become
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
raise RuntimeError(
|
||||
"Can't access the shape of an uninitialized parameter or buffer. "
|
||||
"This error usually happens in `load_state_dict` when trying to load "
|
||||
"an uninitialized parameter into an initialized one. "
|
||||
"Call `forward` to initialize the parameters before accessing their attributes."
|
||||
)
|
||||
|
||||
def share_memory_(self):
|
||||
raise RuntimeError(
|
||||
"Can't share memory on an uninitialized parameter or buffer. "
|
||||
"Call `forward` to initialize the parameters before calling "
|
||||
"`module.share_memory()`."
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<{self.__class__.__name__}>"
|
||||
|
||||
def __reduce_ex__(self, proto):
|
||||
# See Note [Don't serialize hooks]
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return (self.__class__, (self.requires_grad,))
|
||||
|
||||
@classmethod
|
||||
def __torch_function__(cls, func, types, args=(), kwargs=None):
|
||||
# method-wrapper is to detect access to Tensor properties that are
|
||||
# wrapped in descriptors
|
||||
if func in cls._allowed_methods or func.__class__.__name__ == "method-wrapper":
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return super().__torch_function__(func, types, args, kwargs)
|
||||
raise ValueError(
|
||||
f"Attempted to use an uninitialized parameter in {func}. "
|
||||
"This error happens when you are using a `LazyModule` or "
|
||||
f"explicitly manipulating `torch.nn.parameter.{cls.__name__}` "
|
||||
"objects. When using LazyModules Call `forward` with a dummy batch "
|
||||
"to initialize the parameters before calling torch functions"
|
||||
)
|
||||
|
||||
|
||||
def is_lazy(param: Any) -> bool:
|
||||
"""
|
||||
Returns whether ``param`` is an ``UninitializedParameter`` or ``UninitializedBuffer``.
|
||||
|
||||
Args:
|
||||
param (Any): the input to check.
|
||||
"""
|
||||
return isinstance(param, UninitializedTensorMixin)
|
||||
|
||||
|
||||
# pyrefly: ignore [inconsistent-inheritance]
|
||||
class UninitializedParameter(UninitializedTensorMixin, Parameter):
|
||||
r"""A parameter that is not initialized.
|
||||
|
||||
Uninitialized Parameters are a special case of :class:`torch.nn.Parameter`
|
||||
where the shape of the data is still unknown.
|
||||
|
||||
Unlike a :class:`torch.nn.Parameter`, uninitialized parameters
|
||||
hold no data and attempting to access some properties, like their shape,
|
||||
will throw a runtime error. The only operations that can be performed on a uninitialized
|
||||
parameter are changing its datatype, moving it to a different device and
|
||||
converting it to a regular :class:`torch.nn.Parameter`.
|
||||
|
||||
The default device or dtype to use when the parameter is materialized can be set
|
||||
during construction using e.g. ``device='cuda'``.
|
||||
"""
|
||||
|
||||
cls_to_become = Parameter
|
||||
|
||||
def __new__(cls, requires_grad=True, device=None, dtype=None) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
data = torch.empty(0, **factory_kwargs)
|
||||
# pyrefly: ignore [bad-return]
|
||||
return torch.Tensor._make_subclass(cls, data, requires_grad)
|
||||
|
||||
def __deepcopy__(self, memo):
|
||||
if id(self) in memo:
|
||||
return memo[id(self)]
|
||||
else:
|
||||
result = type(self)(self.requires_grad, self.data.device, self.data.dtype)
|
||||
memo[id(self)] = result
|
||||
return result
|
||||
|
||||
|
||||
# Metaclass to combine _TensorMeta and the instance check override for Buffer.
|
||||
class _BufferMeta(torch._C._TensorMeta):
|
||||
# Make `isinstance(t, Buffer)` return True for custom tensor instances that have the _is_buffer flag.
|
||||
def __instancecheck__(self, instance) -> bool:
|
||||
if self is Buffer:
|
||||
if isinstance(instance, torch.Tensor) and getattr(
|
||||
instance, "_is_buffer", False
|
||||
):
|
||||
return True
|
||||
return super().__instancecheck__(instance)
|
||||
|
||||
|
||||
class Buffer(torch.Tensor, metaclass=_BufferMeta):
|
||||
r"""A kind of Tensor that should not be considered a model
|
||||
parameter. For example, BatchNorm's ``running_mean`` is not a parameter, but is part of the module's state.
|
||||
|
||||
Buffers are :class:`~torch.Tensor` subclasses, that have a
|
||||
very special property when used with :class:`Module` s -- when they're
|
||||
assigned as Module attributes they are automatically added to the list of
|
||||
its buffers, and will appear e.g. in :meth:`~torch.nn.Module.buffers` iterator.
|
||||
Assigning a Tensor doesn't have such effect. One can still assign a Tensor as explicitly by using
|
||||
the :meth:`~torch.nn.Module.register_buffer` function.
|
||||
|
||||
Args:
|
||||
data (Tensor): buffer tensor.
|
||||
persistent (bool, optional): whether the buffer is part of the module's
|
||||
:attr:`state_dict`. Default: ``True``
|
||||
"""
|
||||
|
||||
def __new__(cls, data=None, *, persistent=True):
|
||||
if data is None:
|
||||
data = torch.empty(0)
|
||||
|
||||
t = data.detach().requires_grad_(data.requires_grad)
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
t.persistent = persistent
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
t._is_buffer = True
|
||||
return t
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
__torch_function__ = _disabled_torch_function_impl
|
||||
|
||||
|
||||
class UninitializedBuffer(UninitializedTensorMixin, torch.Tensor):
|
||||
r"""A buffer that is not initialized.
|
||||
|
||||
Uninitialized Buffer is a a special case of :class:`torch.Tensor`
|
||||
where the shape of the data is still unknown.
|
||||
|
||||
Unlike a :class:`torch.Tensor`, uninitialized parameters
|
||||
hold no data and attempting to access some properties, like their shape,
|
||||
will throw a runtime error. The only operations that can be performed on a uninitialized
|
||||
parameter are changing its datatype, moving it to a different device and
|
||||
converting it to a regular :class:`torch.Tensor`.
|
||||
|
||||
The default device or dtype to use when the buffer is materialized can be set
|
||||
during construction using e.g. ``device='cuda'``.
|
||||
"""
|
||||
|
||||
cls_to_become = torch.Tensor
|
||||
|
||||
def __new__(
|
||||
cls, requires_grad=False, device=None, dtype=None, persistent=True
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
data = torch.empty(0, **factory_kwargs)
|
||||
ret = torch.Tensor._make_subclass(cls, data, requires_grad)
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
ret.persistent = persistent
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
ret._is_buffer = True
|
||||
# pyrefly: ignore [bad-return]
|
||||
return ret
|
||||
@@ -0,0 +1,43 @@
|
||||
from typing_extensions import TypeIs
|
||||
|
||||
from torch import device, dtype, Tensor
|
||||
|
||||
class Parameter(Tensor):
|
||||
def __init__(self, data: Tensor = ..., requires_grad: bool = ...) -> None: ...
|
||||
|
||||
def is_lazy(
|
||||
param: Tensor,
|
||||
) -> TypeIs[UninitializedParameter | UninitializedBuffer]: ...
|
||||
|
||||
class UninitializedParameter(Tensor):
|
||||
def __init__(self, data: Tensor = ..., requires_grad: bool = ...) -> None: ...
|
||||
def materialize(
|
||||
self,
|
||||
shape: tuple[int, ...],
|
||||
device: device | None = None,
|
||||
dtype: dtype | None = None,
|
||||
) -> None: ...
|
||||
|
||||
class Buffer(Tensor):
|
||||
persistent: bool
|
||||
def __init__(
|
||||
self,
|
||||
data: Tensor = ...,
|
||||
requires_grad: bool = ...,
|
||||
persistent: bool = ...,
|
||||
) -> None: ...
|
||||
|
||||
class UninitializedBuffer(Tensor):
|
||||
persistent: bool
|
||||
def __init__(
|
||||
self,
|
||||
data: Tensor = ...,
|
||||
requires_grad: bool = ...,
|
||||
persistent: bool = ...,
|
||||
) -> None: ...
|
||||
def materialize(
|
||||
self,
|
||||
shape: tuple[int, ...],
|
||||
device: device | None = None,
|
||||
dtype: dtype | None = None,
|
||||
) -> None: ...
|
||||
@@ -0,0 +1,19 @@
|
||||
# flake8: noqa: F401
|
||||
r"""QAT Dynamic Modules.
|
||||
|
||||
This package is in the process of being deprecated.
|
||||
Please, use `torch.ao.nn.qat.dynamic` instead.
|
||||
"""
|
||||
|
||||
from torch.nn.qat import dynamic, modules # noqa: F403
|
||||
from torch.nn.qat.modules import * # noqa: F403
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Linear",
|
||||
"Conv1d",
|
||||
"Conv2d",
|
||||
"Conv3d",
|
||||
"Embedding",
|
||||
"EmbeddingBag",
|
||||
]
|
||||
@@ -0,0 +1,8 @@
|
||||
# flake8: noqa: F401
|
||||
r"""QAT Dynamic Modules.
|
||||
|
||||
This package is in the process of being deprecated.
|
||||
Please, use `torch.ao.nn.qat.dynamic` instead.
|
||||
"""
|
||||
|
||||
from torch.nn.qat.dynamic.modules import * # noqa: F403
|
||||
@@ -0,0 +1,4 @@
|
||||
from torch.nn.qat.dynamic.modules.linear import Linear
|
||||
|
||||
|
||||
__all__ = ["Linear"]
|
||||
@@ -0,0 +1,11 @@
|
||||
# flake8: noqa: F401
|
||||
r"""QAT Modules.
|
||||
|
||||
This file is in the process of migration to `torch/ao/nn/qat/dynamic`, and
|
||||
is kept here for compatibility while the migration process is ongoing.
|
||||
If you are adding a new entry/functionality, please, add it to the
|
||||
appropriate file under the `torch/ao/nn/qat/dynamic/modules`,
|
||||
while adding an import statement here.
|
||||
"""
|
||||
|
||||
from torch.ao.nn.qat.dynamic.modules.linear import Linear
|
||||
@@ -0,0 +1,21 @@
|
||||
# flake8: noqa: F401
|
||||
r"""QAT Modules.
|
||||
|
||||
This package is in the process of being deprecated.
|
||||
Please, use `torch.ao.nn.qat.modules` instead.
|
||||
"""
|
||||
|
||||
from torch.ao.nn.qat.modules.conv import Conv1d, Conv2d, Conv3d
|
||||
from torch.ao.nn.qat.modules.embedding_ops import Embedding, EmbeddingBag
|
||||
from torch.ao.nn.qat.modules.linear import Linear
|
||||
from torch.nn.qat.modules import conv, embedding_ops, linear
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Linear",
|
||||
"Conv1d",
|
||||
"Conv2d",
|
||||
"Conv3d",
|
||||
"Embedding",
|
||||
"EmbeddingBag",
|
||||
]
|
||||
@@ -0,0 +1,11 @@
|
||||
# flake8: noqa: F401
|
||||
r"""QAT Modules.
|
||||
|
||||
This file is in the process of migration to `torch/ao/nn/qat`, and
|
||||
is kept here for compatibility while the migration process is ongoing.
|
||||
If you are adding a new entry/functionality, please, add it to the
|
||||
appropriate file under the `torch/ao/nn/qat/modules`,
|
||||
while adding an import statement here.
|
||||
"""
|
||||
|
||||
from torch.ao.nn.qat.modules.conv import Conv1d, Conv2d, Conv3d
|
||||
@@ -0,0 +1,14 @@
|
||||
# flake8: noqa: F401
|
||||
r"""QAT Modules.
|
||||
|
||||
This file is in the process of migration to `torch/ao/nn/qat`, and
|
||||
is kept here for compatibility while the migration process is ongoing.
|
||||
If you are adding a new entry/functionality, please, add it to the
|
||||
appropriate file under the `torch/ao/nn/qat/modules`,
|
||||
while adding an import statement here.
|
||||
"""
|
||||
|
||||
from torch.ao.nn.qat.modules.embedding_ops import Embedding, EmbeddingBag
|
||||
|
||||
|
||||
__all__ = ["Embedding", "EmbeddingBag"]
|
||||
@@ -0,0 +1,11 @@
|
||||
# flake8: noqa: F401
|
||||
r"""QAT Modules.
|
||||
|
||||
This file is in the process of migration to `torch/ao/nn/qat`, and
|
||||
is kept here for compatibility while the migration process is ongoing.
|
||||
If you are adding a new entry/functionality, please, add it to the
|
||||
appropriate file under the `torch/ao/nn/qat/modules`,
|
||||
while adding an import statement here.
|
||||
"""
|
||||
|
||||
from torch.ao.nn.qat.modules.linear import Linear
|
||||
@@ -0,0 +1 @@
|
||||
from torch.nn.quantizable.modules import * # noqa: F403
|
||||
@@ -0,0 +1,9 @@
|
||||
from torch.ao.nn.quantizable.modules.activation import MultiheadAttention
|
||||
from torch.ao.nn.quantizable.modules.rnn import LSTM, LSTMCell
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LSTM",
|
||||
"LSTMCell",
|
||||
"MultiheadAttention",
|
||||
]
|
||||
@@ -0,0 +1,11 @@
|
||||
# flake8: noqa: F401
|
||||
r"""Quantizable Modules.
|
||||
|
||||
This file is in the process of migration to `torch/ao/nn/quantizable`, and
|
||||
is kept here for compatibility while the migration process is ongoing.
|
||||
If you are adding a new entry/functionality, please, add it to the
|
||||
appropriate file under the `torch/ao/nn/quantizable/modules`,
|
||||
while adding an import statement here.
|
||||
"""
|
||||
|
||||
from torch.ao.nn.quantizable.modules.activation import MultiheadAttention
|
||||
@@ -0,0 +1,11 @@
|
||||
# flake8: noqa: F401
|
||||
r"""Quantizable Modules.
|
||||
|
||||
This file is in the process of migration to `torch/ao/nn/quantizable`, and
|
||||
is kept here for compatibility while the migration process is ongoing.
|
||||
If you are adding a new entry/functionality, please, add it to the
|
||||
appropriate file under the `torch/ao/nn/quantizable/modules`,
|
||||
while adding an import statement here.
|
||||
"""
|
||||
|
||||
from torch.ao.nn.quantizable.modules.rnn import LSTM, LSTMCell
|
||||
@@ -0,0 +1,39 @@
|
||||
from torch.nn.quantized import dynamic, functional, modules # noqa: F403
|
||||
from torch.nn.quantized.modules import * # noqa: F403
|
||||
from torch.nn.quantized.modules import MaxPool2d
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BatchNorm2d",
|
||||
"BatchNorm3d",
|
||||
"Conv1d",
|
||||
"Conv2d",
|
||||
"Conv3d",
|
||||
"ConvTranspose1d",
|
||||
"ConvTranspose2d",
|
||||
"ConvTranspose3d",
|
||||
"DeQuantize",
|
||||
"Dropout",
|
||||
"ELU",
|
||||
"Embedding",
|
||||
"EmbeddingBag",
|
||||
"GroupNorm",
|
||||
"Hardswish",
|
||||
"InstanceNorm1d",
|
||||
"InstanceNorm2d",
|
||||
"InstanceNorm3d",
|
||||
"LayerNorm",
|
||||
"LeakyReLU",
|
||||
"Linear",
|
||||
"LSTM",
|
||||
"MultiheadAttention",
|
||||
"PReLU",
|
||||
"Quantize",
|
||||
"ReLU6",
|
||||
"Sigmoid",
|
||||
"Softmax",
|
||||
# Wrapper modules
|
||||
"FloatFunctional",
|
||||
"FXFloatFunctional",
|
||||
"QFunctional",
|
||||
]
|
||||
@@ -0,0 +1 @@
|
||||
from torch.nn.quantized._reference.modules import * # noqa: F403
|
||||
+39
@@ -0,0 +1,39 @@
|
||||
# flake8: noqa: F401
|
||||
r"""Quantized Reference Modules.
|
||||
|
||||
This module is in the process of migration to
|
||||
`torch/ao/nn/quantized/reference`, and is kept here for
|
||||
compatibility while the migration process is ongoing.
|
||||
If you are adding a new entry/functionality, please, add it to the
|
||||
appropriate file under the `torch/ao/nn/quantized/reference`,
|
||||
while adding an import statement here.
|
||||
"""
|
||||
|
||||
from torch.ao.nn.quantized.reference.modules.conv import (
|
||||
Conv1d,
|
||||
Conv2d,
|
||||
Conv3d,
|
||||
ConvTranspose1d,
|
||||
ConvTranspose2d,
|
||||
ConvTranspose3d,
|
||||
)
|
||||
from torch.ao.nn.quantized.reference.modules.linear import Linear
|
||||
from torch.ao.nn.quantized.reference.modules.rnn import GRUCell, LSTM, LSTMCell, RNNCell
|
||||
from torch.ao.nn.quantized.reference.modules.sparse import Embedding, EmbeddingBag
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Linear",
|
||||
"Conv1d",
|
||||
"Conv2d",
|
||||
"Conv3d",
|
||||
"ConvTranspose1d",
|
||||
"ConvTranspose2d",
|
||||
"ConvTranspose3d",
|
||||
"RNNCell",
|
||||
"LSTMCell",
|
||||
"GRUCell",
|
||||
"LSTM",
|
||||
"Embedding",
|
||||
"EmbeddingBag",
|
||||
]
|
||||
+21
@@ -0,0 +1,21 @@
|
||||
# flake8: noqa: F401
|
||||
r"""Quantized Reference Modules.
|
||||
|
||||
This module is in the process of migration to
|
||||
`torch/ao/nn/quantized/reference`, and is kept here for
|
||||
compatibility while the migration process is ongoing.
|
||||
If you are adding a new entry/functionality, please, add it to the
|
||||
appropriate file under the `torch/ao/nn/quantized/reference`,
|
||||
while adding an import statement here.
|
||||
"""
|
||||
|
||||
from torch.ao.nn.quantized.reference.modules.conv import (
|
||||
_ConvNd,
|
||||
_ConvTransposeNd,
|
||||
Conv1d,
|
||||
Conv2d,
|
||||
Conv3d,
|
||||
ConvTranspose1d,
|
||||
ConvTranspose2d,
|
||||
ConvTranspose3d,
|
||||
)
|
||||
+12
@@ -0,0 +1,12 @@
|
||||
# flake8: noqa: F401
|
||||
r"""Quantized Reference Modules.
|
||||
|
||||
This module is in the process of migration to
|
||||
`torch/ao/nn/quantized/reference`, and is kept here for
|
||||
compatibility while the migration process is ongoing.
|
||||
If you are adding a new entry/functionality, please, add it to the
|
||||
appropriate file under the `torch/ao/nn/quantized/reference`,
|
||||
while adding an import statement here.
|
||||
"""
|
||||
|
||||
from torch.ao.nn.quantized.reference.modules.linear import Linear
|
||||
@@ -0,0 +1,19 @@
|
||||
# flake8: noqa: F401
|
||||
r"""Quantized Reference Modules.
|
||||
|
||||
This module is in the process of migration to
|
||||
`torch/ao/nn/quantized/reference`, and is kept here for
|
||||
compatibility while the migration process is ongoing.
|
||||
If you are adding a new entry/functionality, please, add it to the
|
||||
appropriate file under the `torch/ao/nn/quantized/reference`,
|
||||
while adding an import statement here.
|
||||
"""
|
||||
|
||||
from torch.ao.nn.quantized.reference.modules.rnn import (
|
||||
GRUCell,
|
||||
LSTM,
|
||||
LSTMCell,
|
||||
RNNBase,
|
||||
RNNCell,
|
||||
RNNCellBase,
|
||||
)
|
||||
+12
@@ -0,0 +1,12 @@
|
||||
# flake8: noqa: F401
|
||||
r"""Quantized Reference Modules.
|
||||
|
||||
This module is in the process of migration to
|
||||
`torch/ao/nn/quantized/reference`, and is kept here for
|
||||
compatibility while the migration process is ongoing.
|
||||
If you are adding a new entry/functionality, please, add it to the
|
||||
appropriate file under the `torch/ao/nn/quantized/reference`,
|
||||
while adding an import statement here.
|
||||
"""
|
||||
|
||||
from torch.ao.nn.quantized.reference.modules.sparse import Embedding, EmbeddingBag
|
||||
+18
@@ -0,0 +1,18 @@
|
||||
# flake8: noqa: F401
|
||||
r"""Quantized Reference Modules.
|
||||
|
||||
This module is in the process of migration to
|
||||
`torch/ao/nn/quantized/reference`, and is kept here for
|
||||
compatibility while the migration process is ongoing.
|
||||
If you are adding a new entry/functionality, please, add it to the
|
||||
appropriate file under the `torch/ao/nn/quantized/reference`,
|
||||
while adding an import statement here.
|
||||
"""
|
||||
|
||||
from torch.ao.nn.quantized.reference.modules.utils import (
|
||||
_get_weight_qparam_keys,
|
||||
_quantize_and_dequantize_weight,
|
||||
_quantize_weight,
|
||||
_save_weight_qparams,
|
||||
ReferenceQuantizedModule,
|
||||
)
|
||||
@@ -0,0 +1 @@
|
||||
from torch.ao.nn.quantized.dynamic import * # noqa: F403
|
||||
+43
@@ -0,0 +1,43 @@
|
||||
# flake8: noqa: F401
|
||||
r"""Quantized Dynamic Modules.
|
||||
|
||||
This file is in the process of migration to `torch/ao/nn/quantized/dynamic`,
|
||||
and is kept here for compatibility while the migration process is ongoing.
|
||||
If you are adding a new entry/functionality, please, add it to the
|
||||
appropriate file under the `torch/ao/nn/quantized/dynamic`,
|
||||
while adding an import statement here.
|
||||
"""
|
||||
|
||||
from torch.ao.nn.quantized.dynamic.modules import conv, linear, rnn
|
||||
from torch.ao.nn.quantized.dynamic.modules.conv import (
|
||||
Conv1d,
|
||||
Conv2d,
|
||||
Conv3d,
|
||||
ConvTranspose1d,
|
||||
ConvTranspose2d,
|
||||
ConvTranspose3d,
|
||||
)
|
||||
from torch.ao.nn.quantized.dynamic.modules.linear import Linear
|
||||
from torch.ao.nn.quantized.dynamic.modules.rnn import (
|
||||
GRU,
|
||||
GRUCell,
|
||||
LSTM,
|
||||
LSTMCell,
|
||||
RNNCell,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Linear",
|
||||
"LSTM",
|
||||
"GRU",
|
||||
"LSTMCell",
|
||||
"RNNCell",
|
||||
"GRUCell",
|
||||
"Conv1d",
|
||||
"Conv2d",
|
||||
"Conv3d",
|
||||
"ConvTranspose1d",
|
||||
"ConvTranspose2d",
|
||||
"ConvTranspose3d",
|
||||
]
|
||||
@@ -0,0 +1,28 @@
|
||||
# flake8: noqa: F401
|
||||
r"""Quantized Dynamic Modules.
|
||||
|
||||
This file is in the process of migration to `torch/ao/nn/quantized/dynamic`,
|
||||
and is kept here for compatibility while the migration process is ongoing.
|
||||
If you are adding a new entry/functionality, please, add it to the
|
||||
appropriate file under the `torch/ao/nn/quantized/dynamic/modules`,
|
||||
while adding an import statement here.
|
||||
"""
|
||||
|
||||
from torch.ao.nn.quantized.dynamic.modules.conv import (
|
||||
Conv1d,
|
||||
Conv2d,
|
||||
Conv3d,
|
||||
ConvTranspose1d,
|
||||
ConvTranspose2d,
|
||||
ConvTranspose3d,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Conv1d",
|
||||
"Conv2d",
|
||||
"Conv3d",
|
||||
"ConvTranspose1d",
|
||||
"ConvTranspose2d",
|
||||
"ConvTranspose3d",
|
||||
]
|
||||
@@ -0,0 +1,11 @@
|
||||
# flake8: noqa: F401
|
||||
r"""Quantized Dynamic Modules.
|
||||
|
||||
This file is in the process of migration to `torch/ao/nn/quantized/dynamic`,
|
||||
and is kept here for compatibility while the migration process is ongoing.
|
||||
If you are adding a new entry/functionality, please, add it to the
|
||||
appropriate file under the `torch/ao/nn/quantized/dynamic/modules`,
|
||||
while adding an import statement here.
|
||||
"""
|
||||
|
||||
from torch.ao.nn.quantized.dynamic.modules.linear import Linear
|
||||
@@ -0,0 +1,34 @@
|
||||
# flake8: noqa: F401
|
||||
r"""Quantized Dynamic Modules.
|
||||
|
||||
This file is in the process of migration to `torch/ao/nn/quantized/dynamic`,
|
||||
and is kept here for compatibility while the migration process is ongoing.
|
||||
If you are adding a new entry/functionality, please, add it to the
|
||||
appropriate file under the `torch/ao/nn/quantized/dynamic/modules`,
|
||||
while adding an import statement here.
|
||||
"""
|
||||
|
||||
from torch.ao.nn.quantized.dynamic.modules.rnn import (
|
||||
GRU,
|
||||
GRUCell,
|
||||
LSTM,
|
||||
LSTMCell,
|
||||
pack_weight_bias,
|
||||
PackedParameter,
|
||||
RNNBase,
|
||||
RNNCell,
|
||||
RNNCellBase,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"pack_weight_bias",
|
||||
"PackedParameter",
|
||||
"RNNBase",
|
||||
"LSTM",
|
||||
"GRU",
|
||||
"RNNCellBase",
|
||||
"RNNCell",
|
||||
"LSTMCell",
|
||||
"GRUCell",
|
||||
]
|
||||
@@ -0,0 +1,10 @@
|
||||
r"""nn.quantized.functional.
|
||||
|
||||
Quantized equivalents of the `nn.functional`.
|
||||
|
||||
Note::
|
||||
This location is in the process of being deprecated.
|
||||
Please, use the `torch.ao.nn.quantized.functional` instead.
|
||||
"""
|
||||
|
||||
from torch.ao.nn.quantized.functional import * # noqa: F401,F403
|
||||
@@ -0,0 +1,97 @@
|
||||
r"""Quantized Modules.
|
||||
|
||||
Note::
|
||||
The `torch.nn.quantized` namespace is in the process of being deprecated.
|
||||
Please, use `torch.ao.nn.quantized` instead.
|
||||
"""
|
||||
|
||||
# The following imports are needed in case the user decides
|
||||
# to import the files directly,
|
||||
# s.a. `from torch.nn.quantized.modules.conv import ...`.
|
||||
# No need to add them to the `__all__`.
|
||||
from torch.ao.nn.quantized.modules import (
|
||||
activation,
|
||||
batchnorm,
|
||||
conv,
|
||||
DeQuantize,
|
||||
dropout,
|
||||
embedding_ops,
|
||||
functional_modules,
|
||||
linear,
|
||||
MaxPool2d,
|
||||
normalization,
|
||||
Quantize,
|
||||
rnn,
|
||||
utils,
|
||||
)
|
||||
from torch.ao.nn.quantized.modules.activation import (
|
||||
ELU,
|
||||
Hardswish,
|
||||
LeakyReLU,
|
||||
MultiheadAttention,
|
||||
PReLU,
|
||||
ReLU6,
|
||||
Sigmoid,
|
||||
Softmax,
|
||||
)
|
||||
from torch.ao.nn.quantized.modules.batchnorm import BatchNorm2d, BatchNorm3d
|
||||
from torch.ao.nn.quantized.modules.conv import (
|
||||
Conv1d,
|
||||
Conv2d,
|
||||
Conv3d,
|
||||
ConvTranspose1d,
|
||||
ConvTranspose2d,
|
||||
ConvTranspose3d,
|
||||
)
|
||||
from torch.ao.nn.quantized.modules.dropout import Dropout
|
||||
from torch.ao.nn.quantized.modules.embedding_ops import Embedding, EmbeddingBag
|
||||
from torch.ao.nn.quantized.modules.functional_modules import (
|
||||
FloatFunctional,
|
||||
FXFloatFunctional,
|
||||
QFunctional,
|
||||
)
|
||||
from torch.ao.nn.quantized.modules.linear import Linear
|
||||
from torch.ao.nn.quantized.modules.normalization import (
|
||||
GroupNorm,
|
||||
InstanceNorm1d,
|
||||
InstanceNorm2d,
|
||||
InstanceNorm3d,
|
||||
LayerNorm,
|
||||
)
|
||||
from torch.ao.nn.quantized.modules.rnn import LSTM
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BatchNorm2d",
|
||||
"BatchNorm3d",
|
||||
"Conv1d",
|
||||
"Conv2d",
|
||||
"Conv3d",
|
||||
"ConvTranspose1d",
|
||||
"ConvTranspose2d",
|
||||
"ConvTranspose3d",
|
||||
"DeQuantize",
|
||||
"ELU",
|
||||
"Embedding",
|
||||
"EmbeddingBag",
|
||||
"GroupNorm",
|
||||
"Hardswish",
|
||||
"InstanceNorm1d",
|
||||
"InstanceNorm2d",
|
||||
"InstanceNorm3d",
|
||||
"LayerNorm",
|
||||
"LeakyReLU",
|
||||
"Linear",
|
||||
"LSTM",
|
||||
"MultiheadAttention",
|
||||
"Quantize",
|
||||
"ReLU6",
|
||||
"Sigmoid",
|
||||
"Softmax",
|
||||
"Dropout",
|
||||
"PReLU",
|
||||
# Wrapper modules
|
||||
"FloatFunctional",
|
||||
"FXFloatFunctional",
|
||||
"QFunctional",
|
||||
]
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user