| |
| |
| |
|
|
| import pdb |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from nets.sampler import * |
| from nets.repeatability_loss import * |
| from nets.reliability_loss import * |
|
|
|
|
| class MultiLoss(nn.Module): |
| """Combines several loss functions for convenience. |
| *args: [loss weight (float), loss creator, ... ] |
| |
| Example: |
| loss = MultiLoss( 1, MyFirstLoss(), 0.5, MySecondLoss() ) |
| """ |
|
|
| def __init__(self, *args, dbg=()): |
| nn.Module.__init__(self) |
| assert len(args) % 2 == 0, "args must be a list of (float, loss)" |
| self.weights = [] |
| self.losses = nn.ModuleList() |
| for i in range(len(args) // 2): |
| weight = float(args[2 * i + 0]) |
| loss = args[2 * i + 1] |
| assert isinstance(loss, nn.Module), "%s is not a loss!" % loss |
| self.weights.append(weight) |
| self.losses.append(loss) |
|
|
| def forward(self, select=None, **variables): |
| assert not select or all(1 <= n <= len(self.losses) for n in select) |
| d = dict() |
| cum_loss = 0 |
| for num, (weight, loss_func) in enumerate(zip(self.weights, self.losses), 1): |
| if select is not None and num not in select: |
| continue |
| l = loss_func(**{k: v for k, v in variables.items()}) |
| if isinstance(l, tuple): |
| assert len(l) == 2 and isinstance(l[1], dict) |
| else: |
| l = l, {loss_func.name: l} |
| cum_loss = cum_loss + weight * l[0] |
| for key, val in l[1].items(): |
| d["loss_" + key] = float(val) |
| d["loss"] = float(cum_loss) |
| return cum_loss, d |
|
|
