Weighted loss during ensemble

username: tomcc
cooked: <p>I am ensembing two models with mean pooling but also want to weight the loss of each seperate model at the same time so the less accurate model will contribute less to the final prediction.</p>
<p>Here is a simple example of what I am trying to achieve.</p>
<pre><code class="lang-auto">class WeightLoss(nn.Module):
    def __init__(self, n):
        super().__init__()
        self.n = n
        self.param = nn.Parameter(torch.empty(n), requires_grad=True)

        self.reset_parameters()

    def reset_parameters(self):
        nn.init.constant_(self.param, 1 / self.n)

    def forward(self, x):
        # make sure the params sum to 1
        return torch.mul(x, F.softmax(self.param))


class Net2(nn.Module):
    def __init__(self, modelA, modelB):
        super(Net2, self).__init__()
        self.modelA = modelA
        self.modelB = modelB

        self.fc1 = nn.Linear(10, 8)
        self.fc2 = nn.Linear(10, 8)

        self.weight_loss = WeightLoss(2)
    
    def loss(self, x, y):
        return F.cross_entropy(x, y)
    
    def forward(self, x, y):
        # individual outputs
        output1 = self.fc1(self.modelA(x))
        output2 = self.fc2(self.modelB(x))
        # mean pooling
        final_output = torch.mean(torch.stack([output1, output2]), dim=0)
        final_loss = self.loss(final_output, y)
        # weighted loss for individual models
        loss1 = self.loss(output1, y)
        loss2 = self.loss(output2, y)
        weighted_loss = self.weight_loss(torch.stack([loss1, loss2]))
        weighted_loss = torch.sum(weighted_loss)

        # combine the 2 losses
        final_loss = torch.add(final_loss, weighted_loss)

        return final_loss
</code></pre>
<p>With the following architecture, <code>final_loss</code> do decrease over time but the weights for weighing the losses for each model do not change over time or by very small margin.</p>
<p>Have I done anything inappropriate?</p>

username: ptrblck
cooked: <p>Your code seems to work generally and reduces the weight for the larger loss to decrease the overall global loss:</p>
<pre><code class="lang-python">class WeightLoss(nn.Module):
    def __init__(self, n):
        super().__init__()
        self.n = n
        self.param = nn.Parameter(torch.empty(n), requires_grad=True)

        self.reset_parameters()

    def reset_parameters(self):
        nn.init.constant_(self.param, 1 / self.n)

    def forward(self, x):
        # make sure the params sum to 1
        return torch.mul(x, F.softmax(self.param, dim=0))


class Net2(nn.Module):
    def __init__(self, modelA, modelB):
        super(Net2, self).__init__()
        self.modelA = modelA
        self.modelB = modelB

        self.fc1 = nn.Linear(10, 8)
        self.fc2 = nn.Linear(10, 8)

        self.weight_loss = WeightLoss(2)
    
    def loss(self, x, y):
        return F.cross_entropy(x, y)
    
    def forward(self, x, y):
        # individual outputs
        output1 = self.fc1(self.modelA(x))
        output2 = self.fc2(self.modelB(x))
        # mean pooling
        final_output = torch.mean(torch.stack([output1, output2]), dim=0)
        final_loss = self.loss(final_output, y)
        # weighted loss for individual models
        loss1 = self.loss(output1, y)
        loss2 = self.loss(output2, y)
        print("loss1 {}, loss2 {}".format(loss1.item(), loss2.item()))
        weighted_loss = self.weight_loss(torch.stack([loss1, loss2]))
        weighted_loss = torch.sum(weighted_loss)

        # combine the 2 losses
        final_loss = torch.add(final_loss, weighted_loss)

        return final_loss
    
model = Net2(nn.Linear(10, 10), nn.Linear(10, 10))
x = torch.randn(8, 10)
target = torch.randint(0, 8, (8,))
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(1000):
    optimizer.zero_grad()
    loss = model(x, target)
    loss.backward()
    optimizer.step()
    print("epoch {}, loss {}, weights {}".format(
        epoch, loss.item(), model.weight_loss.param))
</code></pre>
<p>Output:</p>
<pre><code class="lang-python">...
loss1 0.017472071573138237, loss2 0.008564327843487263
epoch 998, loss 0.020118102431297302, weights Parameter containing:
tensor([-0.0558,  1.0558], requires_grad=True)
loss1 0.017431674525141716, loss2 0.00854448415338993
epoch 999, loss 0.020070992410182953, weights Parameter containing:
tensor([-0.0559,  1.0559], requires_grad=True)
</code></pre>
<p>I don’t know if calculating the <code>final_loss</code> first and adding the <code>weighted_loss</code> is beneficial or not.</p>