| """ |
| Adapted from salesforce@LAVIS. Below is the original copyright: |
| Copyright (c) 2023, salesforce.com, inc. |
| All rights reserved. |
| SPDX-License-Identifier: BSD-3-Clause |
| For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause |
| """ |
| import contextlib |
| import logging |
| import os |
| import time |
| import datetime |
|
|
| import torch |
| import torch.nn as nn |
| import torch.distributed as dist |
| import torch.nn.functional as F |
|
|
| from .common import dist_utils |
| |
| |
| from .common.logger import MetricLogger |
| |
| |
| |
| |
|
|
|
|
| def disabled_train(self, mode=True): |
| """Overwrite model.train with this function to make sure train/eval mode |
| does not change anymore.""" |
| return self |
|
|
|
|
| class LayerNorm(nn.LayerNorm): |
| """Subclass torch's LayerNorm to handle fp16.""" |
|
|
| def forward(self, x: torch.Tensor): |
| orig_type = x.dtype |
| ret = super().forward(x.type(torch.float32)) |
| return ret.type(orig_type) |
|
|
|
|
| def compute_sim_matrix(model, data_loader, **kwargs): |
| k_test = kwargs.pop("k_test") |
|
|
| metric_logger = MetricLogger(delimiter=" ") |
| header = "Evaluation:" |
|
|
| logging.info("Computing features for evaluation...") |
| start_time = time.time() |
|
|
| texts = data_loader.dataset.text |
| num_text = len(texts) |
| text_bs = 256 |
| text_ids = [] |
| text_embeds = [] |
| text_atts = [] |
| for i in range(0, num_text, text_bs): |
| text = texts[i : min(num_text, i + text_bs)] |
| text_input = model.tokenizer( |
| text, |
| padding="max_length", |
| truncation=True, |
| max_length=35, |
| return_tensors="pt", |
| ).to(model.device) |
| text_feat = model.forward_text(text_input) |
| text_embed = F.normalize(model.text_proj(text_feat)) |
| text_embeds.append(text_embed) |
| text_ids.append(text_input.input_ids) |
| text_atts.append(text_input.attention_mask) |
|
|
| text_embeds = torch.cat(text_embeds, dim=0) |
| text_ids = torch.cat(text_ids, dim=0) |
| text_atts = torch.cat(text_atts, dim=0) |
|
|
| vit_feats = [] |
| image_embeds = [] |
| for samples in data_loader: |
| image = samples["image"] |
|
|
| image = image.to(model.device) |
| image_feat, vit_feat = model.forward_image(image) |
| image_embed = model.vision_proj(image_feat) |
| image_embed = F.normalize(image_embed, dim=-1) |
|
|
| vit_feats.append(vit_feat.cpu()) |
| image_embeds.append(image_embed) |
|
|
| vit_feats = torch.cat(vit_feats, dim=0) |
| image_embeds = torch.cat(image_embeds, dim=0) |
|
|
| sims_matrix = [] |
| for image_embed in image_embeds: |
| sim_q2t = image_embed @ text_embeds.t() |
| sim_i2t, _ = sim_q2t.max(0) |
| sims_matrix.append(sim_i2t) |
| sims_matrix = torch.stack(sims_matrix, dim=0) |
|
|
| score_matrix_i2t = torch.full( |
| (len(data_loader.dataset.image), len(texts)), -100.0 |
| ).to(model.device) |
|
|
| num_tasks = dist_utils.get_world_size() |
| rank = dist_utils.get_rank() |
| step = sims_matrix.size(0) // num_tasks + 1 |
| start = rank * step |
| end = min(sims_matrix.size(0), start + step) |
|
|
| for i, sims in enumerate( |
| metric_logger.log_every(sims_matrix[start:end], 50, header) |
| ): |
| topk_sim, topk_idx = sims.topk(k=k_test, dim=0) |
| image_inputs = vit_feats[start + i].repeat(k_test, 1, 1).to(model.device) |
| score = model.compute_itm( |
| image_inputs=image_inputs, |
| text_ids=text_ids[topk_idx], |
| text_atts=text_atts[topk_idx], |
| ).float() |
| score_matrix_i2t[start + i, topk_idx] = score + topk_sim |
|
|
| sims_matrix = sims_matrix.t() |
| score_matrix_t2i = torch.full( |
| (len(texts), len(data_loader.dataset.image)), -100.0 |
| ).to(model.device) |
|
|
| step = sims_matrix.size(0) // num_tasks + 1 |
| start = rank * step |
| end = min(sims_matrix.size(0), start + step) |
|
|
| for i, sims in enumerate( |
| metric_logger.log_every(sims_matrix[start:end], 50, header) |
| ): |
| topk_sim, topk_idx = sims.topk(k=k_test, dim=0) |
| image_inputs = vit_feats[topk_idx.cpu()].to(model.device) |
| score = model.compute_itm( |
| image_inputs=image_inputs, |
| text_ids=text_ids[start + i].repeat(k_test, 1), |
| text_atts=text_atts[start + i].repeat(k_test, 1), |
| ).float() |
| score_matrix_t2i[start + i, topk_idx] = score + topk_sim |
|
|
| if dist_utils.is_dist_avail_and_initialized(): |
| dist.barrier() |
| torch.distributed.all_reduce( |
| score_matrix_i2t, op=torch.distributed.ReduceOp.SUM |
| ) |
| torch.distributed.all_reduce( |
| score_matrix_t2i, op=torch.distributed.ReduceOp.SUM |
| ) |
|
|
| total_time = time.time() - start_time |
| total_time_str = str(datetime.timedelta(seconds=int(total_time))) |
| logging.info("Evaluation time {}".format(total_time_str)) |
|
|
| return score_matrix_i2t.cpu().numpy(), score_matrix_t2i.cpu().numpy() |
|
|
