python pytorch自定义_Pytorch 实现自定义参数层的例子

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

官方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_feature