5.4 自定义层-526互联

有时，我们也许需要在神经网络中用到一个框架中没有的层，这就需要我们自定义层。

1. 不带参数的自定义层

import torch 
from torch import nn
from torch.nn import functional as F

class CenteredLayer(nn.Module):
    def __init__(self):
        super().__init__()
        
    def forward(self,X):
        return X - X.mean()

layer = CenteredLayer()
print(layer(torch.FloatTensor([1,2,3,4,5])))

代码如上所示。不带参数的自定义层在__init__()中只需要实现父类的初始化即可，该层要完成的操作在forward()中实现。

我们可以将自己设计的层组件放到更加复杂的模型中去：

net = nn.Sequential(nn.Linear(8, 128), CenteredLayer())

我们做一个健全性检查，检查net是否正常工作：首先生成一组从均匀分布[0,1)中抽样的数据，然后将它们用net变换，即经过一个Linear层，然后经过我们定义的CenteredLayer层。无论Linear层将数据变换成了什么样，经过CenteredLayer层之后，这一组数据的均值Y.mean()应该为0.这是因为，任何一组数据中，每一个值减去这组数据的均值后，新的这组数据的均值都应该是0.

Y = net(torch.rand(4, 8))
Y.mean()

可以看到输出结果为：

2. 带参数的自定义层

作为例子，我们实现一个自定义版本的全连接层。

class MyLinear(nn.Module):
    def __init__(self,in_units,units):
        super().__init__()
        self.weight = nn.Parameter(torch.normal(0,0.01,size=(in_units,units)))
        self.bias = nn.Parameter(torch.zeros(units))
    
    def forward(self,X):
        return torch.matmul(X,self.weight.data) + self.bias.data

注意，在__init__()函数中，self.weight和self.bias的定义用nn.Parameter()包起来了，在forward()中，计算的时候一定是要使用self.weight.data以及self.bias.data，而不是self.weight以及self.bias。

自定义线性层已经定义好了，我们便可以创建它的实例，访问其权重：

可以使用自定义层构建模型：