-
Notifications
You must be signed in to change notification settings - Fork 36
/
discriminative.py
92 lines (85 loc) · 3.23 KB
/
discriminative.py
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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
import torch
import torch.nn as nn
class DiscriminativeLoss(nn.Module):
def __init__(self, delta_d, delta_v,
alpha=1.0, beta=1.0, gamma=0.001,
reduction='mean'):
# TODO: Respect the reduction rule
super(DiscriminativeLoss, self).__init__()
self.alpha = alpha
self.beta = beta
self.gamma = gamma
# Set delta_d > 2 * delta_v
self.delta_d = delta_d
self.delta_v = delta_v
def forward(self, embedded, masks, size):
centroids = self._centroids(embedded, masks, size)
L_v = self._variance(embedded, masks, centroids, size)
L_d = self._distance(centroids, size)
L_r = self._regularization(centroids, size)
loss = self.alpha * L_v + self.beta * L_d + self.gamma * L_r
return loss
def _centroids(self, embedded, masks, size):
batch_size = embedded.size(0)
embedding_size = embedded.size(2)
K = masks.size(2)
x = embedded.unsqueeze(2).expand(-1, -1, K, -1)
masks = masks.unsqueeze(3)
x = x * masks
centroids = []
for i in range(batch_size):
n = size[i]
mu = x[i,:,:n].sum(0) / masks[i,:,:n].sum(0)
if K > n:
m = int(K - n)
filled = torch.zeros(m, embedding_size)
filled = filled.to(embedded.device)
mu = torch.cat([mu, filled], dim=0)
centroids.append(mu)
centroids = torch.stack(centroids)
return centroids
def _variance(self, embedded, masks, centroids, size):
batch_size = embedded.size(0)
num_points = embedded.size(1)
embedding_size = embedded.size(2)
K = masks.size(2)
# Convert input into the same size
mu = centroids.unsqueeze(1).expand(-1, num_points, -1, -1)
x = embedded.unsqueeze(2).expand(-1, -1, K, -1)
# Calculate intra pull force
var = torch.norm(x - mu, 2, dim=3)
var = torch.clamp(var - self.delta_v, min=0.0) ** 2
var = var * masks
loss = 0.0
for i in range(batch_size):
n = size[i]
loss += torch.sum(var[i,:,:n]) / torch.sum(masks[i,:,:n])
return loss
def _distance(self, centroids, size):
batch_size = centroids.size(0)
loss = 0.0
for i in range(batch_size):
n = size[i]
if n <= 1: continue
mu = centroids[i, :n, :]
mu_a = mu.unsqueeze(1).expand(-1, n, -1)
mu_b = mu_a.permute(1, 0, 2)
diff = mu_a - mu_b
norm = torch.norm(diff, 2, dim=2)
margin = 2 * self.delta_d * (1.0 - torch.eye(n))
margin = margin.to(centroids.device)
distance = torch.sum(torch.clamp(margin - norm, min=0.0) ** 2) # hinge loss
distance /= float(n * (n - 1))
loss += distance
loss /= batch_size
return loss
def _regularization(self, centroids, size):
batch_size = centroids.size(0)
loss = 0.0
for i in range(batch_size):
n = size[i]
mu = centroids[i, :n, :]
norm = torch.norm(mu, 2, dim=1)
loss += torch.mean(norm)
loss /= batch_size
return loss