nn.Parameter和F.linear的用法以及参数初始化方式

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} = \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

    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}=\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
        )

# TODO: PartialLinear - maybe in sparse?
posted @ 2021-05-06 16:02  douzujun  阅读(2273)  评论(0编辑  收藏  举报