|
|
import random |
|
|
import torch |
|
|
import torch.nn.functional as F |
|
|
|
|
|
|
|
|
class CLIPSelf: |
|
|
def __call__(self, batch, model, dist_P_VLM, dist_model, loss, device, cast_dtype, distributed, args): |
|
|
if distributed: |
|
|
model = model.module |
|
|
dist_model = dist_model.module |
|
|
images, normed_boxes, image_crops = batch |
|
|
|
|
|
images = images.to(device=device, dtype=cast_dtype, non_blocking=True) |
|
|
normed_boxes = normed_boxes.to(device=device, dtype=cast_dtype, non_blocking=True) |
|
|
image_crops = image_crops.to(device=device, dtype=cast_dtype, non_blocking=True) |
|
|
|
|
|
if args.multiscale: |
|
|
cur_h, cur_w = images.shape[2:] |
|
|
assert cur_h == cur_w |
|
|
if cur_h == 1024: |
|
|
tar_sizes = [320, 640, 896, 1024] |
|
|
elif cur_h == 896: |
|
|
tar_sizes = [336, 448, 672, 896] |
|
|
else: |
|
|
raise NotImplementedError |
|
|
tar_size = random.choice(tar_sizes) |
|
|
images = F.interpolate(images, size=(tar_size, tar_size), mode='bilinear') |
|
|
|
|
|
rois_list = [] |
|
|
crops_list = [] |
|
|
for bboxes_per_image, crops_per_image in zip(normed_boxes, image_crops): |
|
|
valid = bboxes_per_image[:, -1] > 0.5 |
|
|
rois_list.append(bboxes_per_image[valid, :4]) |
|
|
crops_list.append(crops_per_image[valid]) |
|
|
|
|
|
image_crops = torch.cat(crops_list) |
|
|
with torch.no_grad(): |
|
|
teacher_crop_features = dist_model.encode_image(image_crops, normalize=False) |
|
|
student_roi_features = model.encode_pseudo_boxes(images, rois_list, normalize=False, |
|
|
extract_type=args.extract_type) |
|
|
|
|
|
normed_student_features = F.normalize(student_roi_features, dim=-1) |
|
|
normed_teacher_features = F.normalize(teacher_crop_features, dim=-1) |
|
|
|
|
|
loss_cosine = 1.0 - (normed_student_features * |
|
|
normed_teacher_features).sum(-1).mean() |
|
|
losses = dict(loss_cosine=loss_cosine*args.cosine_weight) |
|
|
|
|
|
return losses, len(images), model.logit_scale.exp() |
|
|
|
