自定義layer

不含模型參數的layer

含模型參數的layer

核心都一樣,自定義一個繼承自nn.Module的類,在類的forward函數里實現該layer的計算,不同的是,帶參數的layer需要用到nn.Parameter

不含模型參數的layer

直接繼承nn.Module

import torch

from torch import nn

class CenteredLayer(nn.Module):

def __init__(self, **kwargs):

super(CenteredLayer, self).__init__(**kwargs)

def forward(self, x):

return x - x.mean()

layer = CenteredLayer()

layer(torch.tensor([1, 2, 3, 4, 5], dtype=torch.float))

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

y = net(torch.rand(4, 8))

y.mean().item()

含模型參數的layer

Parameter

ParameterList

ParameterDict

Parameter類其實是Tensor的子類，如果一個Tensor是Parameter，那么它會自動被添加到模型的參數列表里。所以在自定義含模型參數的層時，我們應該將參數定義成Parameter，除了直接定義成Parameter類外，還可以使用ParameterList和ParameterDict分別定義參數的列表和字典。

ParameterList用法和list類似

class MyDense(nn.Module):

def __init__(self):

super(MyDense,self).__init__()

self.params = nn.ParameterList([nn.Parameter(torch.randn(4,4)) for i in range(4)])

self.params.append(nn.Parameter(torch.randn(4,1)))

def forward(self,x):

for i in range(len(self.params)):

x = torch.mm(x,self.params[i])

return x

net = MyDense()

print(net)

輸出

MyDense(

(params): ParameterList(

(0): Parameter containing: [torch.FloatTensor of size 4x4]

(1): Parameter containing: [torch.FloatTensor of size 4x4]

(2): Parameter containing: [torch.FloatTensor of size 4x4]

(3): Parameter containing: [torch.FloatTensor of size 4x4]

(4): Parameter containing: [torch.FloatTensor of size 4x1]

)

)

ParameterDict用法和python dict類似.也可以用.keys(),.items()

class MyDictDense(nn.Module):

def __init__(self):

super(MyDictDense, self).__init__()

self.params = nn.ParameterDict({

'linear1': nn.Parameter(torch.randn(4, 4)),

'linear2': nn.Parameter(torch.randn(4, 1))

})

self.params.update({'linear3': nn.Parameter(torch.randn(4, 2))}) # 新增

def forward(self, x, choice='linear1'):

return torch.mm(x, self.params[choice])

net = MyDictDense()

print(net)

print(net.params.keys(),net.params.items())

x = torch.ones(1, 4)

net(x, 'linear1')

輸出

MyDictDense(

(params): ParameterDict(

(linear1): Parameter containing: [torch.FloatTensor of size 4x4]

(linear2): Parameter containing: [torch.FloatTensor of size 4x1]

(linear3): Parameter containing: [torch.FloatTensor of size 4x2]

)

)

odict_keys(['linear1', 'linear2', 'linear3']) odict_items([('linear1', Parameter containing:

tensor([[-0.2275, -1.0434, -1.6733, -1.8101],

[ 1.7530, 0.0729, -0.2314, -1.9430],

[-0.1399, 0.7093, -0.4628, -0.2244],

[-1.6363, 1.2004, 1.4415, -0.1364]], requires_grad=True)), ('linear2', Parameter containing:

tensor([[ 0.5035],

[-0.0171],

[-0.8580],

[-1.1064]], requires_grad=True)), ('linear3', Parameter containing:

tensor([[-1.2078, 0.4364],

[-0.8203, 1.7443],

[-1.7759, 2.1744],

[-0.8799, -0.1479]], requires_grad=True))])

使用自定義的layer構造模型

layer1 = MyDense()

layer2 = MyDictDense()

net = nn.Sequential(layer2,layer1)

print(net)

print(net(x))

輸出

Sequential(

(0): MyDictDense(

(params): ParameterDict(

(linear1): Parameter containing: [torch.FloatTensor of size 4x4]

(linear2): Parameter containing: [torch.FloatTensor of size 4x1]

(linear3): Parameter containing: [torch.FloatTensor of size 4x2]

)

)

(1): MyDense(

(params): ParameterList(

(0): Parameter containing: [torch.FloatTensor of size 4x4]

(1): Parameter containing: [torch.FloatTensor of size 4x4]

(2): Parameter containing: [torch.FloatTensor of size 4x4]

(3): Parameter containing: [torch.FloatTensor of size 4x4]

(4): Parameter containing: [torch.FloatTensor of size 4x1]

)

)

)

tensor([[-4.7566]], grad_fn=)

總結

以上是生活随笔為你收集整理的pytorch自定义新层demo_从头学pytorch(十一):自定义层的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。