nn.Parameter和F.linear
class TextRNN(nn.Module):
def __init__(self,
input_size = 768,
hidden_size = 164,
output_size = 768,
n_layers = 1,
dropout = 0.1,
args = None
):
super(TextRNN, self).__init__()
self.h0 = nn.Parameter(torch.Tensor(n_layers, hidden_size))
self.c0 = nn.Parameter(torch.Tensor(n_layers, hidden_size))
self.dropout = nn.Dropout(dropout)
self._activate = nn.Tanh()
# block input
self.Wz = nn.Parameter(torch.Tensor(input_size, hidden_size)) # [8, 768, 164]
self.Rz = nn.Parameter(torch.Tensor(hidden_size, hidden_size)) # [8, 164, 164]
self.bz = nn.Parameter(torch.Tensor(hidden_size, hidden_size)) # [8, 164, 164]
# input gate
self.Ai = nn.Parameter(torch.Tensor(input_size, hidden_size))
self.Wi = nn.Parameter(torch.Tensor(input_size, hidden_size))
self.Ri = nn.Parameter(torch.Tensor(hidden_size, hidden_size))
self.Pi = nn.Parameter(torch.Tensor(n_layers, hidden_size))
self.bi = nn.Parameter(torch.Tensor(n_layers, hidden_size))
# forget gate
# input_size = 768, hidden_size = 164
self.Af = nn.Parameter(torch.Tensor(input_size, hidden_size))
self.Wf = nn.Parameter(torch.Tensor(input_size, hidden_size))
self.Rf = nn.Parameter(torch.Tensor(hidden_size, hidden_size))
self.Pf = nn.Parameter(torch.Tensor(n_layers, hidden_size))
self.bf = nn.Parameter(torch.Tensor(n_layers, hidden_size))
# output gate
self.Ao = nn.Parameter(torch.Tensor(input_size, hidden_size))
self.Wo = nn.Parameter(torch.Tensor(input_size, hidden_size))
self.Ro = nn.Parameter(torch.Tensor(hidden_size, hidden_size))
self.Po = nn.Parameter(torch.Tensor(n_layers, hidden_size))
self.bo = nn.Parameter(torch.Tensor(n_layers, hidden_size))
self.reset_weigths()
def reset_weigths(self):
"""reset weights
"""
for weight in self.parameters():
nn.init.xavier_normal_(weight)
def forward(self, input_ids, input_attention): # [8, 164, 768], [8, 164, 768]
# input_ids: [8, 164, 768], input_attention: [8, 164, 768]
batch_size = input_ids.shape[0]
# [8, 164, 164]
z = self._activate(F.linear(input_attention, self.Wz.t()) + torch.mm(self.h0, self.Rz) + self.bz)
# input_gate
# [8, 164, 164]
input_gate = nn.Sigmoid()(F.linear(input_ids, self.Ai.t()) +
F.linear(input_attention, self.Wi.t()) + torch.mm(self.h0, self.Ri)
+ self.Pi * self.c0 + self.bi
)
# print(input_gate.shape)
# forget gate
# [8, 164, 164]
forget_gate = nn.Sigmoid()(F.linear(input_ids, self.Ai.t()) +
F.linear(input_attention, self.Wf.t()) +
torch.mm(self.h0, self.Rf) +
self.Pf * self.c0 + self.bf)
# [8, 164, 164]
new_c = self.c0 * forget_gate + z * input_gate
# print(new_c.shape)
# output_gate
# [8, 164, 164]
output_gate = nn.Sigmoid()(F.linear(input_attention, self.Wo.t()) +
torch.mm(self.h0, self.Ro) +
self.Po * self.c0 + self.bo)
# block output
# [8, 164, 164]
new_h = output_gate * self._activate(new_c)
return new_h, (new_c, new_h)
nn.Linear实现细节
import math
import torch
from torch import Tensor
from torch.nn.parameter import Parameter
from .. import functional as F
from .. import init
from .module import Module
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} = ext{in\_features}`
- Output: :math:`(N, *, H_{out})` where all but the last dimension
are the same shape as the input and :math:`H_{out} = ext{out\_features}`.
Attributes:
weight: the learnable weights of the module of shape
:math:`( ext{out\_features}, ext{in\_features})`. The values are
initialized from :math:`mathcal{U}(-sqrt{k}, sqrt{k})`, where
:math:`k = frac{1}{ ext{in\_features}}`
bias: the learnable bias of the module of shape :math:`( ext{out\_features})`.
If :attr:`bias` is ``True``, the values are initialized from
:math:`mathcal{U}(-sqrt{k}, sqrt{k})` where
:math:`k = frac{1}{ ext{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
def __init__(self, in_features: int, out_features: int) -> None:
super().__init__(in_features, out_features, bias=True)
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}= ext{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}= ext{in2\_features}`.
- Output: :math:`(N, *, H_{out})` where :math:`H_{out}= ext{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:`( ext{out\_features}, ext{in1\_features}, ext{in2\_features})`.
The values are initialized from :math:`mathcal{U}(-sqrt{k}, sqrt{k})`, where
:math:`k = frac{1}{ ext{in1\_features}}`
bias: the learnable bias of the module of shape :math:`( ext{out\_features})`.
If :attr:`bias` is ``True``, the values are initialized from
:math:`mathcal{U}(-sqrt{k}, sqrt{k})`, where
:math:`k = frac{1}{ ext{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
)
# TODO: PartialLinear - maybe in sparse?