【动手学深度学习】自定义层

【动手学深度学习】自定义层

不带参数的层

 下面的类不带任何参数

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

class CenteredLayer(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self,X):
        return X  - X.mean()

layer = CenteredLayer()

# 测试
layer(torch.FloatTensor([1,2,3,4,5]))

#  将层作为组件组合到更加复杂的模型中
net = nn.Sequential(nn.Linear(8,128),CenteredLayer())

带参数的层

定义自带参数的层 输入和输出

# 定义自带参数的层  输入和输出  
class MyLinear(nn.Module):
    def __init__(self,in_units,units):
        super().__init__()
        # 初始化权重参数
        self.weight = nn.Parameter(torch.randn(in_units,units))
        # 初始化偏置
        self.bias = nn.Parameter(torch.randn(units,))

    def forward(self,X):
        #  计算线性层
        linear = torch.matmul(X,self.weight.data) + self.bias.data
        return F.relu(linear)
    
#  权重矩阵  5 x 3
linear = MyLinear(5,3)
print(linear.weight)


linear(torch.randn(2,5))  # 两个样本 五个特征

自定义层构建模型，使用内置的全连接层

# 使用自定义岑构建模型 
net = nn.Sequential(MyLinear(64,8),MyLinear(8,1))

net(torch.rand(2,64))