1. 不带参数的层

与块类的实现类似，继承nn.Module类并实现前向传播的forward()方法即可定义一个新的层。
例如，定义一个层，读取输入并与均值相减后输出：

import torch
from torch import nn


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

    @staticmethod
    def forward(data):
        return data - data.mean()  # mean() 操作时，数据类型应为浮点或复数型


if __name__ == '__main__':
    layer = CenteredLayer()
    print(layer(torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32)))

输出结果：

1	tensor([-2., -1., 0., 1., 2.])

2. 带参数的层

自定义的层可能需要记录在训练过程中不断更新的参数（如权重和偏置），手动管理参数的过程很繁琐。现代深度学习框架内置的参数管理 API 可以很方便地管理访问、初始化、共享、保存和加载参数。
在 PyTorch 中，nn.Parameter()用于将定义的张量自动添加到Module.parameters参数迭代器，实现对参数的注册。使参数自动参与梯度计算requires_grad=True与反向传播，保证层间的参数共享。
例如，定义一个全连接层，用weight参数表示权重，用bias参数表示偏置。使用 ReLU 函数激活输出值。同时，可为层指定输出形状的形参in_features和out_features。

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


class LinearReLU(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(in_features, out_features))
        self.bias = nn.Parameter(torch.randn(out_features))

    def forward(self, x):
        return F.relu(x @ self.weight.data + self.bias.data)


if __name__ == '__main__':
    linear = LinearReLU(5, 3)
    print(f'{linear.weight = }')
    print(f'{linear(torch.rand(2, 5)) = }')

输出结果：

linear.weight = Parameter containing:
tensor([[ 7.3203e-01, -5.7376e-01,  5.4769e-01],
        [-9.2961e-02, -5.1345e-01, -2.9226e+00],
        [-1.6591e-03,  1.1632e+00,  1.5749e+00],
        [-1.7442e-01,  4.9862e-01,  1.7185e+00],
        [ 8.3327e-01, -1.0872e+00, -3.3840e-02]], requires_grad=True)
linear(torch.rand(2, 5)) = tensor([[2.0723, 0.4304, 0.0000],
        [1.7495, 1.1964, 0.0000]])