How to create and use PyTorch learnable scalar variables outside of nn.Module?

user_4561314 0
Viewed 799

I am working on a multi-objective problem where I have multiple losses that I need to compute, and the total loss is just the sum of the losses. I want to have PyTorch learnable floating-point parameters alpha, and beta that act as coefficients to the individual losses. Note that the summation of losses occurs outside my NN model in the training loop:

optimizer = AdamW(model.parameters(), lr=2e-5, eps=1e-8)

for batch in dl:

    optimizer.zero_grad()

    result = model(batch)

    loss1 = loss_fn_1(result)
    loss2 = loss_fn_2(result)
    loss3 = loss_fn_3(result)

    loss = alpha*loss1 + beta*loss2 + (1-beta)*loss3 # How to optimize alpha, beta?
                        
    loss.backward()
    optimizer.step()

How would I declare and use the learnable parameters alpha and beta?

pytorch machine-learning
1 Answers

You can put them into a list and add them to an optimizer, for example,

optimizer_for_my_params = torch.Adam([alpha, beta], lr=1e-3)

or separately,

optimizer_alpha = torch.Adam([alpha], lr=1e-3)
optimizer_beta = torch.Adam([beta], lr=1e-3)

and at each step, call zero_grad and step on all optimizers.

Or you can put them in an nn.Module and declare them as parameters:

class MyParams(nn.Module):
  def __init__(self):
    super(MyParams, self).__init__()

    self.alpha = nn.Parameter(torch.tensor(0.))
    self.beta = nn.Parameter(torch.tensor(0.))

  def forward(self, loss1, loss2, loss3):
    loss = self.alpha*loss1 + self.beta*loss2 + (1 - self.beta)*loss3
    return loss

When using it, defining a separate optimizer for the class object should do the job.

Update: Here is a more comprehensive example for the first method.

import torch
import torch.optim as optim

alpha = torch.tensor(0.)
alpha.requires_grad = True
optimizer_alpha = optim.Adam([alpha], lr=1e-3)

print(optimizer_alpha)
# Adam (
# Parameter Group 0
#     amsgrad: False
#     betas: (0.9, 0.999)
#     eps: 1e-08
#     lr: 0.001
#     weight_decay: 0
# )

out = alpha + 1

# test backward()
optimizer_alpha.zero_grad()
out.backward()
print(alpha.grad)
# tensor(1.)

# test step()
optimizer_alpha.step()
print(alpha)
# tensor(-0.0010, requires_grad=True)
user_117906370
Related