I am creating a custom loss function that uses two metric
def loss_fn(preds,target,a,b) :
loss1 = loss_func1(preds,target)
loss2 = loss_func2(preds,target)
loss = a*loss1 + b*loss2
How do i make variable a and b, in the above function, learnable in the training loop
You need to make two
nn.Parameters for a and b (on the device
nn.Paramter(torch.tensor(1.0, device=my_device))) or so and pass those to the optimizer (e.g.
list(model.parameters()) + [a, b] instead of just the parameters).
Note that generally it is mathematically difficult to train loss weights, but whether this applies in your case is hard to tell at this level of detail.
so i just initialize the variables in the constructor of my network class and use them as netobject.a or do i have to write something in the backward function also?
Yeah, you could assign them to your own Loss module or something.
No, backward will be automatically.