pytorch自定义新层demo_Pytorch 自定义模型参数及更新

被这个问题困扰了很久，用nn.parameter()定义了参数，但该参数没有更新，.grad() 为none, is_leaf 为False， 其了个怪了。原来是在参数初始化的时候没有正确初始化，我好菜~~~~~。

1) 先看正确的例子

import torch

import torch.nn as nn

class Mask(nn.Module):

def __init__(self):

super(Mask, self).__init__()

self.weight = (torch.nn.Parameter(data=torch.Tensor(1, 1, 1, 1), requires_grad=True))

self.weight.data.uniform_(-1, 1)

print(self.weight.is_leaf)

def forward(self, x):

masked_wt = (self.weight.mul(1)).cuda()

return masked_wt

class Model(nn.Module):

def __init__(self):

super(Model, self).__init__()

self.Mask = Mask()

def forward(self, x):

x = Mask(x)

return x

model = Model()

for name, param in model.named_parameters():

print(name, param)

输出为

True

Mask.weight Parameter containing:

tensor([[[[0.7625]]]], requires_grad=True)