Apply math operation dynamically between two nn.Modules or loss functions

Question

I would like to apply math operations dynamically between two loss functions or nn.Modules or python objects. It could be also treated as a problem to generate dynamic graphs in pytorch.

For example: In the below example, I would like to add two loss functions.

nn.L1Loss() + nn.CosineEmbeddingLoss()

If I do this, it gives me an error:

----> 1 nn.L1Loss() + nn.CosineEmbeddingLoss()
TypeError: unsupported operand type(s) for +: 'L1Loss' and 'CosineEmbeddingLoss'

I also tried creating a wrapper with forward function and torch operations like below, but it doesn’t work either. In the below case x and y can be any loss functions and op can be any math operation like addition, and subtraction, and so on.

class Execute_Op(nn.Module):
    def __init__(self):
        super().__init__()
        
    def forward(self, x, y, op):
        if op == 'add':
            return torch.add(x, y)
        elif op == 'subtract':
            return torch.subtract(x - y)

exec_op = Execute_Op()
exec_op(nn.L1Loss(), nn.CosineEmbeddingLoss(), 'add')

It gives error like the below:

Execute_Op.forward(self, x, y, op)
      5 def forward(self, x, y, op):
      6     if op == 'add':
----> 7         return torch.add(x, y)
      8     elif op == 'subtract':
      9         return torch.subtract(x - y)

TypeError: add(): argument 'input' (position 1) must be Tensor, not L1Loss

I am aware of functional APIs and the general way to pass truth values and predicted values to the loss function. But in that case, I cannot combine loss functions dynamically at run time.

I am not sure how exactly to implement it. But any help is really appreciated. Also, if there is a pythonic way or Pytorch way to do this, it would be great.

Edited:

I would like to call this function/class recursively.

user_2790047 · Answer

The problem here is that loss functions don't have a standard calling convention. Different loss functions take different arguments with different semantics. For example nn.L1Loss() takes two arguments: an input and a target of the same shape. On the other hand nn.CosineEmbeddingLoss() takes three arguments: two inputs of the same shape and a target of a different shape.

This means that the calling convention for ExecuteOp would depend on the specific loss functions chosen. For this reason I don't think its really a good idea to make it general with respect to the loss functions.

A simple solution without the additional level of generalization would be

class ExecuteOp(nn.Module):
    def init(self):
        super().__init__()
        self.l1_loss = nn.L1Loss()
        self.emb_loss = nn.CosineEmbeddingLoss()

    def forward(self, op, input_l1, target_l1, input1_emb, input2_emb, target_emb):
        assert op in ('add', 'subtract')
        loss1 = self.l1_loss(input_l1, target_l1)
        loss2 = self.emb_loss(input1_emb, input2_emb, target_emb)
        if op == 'add':
            return loss1 + loss2
        return loss1 - loss2

Then you could use this as follows:

exec_op = ExecuteOp()

...

# at training/val time, assuming inputs and targets are provided from your model and dataloader
loss = exec_op('add', input_l1, target_l1, input1_emb, input2_emb, target_emb)

Consider how you would even call exec_op at runtime if it were general w.r.t. loss functions. What arguments would you provide it, and in what order?

Apply math operation dynamically between two nn.Modules or loss functions

1 Answers