File size: 1,746 Bytes
404d2af
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8b973ee
 
404d2af
8b973ee
404d2af
 
 
8b973ee
404d2af
 
8b973ee
404d2af
 
8b973ee
 
 
404d2af
 
 
 
 
8b973ee
404d2af
 
8b973ee
 
 
 
404d2af
 
 
8b973ee
404d2af
8b973ee
 
 
404d2af
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