自定義layer

不含模型參數(shù)的layer

含模型參數(shù)的layer

核心都一樣,自定義一個(gè)繼承自nn.Module的類,在類的forward函數(shù)里實(shí)現(xiàn)該layer的計(jì)算,不同的是,帶參數(shù)的layer需要用到nn.Parameter

不含模型參數(shù)的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()

含模型參數(shù)的layer

Parameter

ParameterList

ParameterDict

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

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構(gòu)造模型

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=)

總結(jié)

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

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。