pytorch 加权 MSE 损失

Question

I wanted to apply a weighted MSE to my pytorch model, but I ran into some spots where I do not know how to adapt it correctly. The original lines of code are:

self.mse_criterion = torch.nn.MSELoss(reduction='none')
loss_mot_rec = self.mse_criterion(self.fake_noise, self.real_noise).mean(dim=-1)

def to(self, device):
    if self.opt.is_train:
        self.mse_criterion.to(device)
    self.encoder = self.encoder.to(device)

The function for my weighted MSE loss is:

def weighted_mse_loss(input, target):
    weight=torch.FloatTensor([2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1])
    return (weight * (input - target) ** 2)

So I am confused how to replace the mse_criterion with my function. Any help would be great. The entire original code can be found here Thanks

I tried what @DerekG suggested and now I am getting this error. File "/content/MotionDiffuse/text2motion/trainers/ddpm_trainer.py",

   line 159, in to self.mse_criterion.to(device)
File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 987, in to
   return self._apply(convert)
File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 638, in _apply
   for module in self.children():
File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1792, in children
   for name, module in self.named_children():
File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1811, in named_children
   for name, module in self._modules.items():
File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1265, in __getattr__
   raise AttributeError("'{}' object has no attribute '{}'".format(
AttributeError: 'weighted_MSELoss' object has no attribute '_modules'

Answer 1

Your loss criterion looks fine. Just wrap it in a nn.module and it should be good to use.

class weighted_MSELoss(nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self,inputs,targets,weights):
        return ((input - target)**2 ) * weight

And then call it as the other loss function more or less (I moved the weights outside as I assume you want to generate the weights as some function of the inputs that I am not aware of:

mse_criterion = weighted_MSELoss()
loss_mot_rec = mse_criterion(self.fake_noise,self.real_noise,weights).mean(dim = -1)