Pytorch 复习总结 4-二. 自定义层

和自定义块一样，自定义层也需要实现构造函数和前向传播函数。

1. 无参数层

import torch
from torch import nn

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

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

net = nn.Sequential(nn.Linear(8, 128), CenteredLayer())
X = torch.rand(4, 8)
output = net(X)
print(output.mean())	# tensor(0., grad_fn=<MeanBackward0>)

2. 有参数层

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

class MyLinear(nn.Module):
    def __init__(self, in_units, out_units):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(in_units, out_units))
        self.bias = nn.Parameter(torch.randn(out_units,))
    def forward(self, X):
        linear = torch.matmul(X, self.weight.data) + self.bias.data
        return F.relu(linear)

net = nn.Sequential(
    MyLinear(64, 8), 
    MyLinear(8, 1)
)
X = torch.rand(2, 64)
output = net(X)
print(output)       # tensor([[11.9497], [13.9729]])