Добавление новых параметров для обучения

Question

Недавно я пытался переопределить в PyTorch некоторые статьи, в которых они реализуют новый способ использования ядер: https://www.sciencedirect.com/science/article/abs/pii/S1051200419301873, и до этого проекта я работал только с ядрами и слоями, поддерживаемыми Pytorch. В этом случае я должен добавить новые параметры, которые будут обучаться вместе с остальными обычными слоями, но мне кажется, что я что-то не делаю должным образом, так как параметры не обновляются.

class PointPlaneResnet(nn.Module):
    ''' PointPlaneNet-based encoder network with ResNet blocks. \n
    Args:
        c_dim (int): dimension of latent code c,  defined by config's model.c_dim = 512
        dim (int): input points dimension, in our case 
        hidden_dim (int): hidden dimension of the network
        k (int) : number of neighbours in grouping layer
        channels (int) : number of planes/channels
    '''

    def __init__(self, 
                c_dim=128, 
                dim=3, 
                hidden_dim=128, 
                k = 40,
                channels = 3):

        super().__init__()
        self.c_dim = c_dim
        self.k = k #grouping layer
        self.channels = channels

        hidden_dim = 512
        ##Parameters are Tensor subclasses, that have a very special property when used 
        # with Module s - when they’re 
        # assigned as Module attributes they are automatically added 
        # to the list of its parameters, and will appear e.g. in parameters() iterator. 
        self.plane_weights = torch.nn.Parameter(torch.randn(channels, 4).cuda())

        torch.nn.init.xavier_normal_(self.plane_weights)


        # self.fc_pos = nn.Linear(dim, 2*hidden_dim)
        self.plane_conv = convolution

        self.mlp = MLP(channels = channels)
        self.block_0 = ResnetBlockFC(2*hidden_dim, hidden_dim)
        self.block_1 = ResnetBlockFC(2*hidden_dim, hidden_dim)
        self.block_2 = ResnetBlockFC(2*hidden_dim, hidden_dim)
        self.block_3 = ResnetBlockFC(2*hidden_dim, hidden_dim)
        self.block_4 = ResnetBlockFC(2*hidden_dim, hidden_dim)
        self.fc_c = nn.Linear(hidden_dim, c_dim)

        self.actvn = nn.ReLU()
        self.pool = maxpool

    def forward(self, p):
        batch_size, T, D = p.size()

        # output size: B x T X F
        # net = self.fc_pos(p)

        # print(f'Weight planes {self.plane_weights}')

        net_batch = []
        for i in range(batch_size):
            # print(f'Weight planes {self.plane_weights}')

            net_sample = self.plane_conv(p[i,:,:], self.k, self.plane_weights, self.channels)
            net_batch.append(net_sample)
            # print(f'net_sample is {net_sample}')

        # print(f'net_batch: {net_batch}')
        net = torch.stack(net_batch)
        net = self.mlp(net)


        net = self.block_0(net)
        pooled = self.pool(net, dim=1, keepdim=True).expand(net.size())
        net = torch.cat([net, pooled], dim=2)

        net = self.block_1(net)
        pooled = self.pool(net, dim=1, keepdim=True).expand(net.size())
        net = torch.cat([net, pooled], dim=2)

        net = self.block_2(net)
        pooled = self.pool(net, dim=1, keepdim=True).expand(net.size())
        net = torch.cat([net, pooled], dim=2)

        net = self.block_3(net)
        pooled = self.pool(net, dim=1, keepdim=True).expand(net.size())
        net = torch.cat([net, pooled], dim=2)

        net = self.block_4(net)

        # Recude to  B x F
        net = self.pool(net, dim=1)

        c = self.fc_c(self.actvn(net))

        return c

Как вы можете видеть здесь, я добавил self.plane_weights в качестве новых параметров, которые будут использоваться в функции plane_conv / convolution.

Но когда я начинаю тренироваться, я не вижу обновляемых plane_weights:

from im2mesh.encoder import point_plane_net
new_model = point_plane_net.PointPlaneResnet(k = 5).cuda()

#random input
input = torch.randn(10,10, 3).cuda()
labels = torch.randn(10,128).cuda() * 1000
input.size()

optimizer = optim.SGD(new_model.parameters(), lr=0.001, momentum=0.9)
optimizer.zero_grad()
criterion = nn.BCELoss()

for i in range(5):
    a = list(new_model.parameters())[0]
    print(a)
    outputs = new_model(input)
    outputs.shape

    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()
    b = list(new_model.parameters())[0]
    print(b)

    print(torch.equal(a.data, b.data))

Может ли кто-нибудь увидеть отсюда, что я делаю что-то не так, не должен ли autograd отвечать за обновление весов или я должен переключиться на другую потерю тренировки?