pytorch 中的加权 mse 损失

Question

def weighted_mse_loss(input_tensor, target_tensor, weight = 1):
    observation_dim = input_tensor.size()[-1]
    streched_tensor = ((input_tensor - target_tensor) ** 2).view(-1, observation_dim)
    entry_num = float(streched_tensor.size())[0]
    non_zero_entry_num = torch.sum(streched_tensor[:,0] != 0).float()
    weighted_tensor = torch.mm(
        ((input_tensor - target_tensor)**2).view(-1, observation_dim),
        (torch.diag(weight.float().view(-1)))
    )
    return torch.mean(weighted_tensor) * weight.nelement() * entry_num / non_zero_entry_num

I can't understand how the code gives weighted Mean Square Error loss. I get that observation_dim is the final output dimension, (the class number I guess), and after that line, I don't get it. Could someone help me figure out how the code calculates the loss?

Thanks a lot.

Answer 1

    def weighted_mse_loss(input, target, weight):
        return (weight * (input - target) ** 2).mean()

try this, hope this can help. All arguments need tensored.