Spaces:
Runtime error
Runtime error
File size: 2,524 Bytes
be13417 |
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 |
'''
* Copyright (c) 2023, salesforce.com, inc.
* All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
* For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
* By Le Xue
'''
import torch
import torch.nn as nn
import torch.nn.functional as F
from lavis.models.ulip_models.utils import utils
class ULIPWithImageLoss(nn.Module):
def __init__(self):
super().__init__()
self.labels = None
self.last_local_batch_size = None
def forward(self, outputs):
pc_embed = outputs['pc_embed']
text_embed = outputs['text_embed']
image_embed = outputs['image_embed']
logit_scale = outputs['logit_scale']
local_batch_size = pc_embed.size(0)
if local_batch_size != self.last_local_batch_size:
self.labels = local_batch_size * utils.get_rank() + torch.arange(
local_batch_size, device=pc_embed.device
)
self.last_local_batch_size = local_batch_size
# normalized features
pc_embed = F.normalize(pc_embed, dim=-1, p=2)
text_embed = F.normalize(text_embed, dim=-1, p=2)
image_embed = F.normalize(image_embed, dim=-1, p=2)
# gather features from all GPUs
pc_embed_all, text_embed_all, image_embed_all = \
utils.all_gather_batch([pc_embed, text_embed, image_embed])
# cosine similarity as logits
logits_per_pc_text = logit_scale * pc_embed @ text_embed_all.t()
logits_per_text_pc = logit_scale * text_embed @ pc_embed_all.t()
logits_per_pc_image = logit_scale * pc_embed @ image_embed_all.t()
logits_per_image_pc = logit_scale * image_embed @ pc_embed_all.t()
loss = (F.cross_entropy(logits_per_pc_text, self.labels) + \
F.cross_entropy(logits_per_text_pc, self.labels)) / 2 + \
(F.cross_entropy(logits_per_pc_image, self.labels) + F.cross_entropy(logits_per_image_pc, self.labels)) / 2
# compute accuracy
with torch.no_grad():
pred = torch.argmax(logits_per_pc_text, dim=-1)
correct = pred.eq(self.labels).sum()
pc_text_acc = 100 * correct / local_batch_size
pred = torch.argmax(logits_per_pc_image, dim=-1)
correct = pred.eq(self.labels).sum()
pc_image_acc = 100 * correct / local_batch_size
return {'loss': loss, 'ulip_loss': loss, 'ulip_pc_image_acc': pc_image_acc, 'ulip_pc_text_acc': pc_text_acc}
|