|
|
|
|
|
|
|
|
|
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 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
