Pytorch training with unbalanced class sizes while validation with balanced classes isn't working

pygene: <p>I’ve been coding a pytorch CNN on a binary classification task where my dataset is unbalanced – the ratio of the two classes is 30:1. Within the training loop, every 5 epochs, I calculate the loss on a small set of validation data to display along with the loss on the much larger training set to monitor progress and help tune the hyperparameters.</p>
<p>At first, to keep things simple, I only used a subset of the training data so that the two classes would be balanced. That was working and I saw some moderate improvements in loss in both training and validation before overfitting set in.</p>
<p>In order to improve training, I then tried to use all of the training data and switched from the original loss function <code>torch.nn.BCELoss</code> to <code>torch.nn.BCEWithLogitsLoss(pos_weight=30.0)</code>, keeping in mind to drop the final <code>torch.sigmoid</code> transform. Since the validation data continued to be balanced, I used <code>torch.nn.BCEWithLogitsLoss(pos_weight=1.0)</code> to calculate the validation loss.</p>
<p>With this setup, although my training loss decreases right away, my validation loss only goes up.</p>
<p>Here is the relevant code:</p>
<pre><code class="lang-auto">def train(epochs, optimizer, model, train_loss_fn, train_loader, test_loss_fn, test_loader):
    for epoch in range(1, epochs + 1):
        train_losses = []
        model.train()

        for images, classes in train_loader:
            optimizer.zero_grad()
            images = images.to(device)
            classes = classes.to(device)
            outputs = model(images)

            loss = train_loss_fn(torch.flatten(outputs), classes)
            loss.backward()
            optimizer.step()
            train_losses.append( loss.item() )

        if (epoch == 1) | (epoch % 5 == 0):
            now = datetime.datetime.now()
            test_loss = test(model, test_loader, test_loss_fn)
            test_loss = np.round(test_loss, 4)
            train_loss = np.round(np.mean(np.asarray(train_losses)),4)
            lr = np.format_float_scientific(np.squeeze(get_lr(optimizer)), precision=4)
            print(f"{now}: Epoch {epoch}, Train Loss: {train_loss}, Test Loss: {test_loss}, Learning Rate: {lr}")

def test(model, test_loader, lf):
    model.eval()
    test_losses = []
    with torch.no_grad():
        for i, (images,classes) in enumerate(test_loader):
            images = images.to(device)
            classes = classes.to(device)
            outputs = model(images)

            loss = lf(torch.flatten(outputs), torch.flatten(classes))
            test_losses.append( loss.item() )

    mean_loss = np.mean(np.asarray(test_losses))
    return mean_loss


Class_1_Weighting = torch.tensor(30.)
TRAIN_loss_fn = torch.nn.BCEWithLogitsLoss(reduction='mean', pos_weight=Class_1_Weighting)
TEST_loss_fn = torch.nn.BCEWithLogitsLoss(reduction='mean')

# train(...)
</code></pre>
<p><em>cross-posted from stackexchange</em></p>

hca: <p><a class="mention" href="/u/pygene">@pygene</a> did you find a solution for this issue? I am observing the same behavior with my training code and I am also using torch.nn.BCEWithLogitsLoss(pos_weight =W) for my unbalanced training data.</p>