-
Notifications
You must be signed in to change notification settings - Fork 10
/
loss.py
172 lines (150 loc) · 8.15 KB
/
loss.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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
import torch
import torch.nn as nn
from torch.nn import functional as F
try:
import torch.distributed.nn
from torch import distributed as dist
has_distributed = True
except ImportError:
has_distributed = False
try:
import horovod.torch as hvd
except ImportError:
hvd = None
def gather_features(
image_features,
text_features,
text_features_no,
local_loss=False,
gather_with_grad=False,
rank=0,
world_size=1,
use_horovod=False
):
assert has_distributed, 'torch.distributed did not import correctly, please use a PyTorch version with support.'
if use_horovod:
assert hvd is not None, 'Please install horovod'
if gather_with_grad:
all_image_features = hvd.allgather(image_features)
#all_image_features_no = hvd.allgather(image_features_no)
all_text_features = hvd.allgather(text_features)
all_text_features_no = hvd.allgather(text_features_no)
else:
with torch.no_grad():
all_image_features = hvd.allgather(image_features)
#all_image_features_no = hvd.allgather(image_features_no)
all_text_features = hvd.allgather(text_features)
all_text_features_no = hvd.allgather(text_features_no)
if not local_loss:
# ensure grads for local rank when all_* features don't have a gradient
gathered_image_features = list(all_image_features.chunk(world_size, dim=0))
#gathered_image_features_no = list(all_image_features_no.chunk(world_size, dim=0))
gathered_text_features = list(all_text_features.chunk(world_size, dim=0))
gathered_text_features_no = list(all_text_features_no.chunk(world_size, dim=0))
gathered_image_features[rank] = image_features
#gathered_image_features_no[rank] = image_features_no
gathered_text_features[rank] = text_features
gathered_text_features_no[rank] = text_features_no
all_image_features = torch.cat(gathered_image_features, dim=0)
#all_image_features_no = torch.cat(gathered_image_features_no, dim=0)
all_text_features = torch.cat(gathered_text_features, dim=0)
all_text_features_no = torch.cat(gathered_text_features_no, dim=0)
else:
# We gather tensors from all gpus
if gather_with_grad:
all_image_features = torch.cat(torch.distributed.nn.all_gather(image_features), dim=0)
#all_image_features_no = torch.cat(torch.distributed.nn.all_gather(image_features_no), dim=0)
all_text_features = torch.cat(torch.distributed.nn.all_gather(text_features), dim=0)
all_text_features_no = torch.cat(torch.distributed.nn.all_gather(text_features_no), dim=0)
else:
gathered_image_features = [torch.zeros_like(image_features) for _ in range(world_size)]
#gathered_image_features_no = [torch.zeros_like(image_features_no) for _ in range(world_size)]
gathered_text_features = [torch.zeros_like(text_features) for _ in range(world_size)]
gathered_text_features_no = [torch.zeros_like(text_features_no) for _ in range(world_size)]
dist.all_gather(gathered_image_features, image_features)
#dist.all_gather(gathered_image_features_no, image_features_no)
dist.all_gather(gathered_text_features, text_features)
dist.all_gather(gathered_text_features_no, text_features_no)
if not local_loss:
# ensure grads for local rank when all_* features don't have a gradient
gathered_image_features[rank] = image_features
#gathered_image_features_no[rank] = image_features_no
gathered_text_features[rank] = text_features
gathered_text_features_no[rank] = text_features_no
all_image_features = torch.cat(gathered_image_features, dim=0)
#all_image_features_no = torch.cat(gathered_image_features_no, dim=0)
all_text_features = torch.cat(gathered_text_features, dim=0)
all_text_features_no = torch.cat(gathered_text_features_no, dim=0)
return all_image_features, all_text_features, all_text_features_no
class ClipLoss(nn.Module):
def __init__(
self,
local_loss=False,
gather_with_grad=False,
cache_labels=False,
rank=0,
world_size=1,
use_horovod=False,
):
super().__init__()
self.local_loss = local_loss
self.gather_with_grad = gather_with_grad
self.cache_labels = cache_labels
self.rank = rank
self.world_size = world_size
self.use_horovod = use_horovod
# cache state
self.prev_num_logits = 0
self.labels = {}
self.eyes = {}
def forward(self, image_features, text_features, text_features_no, logit_scale):
device = image_features.device
if self.world_size > 1:
all_image_features, all_text_features, all_text_features_no = gather_features(
image_features, text_features, text_features_no,
self.local_loss, self.gather_with_grad, self.rank, self.world_size, self.use_horovod)
if self.local_loss:
logits_per_image = logit_scale * image_features @ all_text_features.T
logits_per_text = logit_scale * text_features @ all_image_features.T
logits_per_text_no = logit_scale * text_features_no @ all_image_features.T
else:
logits_per_image = logit_scale * all_image_features @ all_text_features.T
logits_per_text = logits_per_image.T
logits_per_image_no = logit_scale * all_image_features @ all_text_features_no.T
#logits_intra_no = logit_scale * all_text_features_no @ all_text_features_no.T
else:
logits_per_image = logit_scale * image_features @ text_features.T
logits_per_text = logit_scale * text_features @ image_features.T
# calculated ground-truth and cache if enabled
num_logits = logits_per_image.shape[0]
if self.prev_num_logits != num_logits or device not in self.labels:
labels = torch.arange(num_logits, device=device, dtype=torch.long)
eyes = torch.eye(num_logits, device=device)
if self.world_size > 1 and self.local_loss:
labels = labels + num_logits * self.rank
if self.cache_labels:
self.labels[device] = labels
self.eyes[device] = eyes
self.prev_num_logits = num_logits
else:
labels = self.labels[device]
eyes = self.eyes[device]
logits_per_image_yes_no = torch.cat([logits_per_image.unsqueeze(-1), logits_per_image_no.unsqueeze(-1)], dim=-1)
logits_per_image_yes_no = F.softmax(logits_per_image_yes_no, dim=-1)
loss_bin_yes, loss_bin_no = self.image_text_binary_opposite_loss(logits_per_image_yes_no, eyes)
loss_tso = self.text_semantic_opposite_loss(all_text_features, all_text_features_no)
return loss_bin_yes, loss_bin_no, loss_tso
def text_semantic_opposite_loss(self, all_text_features, all_text_features_no, mode="L2"):
if mode == "L2":
l2_distance = 2 - 2 * (all_text_features * all_text_features_no).sum(-1) + 1e-4# epsilon = 1e-4, used to get rid of inifity gradient
loss = 2 - torch.sqrt(l2_distance) # \in [0,2]
if mode == "cosine":
loss = (all_text_features * all_text_features_no).sum(-1) + 1.0 # \in [0,2]
return loss.mean()
def image_text_binary_opposite_loss(self, logits_per_image_yes_no, eyes):
N = logits_per_image_yes_no.shape[0]
binary_yes_no = eyes * logits_per_image_yes_no[:,:,0] + (1-eyes) * logits_per_image_yes_no[:,:,1]
loss_bin = - torch.log(binary_yes_no)
loss_bin_yes = (eyes * loss_bin).view(-1).sum() / N
loss_bin_no = ((1-eyes) * loss_bin).view(-1).sum() / (N**2 - N)
return loss_bin_yes, loss_bin_no