Loss for Multi-label Classification
I am working on a multi-label classification problem. My gt labels are of shape 14 x 10 x 128, where 14 is the batch_size, 10 is the sequence_length, and 128 is the vector with values 1 if the item in sequence belongs to the object and 0 otherwise.
My output is also of same shape: 14 x 10 x 128. Since, my input sequence was of varying length I had to pad it to make it of fixed length 10. I'm trying to find the loss of the model as follows:
total_loss = 0.0
unpadded_seq_lengths = [3, 4, 5, 7, 9, 3, 2, 8, 5, 3, 5, 7, 7, ...] # true lengths of sequences

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.BCEWithLogitsLoss()

for data in training_dataloader:
    optimizer.zero_grad()

    # shape of input 14 x 10 x 128
    output = model(data)
   
    batch_loss = 0.0
    for batch_idx, sequence in enumerate(output):
        # sequence shape is 10 x 128
        true_seq_len = unpadded_seq_lengths[batch_idx]
        
        # only keep unpadded gt and predicted labels since we don't want loss to be influenced by padded values
        predicted_labels = sequence[:true_seq_len, :] # for example, 3 x 128
        gt_labels =  gt_labels_padded[batch_idx, :true_seq_len, :] # same shape as above, gt_labels_padded has shape 14 x 10 x 128
        
        # loop through unpadded predicted and gt labels and calculate loss  
        for item_idx, predicted_labels_seq_item in enumerate(predicted_labels):
            
            # predicted_labels_seq_item and gt_labels_seq_item are 1D vectors of length 128
            gt_labels_seq_item = gt_labels[item_idx]
            current_loss = criterion(predicted_labels_seq_item, gt_labels_seq_item)                     
            total_loss += current_loss
            batch_loss += current_loss


    batch_loss.backward()
    optimizer.step()


Can anybody please check to see if I'm calculating loss correctly. Thanks
Update:
Is this the correct approach for calculating accuracy metrics?
# batch size: 14
# seq length: 10

for epoch in range(10):
    TP = FP = TN = FN = 0.
    for x, y, mask in tr_dl:
        # mask shape: (10,)
        out = model(x) # out shape: (14, 10, 128)
        y_pred = (torch.sigmoid(out) &gt;= 0.5).float().type(torch.int64) # consider all predictions above 0.5 as 1, rest 0
        y_pred = y_pred[mask] # y_pred shape: (14, 10, 10, 128)
        y_labels = y[mask] # y_labels shape: (14, 10, 10, 128)

        # do I flatten y_pred and y_labels?
        y_pred = y_pred.flatten()
        y_labels = y_labels.flatten()
        
        for idx, prediction in enumerate(y_pred):
            if prediction == 1 and y_labels[idx] == 1:
                # calculate IOU (overlap of prediction and gt bounding box)
                iou = 0.78 # assume we get this iou value for objects at idx
                if iou &gt;= 0.5:
                    TP += 1
                else:
                    FP += 1
            elif prediction == 1 and y_labels[idx] == 0:
                FP += 1
            elif prediction == 0 and y_labels[idx] == 1:
                FN += 1
            else:
                TN += 1
          
     EPOCH_ACC = (TP + TN) / (TP + TN + FP + FN)   



It is usually recommended to stick with batch-wise operations and avoid going into single-element processing steps while in the main training loop. One way to handle this case is to make your dataset return padded inputs and labels with additionally a mask that will come useful for loss computation. In other words, to compute the loss term with sequences of varying sizes, we will use a mask instead of doing individual slices.
Dataset
The way to proceed is to make sure you build the mask in the dataset and not in the inference loop. Here I am showing a minimal implementation that you should be able to transfer to your dataset without much hassle:
class Dataset(data.Dataset):
    def __init__(self):
        super().__init__()

    def __len__(self):
        return 100

    def __getitem__(self, index):
        i = random.randint(5, SEQ_LEN) # for demo puporse, generate x with random length
        x = torch.rand(i, EMB_SIZE)
        y = torch.randint(0, N_CLASSES, (i, EMB_SIZE))

        # pad data to fit in batch
        pad = torch.zeros(SEQ_LEN-len(x), EMB_SIZE)
        x_padded = torch.cat((pad, x))
        y_padded = torch.cat((pad, y))

        # construct tensor to mask loss
        mask = torch.cat((torch.zeros(SEQ_LEN-len(x)), torch.ones(len(x))))

        return x_padded, y_padded, mask

Essentially in the __getitem__, we not only pad the input x and target y with zero values, we also construct a simple mask containing the positions of the padded values in the currently processed element.
Notice how:

x_padded, shaped (SEQ_LEN, EMB_SIZE)
y_padded, shaped (SEQ_LEN, N_CLASSES)
mask, shaped (SEQ_LEN,)

are all three tensors which are shape invariant across the dataset, yet mask contains the padding information necessary for us to compute the loss function appropriately.
Inference
The loss you've used nn.BCEWithLogitsLoss, is the correct one since it's a multi-dimensional loss used for binary classification. In other words, you can use it here in this multi-label classification task, considering each one of the 128 logits as an individual binary prediction. Do not use nn.CrossEntropyLoss) as suggested elsewhere, since the softmax will push a single logit (i.e. class), which is the behaviour required for single-label classification tasks.
Therefore, in the training loop, we simply have to apply the mask to our loss.
for x, y, mask in dl:
   y_pred = model(x)
   loss = mask*bce(y_pred, y)
   # backpropagation, loss postprocessing, logs, etc.

