Learnable scalar weight in PyTorch

Question

I have two neural networks running in parallel. Each gives a features map of same size say Nx1. Now I want weighted average of these embedding like this w1 * embed1 + w2 * embed2. I have tried these 1 2.But the weights are not updating. Any help would be appreciated. Here is how I am trying to do it:

class LinearWeightedAvg(nn.Module):
      def __init__(self, n_inputs):
        super(LinearWeightedAvg, self).__init__()
        self.weight1 = Variable(torch.randn(1), requires_grad=True).cuda()
        self.weight2 = Variable(torch.randn(1), requires_grad=True).cuda()

    def forward(self, inp_embed):
        return self.weight1 * inp_embed[0] + self.weight2 * inp_embed[1]

class EmbedBranch(nn.Module):
    def __init__(self, feat_dim, embedding_dim):
        super(EmbedBranch, self).__init__()
        fc_layer1 = fc_layer
    def forward(self, x):
        x = self.fc_layer1(x)
        return x

class EmbeddingNetwork(nn.Module):
    def __init__(self, args, N):
        super(EmbeddingNetwork, self).__init__()
        embedding_dim = N

        self.embed1 = EmbedBranch(N, N)
        self.embed2 = EmbedBranch(N, N)
        self.comb_branch = LinearWeightedAvg(metric_dim)
        
        self.args = args
        if args.cuda:
            self.cuda()

    def forward(self, emb1, emb2):
        embeds1 = self.text_branch(emb1)
        embeds2 = self.image_branch(emb2)
        combined = self.comb_branch([embeds1, embeds2])
        return combined

    def train_forward(self, embed1, embed2):

        combined = self(embed1, embed2)

embeds = model.train_forward(embed1, embed2)
loss = loss_func(embeds, labels)
running_loss.update(loss.data.item())
optimizer.zero_grad()
loss.backward()

Also I want the weight to be within 0-1 range.

Thanks,

Answer 1

You should use self.weightx = torch.nn.Parameter(your_inital_tensor) to register a tensor as a learnable parameter of the model.