Pytorch 实现自定义参数层

谁践踏了优雅 2021-06-24 16:00 494阅读 0赞

注意，一般官方接口都带有可导功能，如果你实现的层不具有可导功能，就需要自己实现梯度的反向传递。

官方Linear层：
    class Linear(Module):
        def __init__(self, in_features, out_features, bias=True):
            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):
            stdv = 1. / math.sqrt(self.weight.size(1))
            self.weight.data.uniform_(-stdv, stdv)
            if self.bias is not None:
                self.bias.data.uniform_(-stdv, stdv)
    
        def forward(self, input):
            return F.linear(input, self.weight, self.bias)
    
        def extra_repr(self):
            return 'in_features={}, out_features={}, bias={}'.format(
                self.in_features, self.out_features, self.bias is not None
            )

实现view层
    class Reshape(nn.Module):
        def __init__(self, *args):
            super(Reshape, self).__init__()
            self.shape = args
    
        def forward(self, x):
            return x.view((x.size(0),)+self.shape)

实现LinearWise层
    class LinearWise(nn.Module):
        def __init__(self, in_features, bias=True):
            super(LinearWise, self).__init__()
            self.in_features = in_features
    
            self.weight = nn.Parameter(torch.Tensor(self.in_features))
            if bias:
                self.bias = nn.Parameter(torch.Tensor(self.in_features))
            else:
                self.register_parameter('bias', None)
            self.reset_parameters()
    
        def reset_parameters(self):
            stdv = 1. / math.sqrt(self.weight.size(0))
            self.weight.data.uniform_(-stdv, stdv)
            if self.bias is not None:
                self.bias.data.uniform_(-stdv, stdv)
    
        def forward(self, input):
            x = input * self.weight
            if self.bias is not None:
                x = x + self.bias
            return x