Is it possible to vectorize MSE for multiple outputs

1

I have a regression model with multiple outputs, so I have a custom mse function

def custom_mse(predicted, target):
    total_mse = 0
    for i in range(target.shape[1]):
        total_mse+=nn.MSELoss()(predicted[i], target[i])
    return total_mse

the purpose is to add up all the mse for each individual output, and return the total as the loss. This works well enough, but is there a way to get rid of the for loop and vectorize the total_mse calculation?

2

nn.MSELoss accepts batches of your model output and targets.

I’m a bit confused about the indexing. You are iterating using target.shape[1], but are indexing both tensors in dim0. Is this a typo?

3

The data has target.shape = (n_samples, k) where k > 1. I’m trying to have a model predict the soft pseudolabels generated by a teacher model, so the targets are the predicted logits for each class, and I’m using the sum of MSE for each class as the loss function

I’m a bit confused about the indexing. You are iterating using target.shape[1] , but are indexing both tensors in dim0. Is this a typo?

Yes, that appears to be an error in the code, it should be indexing [:,i]

I did some further digging, it looks like I had an error in my original implementation of nn.MSELoss. Fixing that seemed to give me the results I’m looking for.