mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
9b52654620
Summary: Fixes https://github.com/pytorch/pytorch/issues/45771 Pull Request resolved: https://github.com/pytorch/pytorch/pull/45772 Reviewed By: mruberry Differential Revision: D24682013 Pulled By: albanD fbshipit-source-id: e32bc4fe9c586c079f7070924a874c70f3d127fa
229 lines
8.6 KiB
Python
229 lines
8.6 KiB
Python
import math
|
|
|
|
import torch
|
|
from torch import Tensor
|
|
from torch.nn.parameter import Parameter, UninitializedParameter
|
|
from .. import functional as F
|
|
from .. import init
|
|
from .module import Module
|
|
from .lazy import LazyModuleMixin
|
|
|
|
|
|
class Identity(Module):
|
|
r"""A placeholder identity operator that is argument-insensitive.
|
|
|
|
Args:
|
|
args: any argument (unused)
|
|
kwargs: any keyword argument (unused)
|
|
|
|
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, **kwargs):
|
|
super(Identity, self).__init__()
|
|
|
|
def forward(self, input: Tensor) -> Tensor:
|
|
return input
|
|
|
|
|
|
class Linear(Module):
|
|
r"""Applies a linear transformation to the incoming data: :math:`y = xA^T + b`
|
|
|
|
This module supports :ref:`TensorFloat32<tf32_on_ampere>`.
|
|
|
|
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:`(N, *, H_{in})` where :math:`*` means any number of
|
|
additional dimensions and :math:`H_{in} = \text{in\_features}`
|
|
- Output: :math:`(N, *, H_{out})` where all but the last dimension
|
|
are the same shape as the input and :math:`H_{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) -> None:
|
|
super(Linear, self).__init__()
|
|
self.in_features = in_features
|
|
self.out_features = out_features
|
|
self.weight = Parameter(torch.Tensor(out_features, in_features))
|
|
if bias:
|
|
self.bias = Parameter(torch.Tensor(out_features))
|
|
else:
|
|
self.register_parameter('bias', None)
|
|
self.reset_parameters()
|
|
|
|
def reset_parameters(self) -> None:
|
|
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)
|
|
init.uniform_(self.bias, -bound, bound)
|
|
|
|
def forward(self, input: Tensor) -> Tensor:
|
|
return F.linear(input, self.weight, self.bias)
|
|
|
|
def extra_repr(self) -> str:
|
|
return 'in_features={}, out_features={}, bias={}'.format(
|
|
self.in_features, self.out_features, self.bias is not None
|
|
)
|
|
|
|
|
|
# This class exists solely for Transformer; it has an annotation stating
|
|
# that bias is never None, which appeases TorchScript
|
|
class _LinearWithBias(Linear):
|
|
bias: Tensor # type: ignore
|
|
|
|
def __init__(self, in_features: int, out_features: int) -> None:
|
|
super().__init__(in_features, out_features, bias=True) # type: ignore
|
|
|
|
|
|
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
|
|
in2_features: size of each second 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:
|
|
- Input1: :math:`(N, *, H_{in1})` where :math:`H_{in1}=\text{in1\_features}` and
|
|
:math:`*` means any number of additional dimensions. All but the last dimension
|
|
of the inputs should be the same.
|
|
- Input2: :math:`(N, *, H_{in2})` where :math:`H_{in2}=\text{in2\_features}`.
|
|
- Output: :math:`(N, *, H_{out})` where :math:`H_{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) -> None:
|
|
super(Bilinear, self).__init__()
|
|
self.in1_features = in1_features
|
|
self.in2_features = in2_features
|
|
self.out_features = out_features
|
|
self.weight = Parameter(torch.Tensor(out_features, in1_features, in2_features))
|
|
|
|
if bias:
|
|
self.bias = Parameter(torch.Tensor(out_features))
|
|
else:
|
|
self.register_parameter('bias', None)
|
|
self.reset_parameters()
|
|
|
|
def reset_parameters(self) -> None:
|
|
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:
|
|
return F.bilinear(input1, input2, self.weight, self.bias)
|
|
|
|
def extra_repr(self) -> str:
|
|
return 'in1_features={}, in2_features={}, out_features={}, bias={}'.format(
|
|
self.in1_features, self.in2_features, self.out_features, self.bias is not None
|
|
)
|
|
|
|
|
|
class LazyLinear(LazyModuleMixin, Linear):
|
|
r"""A :class:`torch.nn.Linear` module with lazy initialization.
|
|
|
|
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.
|
|
|
|
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]
|
|
weight: UninitializedParameter
|
|
|
|
def __init__(self, out_features: int, bias: bool = True) -> None:
|
|
super().__init__(0, out_features, bias)
|
|
self.weight = UninitializedParameter()
|
|
|
|
def reset_parameters(self) -> None:
|
|
if not self.has_uninitialized_params() and self.in_features != 0:
|
|
super().reset_parameters()
|
|
|
|
def initialize_parameters(self, input) -> None: # type: ignore
|
|
if self.has_uninitialized_params():
|
|
with torch.no_grad():
|
|
self.in_features = input.shape[-1]
|
|
self.weight.materialize((self.out_features, self.in_features))
|
|
self.reset_parameters()
|
|
# TODO: PartialLinear - maybe in sparse?
|