pytorch multilabel classification network not training
I'm trying a simple multi label classification example but the network does not seem to be training correctly as the loss is stagnant.
I've used multilabel_soft_margin_loss as the pytorch docs suggest, but there isn't much else to go on..can't find any proper examples in the docs.
Can anyone peer into this and point out whats wrong with it? Fully working example below (also question on prediction below)
Fully working example code
from __future__ import print_function
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR
from sklearn.datasets import make_multilabel_classification
from torch.utils.data import TensorDataset, DataLoader
from sklearn.model_selection import train_test_split
import xgboost as xgb
from sklearn.metrics import accuracy_score
num_classes = 3
X, y = make_multilabel_classification(n_samples=1000,n_classes=num_classes)
X_tensor, y_tensor = torch.tensor(X), torch.tensor(y)
print(&quot;X Shape :{}&quot;.format(X_tensor.shape))
print(&quot;y Shape :{}&quot;.format(y_tensor.shape))


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(X.shape[1], 300)
        self.fc2 = nn.Linear(300, 10)
        self.fc3 = nn.Linear(10, num_classes)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x
    

device = torch.device(&quot;cpu&quot;) 
lr = 1
batch_size = 128
gamma = 0.9
epochs = 100
args = {'log_interval': 10, 'dry_run':False}
kwargs = {'batch_size': batch_size}
kwargs.update({'num_workers': 1,
               'pin_memory': True,
               'shuffle': True},
                 )

model = Net().to(device)
optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=0.1)    
scheduler = StepLR(optimizer, step_size=1, gamma=gamma)

# data loader
my_dataset = TensorDataset(X_tensor,y_tensor) # create tensor dataset 

train_dataset, test_dataset, = train_test_split(
    my_dataset, test_size=0.2, random_state=42)

train_loader = DataLoader(train_dataset,**kwargs) 
test_loader = DataLoader(test_dataset,**kwargs) 

## Train step ##
for epoch in range(1, epochs + 1):
    model.train() # set model to train
    
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data.float())
        loss = F.multilabel_soft_margin_loss(output,target)
        loss.backward()
        optimizer.step()
        if batch_idx % args['log_interval'] == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))
            if args['dry_run']:
                break
    
    scheduler.step()

Training Loss Progress
Train Epoch: 1 [0/800 (0%)] Loss: 0.694400
Train Epoch: 2 [0/800 (0%)] Loss: 0.697095
Train Epoch: 3 [0/800 (0%)] Loss: 0.705593
Train Epoch: 4 [0/800 (0%)] Loss: 0.651981
Train Epoch: 5 [0/800 (0%)] Loss: 0.704895
Train Epoch: 6 [0/800 (0%)] Loss: 0.650302
Train Epoch: 7 [0/800 (0%)] Loss: 0.658809
Train Epoch: 8 [0/800 (0%)] Loss: 0.904834
Train Epoch: 9 [0/800 (0%)] Loss: 0.655516
Train Epoch: 10 [0/800 (0%)]    Loss: 0.662808
Train Epoch: 11 [0/800 (0%)]    Loss: 0.664752
Train Epoch: 12 [0/800 (0%)]    Loss: 0.656390
Train Epoch: 13 [0/800 (0%)]    Loss: 0.664982
Train Epoch: 14 [0/800 (0%)]    Loss: 0.664430
Train Epoch: 15 [0/800 (0%)]    Loss: 0.664603 # stagnates

On top of that, how would I obtain predictions for this? It's not the same as taking the argmax anymore as its a multi label problem right? (Example output of the network below)
Output
tensor([[ 0.2711,  0.1754, -0.3354],
    [ 0.2711,  0.1754, -0.3354],
    [ 0.2711,  0.1754, -0.3354],
    [ 0.2711,  0.1754, -0.3354],
    [ 0.2711,  0.1754, -0.3354],
    [ 0.2711,  0.1754, -0.3354],
    [ 0.2711,  0.1754, -0.3354]]

Thanks!


On top of that, how would I obtain predictions for this?

If it's a multilabel task and you are outputting logits (as you are) then simply do:
output = model(data.float())
labels = output &gt; 0


point out whats wrong with it?

It is hard and opinionated, what I would do in order:

validate your data. Your neural network response is the same for every input (given your example output is real). Maybe you are passing the same single sample (though seems unlikely as it's sklearn created data)
start simple; no LR scheduler, no weight decay, simple neural network and optimizer only (Adam can stay). Use weight decay if your model is overfitting, it clearly isn't right now.
fix your learning rate; it is one of the most important hyperparameters. 1 is probably too high, start with something like 3e-4 or 1e-3.
try to overfit (loss ~0.0) on small amount of samples (say 32 samples). If you can't, your neural network probably doesn't have enough capacity or there is an error in your code (didn't spot it from quick glance, besides what I've mentioned above). You should verify input and output shapes are correct and returned values manually (it seems for each sample network returns the same logits?).
if you are sure there is no error increase network capacity. Add new hidden layer or two (there is only one) and overfit on single batch. If it's capable go with more data


I've used multilabel_soft_margin_loss as the pytorch docs suggest,

It is the same thing as using torch.nn.BCEWithLogitsLoss which I think is more common, but that's an addendum.
