在学习笔记3中我们知道Parameter 也是属于Tensor的一个子类,如果定义了一个Parameter,他会被自动添加到模型的参数中。所以在定义含模型参数的层时,使用Parameter即可,或者是我们也可以使用ParameterList或ParameterDict。下面以ParameterList为例,描述使用方法,ParameterDict的使用方法参见这里。
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(3)])
self.params.append(nn.Parameter(torch.randn(4, 1)))
def forward(self, x, index):
x = torch.mm(x, self.params[index])
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 4x1]
)
)
'''
这样就可以根据索引来访问某个参数,并进行不同的前向传播:
x = torch.ones(1, 4)
print(net(x, 0))
print(net(x, 1))
print(net(x, 2))
'''输出
tensor([[ 0.2612, -0.8207, 2.6069, 3.1468]], grad_fn=<MmBackward>)
tensor([[-4.4122, -0.6476, 3.8580, -2.4974]], grad_fn=<MmBackward>)
tensor([[ 0.6796, -2.4097, -1.8445, 2.7097]], grad_fn=<MmBackward>)
'''