Rani
Rani

Reputation: 503

PyTorch: How to write a neural network that only returns the weights?

I'm training a neural network that learns some weights and based on those weights, I compute transformations that produce the predicted model in combination with the weights. My network doesn't learn properly and therefore I'm writing a different network that does nothing but returning the weights independent from the input x (after normalization with softmax and transpose). This way, I want to find out whether the problem lies in the network or in the transformation estimation outside the network. But this doesn't work. This is what I've got.

class DoNothingNet(torch.nn.Module):
    def __init__(self, n_vertices=6890, n_joints=14):
        super(DoNothingNet, self).__init__()
        self.weights = nn.parameter.Parameter(torch.randn(n_vertices, n_joints))

    def forward(self, x, indices):
        self.weights = F.softmax(self.weights, dim=1)
        return self.weights.transpose(0,1)

But the line self.weights = F.softmax(self.weights, dim=1) doesn't work and produces the error TypeError: cannot assign 'torch.cuda.FloatTensor' as parameter 'weights' (torch.nn.Parameter or None expected). How do I fix this? And does the code even make sense?

Upvotes: 1

Views: 467

Answers (1)

antoleb
antoleb

Reputation: 323

nn.Module tracks all fields of type nn.Parameter for training. In your code every forward call you try to change parameters weights by assigning it to Tensor type, so the error occurs.

The following code outputs normalised weights without changing the stored ones. Hope this will help.

import torch
from torch import nn
from torch.nn import functional as F

class DoNothingNet(torch.nn.Module):
    def __init__(self, n_vertices=6890, n_joints=14):
        super(DoNothingNet, self).__init__()
        self.weights = nn.parameter.Parameter(torch.randn(n_vertices, n_joints))

    def forward(self, x, indices):
        output = F.softmax(self.weights, dim=1)
        return output.transpose(0,1)

Upvotes: 1

Related Questions