PyTorchBCNN/utils.py at master · as595/PyTorchBCNN · GitHub

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
54
55
56
57
58
59
60
61
62
63
64
# -------------------------------------------------------------------------
# History:
# [AMS - 200601] created
# -------------------------------------------------------------------------

import torch
import torch.nn as nn
import torch.nn.functional as F

import numpy as np


def bcnn_loss(c1, c2, f1, y_train, weights, device="cpu"):
    """
        Function to calculate weighted 3 term loss function for BCNN
    """

    if device=="cpu":
      y_c1_train = l1_labels(y_train)
      y_c2_train = l2_labels(y_train)
    else:
      y_c1_train = l1_labels(y_train.to("cpu")).to(device)
      y_c2_train = l2_labels(y_train.to("cpu")).to(device)
      weights = weights.to(device)

    l1 = F.cross_entropy(c1, y_c1_train)
    l2 = F.cross_entropy(c2, y_c2_train)
    l3 = F.cross_entropy(f1, y_train)

    loss = weights[0]*l1 + weights[1]*l2 + weights[2]*l3

    return loss


def l1_labels(labels):

    """
        0: vehicle (0:plane, 1:car, 4:truck)
        1: animal (2:bird, 3:horse)
    """

    l1_labels = np.zeros(labels.size())
    l1_labels[np.where(labels==2)]=1
    l1_labels[np.where(labels==3)]=1

    return torch.tensor(l1_labels, dtype=torch.long)


def l2_labels(labels):

    """
        0: air (0:plane)
        1: ground (1:car, 4:truck)
        2: bird (2:bird)
        3: horse (3:horse)
    """

    l2_labels = np.zeros(labels.size())
    l2_labels[np.where(labels==1)]=1
    l2_labels[np.where(labels==2)]=2
    l2_labels[np.where(labels==3)]=3
    l2_labels[np.where(labels==4)]=1

    return torch.tensor(l2_labels, dtype=torch.long)