Spaces:
Running
Running
File size: 1,746 Bytes
a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb c74a070 a80d6bb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
# Copyright 2019-present NAVER Corp.
# CC BY-NC-SA 3.0
# Available only for non-commercial use
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
|