# 创建不带有参数层
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())
Y=net(torch.rand(4,8))
Y.mean()
tensor(8.3819e-09, grad_fn=<MeanBackward0>)
# 创建带有参数层
class MyLinear(nn.Module):
    def __init__(self,in_units,units):
        '''
        使用参数自定义层
        使用Parameter类自定义线性层
        :param in_units 前一层神经元数量
        :param 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)

dense=MyLinear(5,3)
dense.weight

# 使用自定义层直接执行前向传播计算
dense(torch.rand(2,5))

# 使用自定义层构建网络
net=nn.Sequential(MyLinear(64,8),MyLinear(8,1))
net(torch.rand(2,64))
tensor([[5.6085],
        [7.0500]])
Logo

技术共进,成长同行——讯飞AI开发者社区

更多推荐