import copy from typing import Optional, Tuple import torch import torch.nn.functional as F import numpy as np from torch import nn from transformers import OwlViTConfig # from transformers.models.owlvit.modeling_owlvit import OwlViTVisionTransformer class OwlViTBoxPredictionHead(nn.Module): def __init__(self, config: OwlViTConfig): super().__init__() width = config.vision_config.hidden_size self.dense0 = nn.Linear(width, width) self.dense1 = nn.Linear(width, width) self.dense2 = nn.Linear(width, width) self.dense3 = nn.Linear(width, width) self.gelu = nn.GELU() self.dense4 = nn.Linear(width, 4) def forward(self, image_features: torch.Tensor) -> torch.FloatTensor: output = self.dense0(image_features) output = self.gelu(output) output = self.dense1(output) output = self.gelu(output) output = self.dense2(output) output = self.gelu(output) output = self.dense3(output) output = self.gelu(output) output = self.dense4(output) output = self.gelu(output) return output class OwlViTClassPredictionHead(nn.Module): def __init__(self, config: OwlViTConfig): super().__init__() out_dim = config.text_config.hidden_size self.query_dim = config.vision_config.hidden_size self.dense0 = nn.Linear(self.query_dim, out_dim) self.logit_shift = nn.Linear(self.query_dim, 1) self.logit_scale = nn.Linear(self.query_dim, 1) self.elu = nn.ELU() def forward( self, image_embeds: torch.FloatTensor, query_embeds: Optional[torch.FloatTensor], query_mask: Optional[torch.Tensor], ) -> Tuple[torch.FloatTensor]: image_class_embeds = self.dense0(image_embeds) if query_embeds is None: device = image_class_embeds.device batch_size, num_patches = image_class_embeds.shape[:2] pred_logits = torch.zeros((batch_size, num_patches, self.query_dim)).to(device) return (pred_logits, image_class_embeds) # Normalize image and text features image_class_embeds = F.normalize(image_class_embeds, dim=-1) + 1e-6 query_embeds = F.normalize(query_embeds, dim=-1) + 1e-6 # Get class predictions pred_logits = torch.einsum("...pd,...qd->...pq", image_class_embeds, query_embeds) # Apply a learnable shift and scale to logits logit_shift = self.logit_shift(image_embeds) logit_scale = self.logit_scale(image_embeds) logit_scale = self.elu(logit_scale) + 1 pred_logits = (pred_logits + logit_shift) * logit_scale if query_mask is not None: if query_mask.ndim > 1: query_mask = torch.unsqueeze(query_mask, dim=-2) pred_logits = pred_logits.to(torch.float64) pred_logits = torch.where(query_mask == 0, -1e6, pred_logits) pred_logits = pred_logits.to(torch.float32) return (pred_logits, image_class_embeds) class OwlViTPredictionHead(nn.Module): def __init__(self, config: OwlViTConfig, num_classes: int, finetuned: bool): super().__init__() out_dim = config.text_config.hidden_size self.query_dim = config.vision_config.hidden_size self.finetuned = finetuned self.num_classes = num_classes self.mlp_image = nn.Sequential( nn.Flatten(), nn.Linear(in_features=self.query_dim, out_features=self.query_dim), nn.GELU(), nn.Linear(in_features=self.query_dim, out_features=self.query_dim), nn.GELU(), nn.Linear(in_features=self.query_dim, out_features=out_dim), nn.GELU(), ) # if self.finetuned: # self.cls_head = nn.Sequential( # nn.GELU(), # nn.Linear(in_features=out_dim, out_features=out_dim), # nn.GELU() # ) def forward(self, image_embeds: torch.FloatTensor, query_embeds: torch.FloatTensor, topk_idxs: torch.FloatTensor, ) -> Tuple[torch.FloatTensor]: # Get class predictions: topk_idxs (batch_size, n_parts, 1), one_hot (batch_size, n_parts, n_patches*n_patches) topk_idxs = torch.swapaxes(topk_idxs, 1, 2) one_hot = torch.zeros(topk_idxs.shape[0], topk_idxs.shape[1], image_embeds.shape[1]).to(image_embeds.device).scatter_(2, topk_idxs, 1) batch_size, n_parts = one_hot.shape[0], one_hot.shape[1] # (batch_size, n_parts, 3600, 1) * (batch_size, 1, 3600, 1024) = (batch_size, n_parts, 3600, 1024).sum(dim=-2) image_embeds = (one_hot.unsqueeze(-1) * image_embeds.unsqueeze(1)).sum(dim=-2) # image_embeds = self.dense0(image_embeds) # (batch_size, n_patches, 1024) --> (.., .., 768) image_embeds = self.mlp_image(image_embeds.view(-1, image_embeds.shape[-1])).view(batch_size, n_parts, -1) query_embeds = query_embeds.view(batch_size, -1, query_embeds.shape[-1]) # if self.finetuned: # image_embeds = self.cls_head(image_embeds) # query_embeds = query_embeds.view(batch_size, -1, query_embeds.shape[-1]) # Normalize image and text features image_embeds = F.normalize(image_embeds, dim=-1) + 1e-6 # (batch_size, n_parts, 768) query_embeds = F.normalize(query_embeds, dim=-1) + 1e-6 # (batch_size, num_classes * n_parts, 768) # Shape: torch.Size([bs, num_boxes, num_classes * num_parts]) image_text_logits = torch.einsum('bnd, bid -> bni', image_embeds, query_embeds) image_text_logits_reshaped = image_text_logits.view(-1, image_text_logits.shape[-1]) # Shape: (bs, num_classes * num_parts, num_boxes) --> (bs, num_classes, num_parts, num_boxes) pred_logits = image_text_logits.swapaxes(axis0=1, axis1=2).view(batch_size, self.num_classes, n_parts, -1) pred_logits = torch.diagonal(pred_logits, dim1=-2, dim2=-1) # --> torch.Size([bs, num_classes, 12]) #DEBUG: try add sigmoid here to see if it helps. PEIJIE: It does not help. # pred_logits = pred_logits.sigmoid() # pred_logits = abs(pred_logits) # for debugging final_pred_logits = torch.sum(pred_logits, dim=-1) return (image_text_logits_reshaped, final_pred_logits, pred_logits) class OwlViTForClassification(nn.Module): config_class = OwlViTConfig def __init__(self, owlvit_det_model, num_classes, weight_dict, device, freeze_box_heads=False, train_box_heads_only=False, network_type=None, logits_from_teacher=False, finetuned: bool = False, custom_box_head: bool = False): super(OwlViTForClassification, self).__init__() self.config = owlvit_det_model.config self.num_classes = num_classes self.num_parts = 12 self.device = device self.sigmoid = nn.Sigmoid() self.ce_loss = torch.nn.CrossEntropyLoss() # Use CE loss for classification OR only train with contrastive loss self.network_type = network_type self.logits_from_teacher = logits_from_teacher # Initialize OwlViT model from the teacher model self.owlvit = copy.deepcopy(owlvit_det_model.owlvit) self.layer_norm = copy.deepcopy(owlvit_det_model.layer_norm) # For image-level classification self.cls_head = OwlViTPredictionHead(self.config, self.num_classes, finetuned=finetuned) # For box prediction if custom_box_head: self.box_head = OwlViTBoxPredictionHead(self.config) else: self.box_head = copy.deepcopy(owlvit_det_model.box_head) # For box-level classification # Why don't just: # self.class_head = copy.deepcopy(owlvit_det_model.class_head) self.class_head = OwlViTClassPredictionHead(self.config) self.class_head.dense0.load_state_dict(owlvit_det_model.class_head.dense0.state_dict()) self.class_head.logit_shift.load_state_dict(owlvit_det_model.class_head.logit_shift.state_dict()) self.class_head.logit_scale.load_state_dict(owlvit_det_model.class_head.logit_scale.state_dict()) # OwlViT: set equal weights for the bounding box, gIoU and classification losses # self.matcher = DetrHungarianMatcher(class_cost=1, bbox_cost=1, giou_cost=1) # Losses for the criterion in DETR/OwlViT self.weight_dict = weight_dict losses = ["cardinality"] losses += ["boxes"] if weight_dict["loss_bbox"] > 0 else [] losses += ["labels"] if weight_dict["loss_ce"] > 0 else [] self.criterion = DetrLoss( matcher=None, num_parts=self.num_parts, eos_coef=0.1, # Following facebook/detr-resnet-50 losses=losses, ) self.freeze_parameters(freeze_box_heads, train_box_heads_only) del owlvit_det_model def freeze_parameters(self, freeze_box_heads, train_box_heads_only): # OwlViT's text encoder is frozen by default for param in self.owlvit.text_model.parameters(): param.requires_grad = False for param in self.owlvit.text_projection.parameters(): param.requires_grad = False # SKIP finetuning box heads if freeze_box_heads: for param in self.box_head.parameters(): param.requires_grad = False for param in self.class_head.parameters(): param.requires_grad = False # SKIP finetuning vision encoder and MLP head for classification --> Adjust weights of box heads only if train_box_heads_only: for param in self.owlvit.parameters(): param.requires_grad = False for param in self.layer_norm.parameters(): param.requires_grad = False for param in self.cls_head.parameters(): param.requires_grad = False def update_num_classes(self, num_classes): self.num_classes = num_classes self.cls_head.num_classes = num_classes def image_text_embedder(self, input_ids: torch.Tensor, pixel_values: torch.FloatTensor, attention_mask: torch.Tensor, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, ) -> Tuple[torch.FloatTensor]: # Encode text and image outputs = self.owlvit( pixel_values=pixel_values, input_ids=input_ids, attention_mask=attention_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=True, ) # Get image embeddings last_hidden_state = outputs.vision_model_output[0] # 0: last_hidden_state; 1: pooled_output image_embeds = self.owlvit.vision_model.post_layernorm(last_hidden_state) # Resize class token new_size = tuple(np.array(image_embeds.shape) - np.array((0, 1, 0))) class_token_out = torch.broadcast_to(image_embeds[:, :1, :], new_size) # Merge image embedding with class tokens image_embeds = image_embeds[:, 1:, :] * class_token_out image_embeds = self.layer_norm(image_embeds) # Resize to [batch_size, num_patches, num_patches, hidden_size] new_size = ( image_embeds.shape[0], int(np.sqrt(image_embeds.shape[1])), int(np.sqrt(image_embeds.shape[1])), image_embeds.shape[-1], ) image_embeds = image_embeds.reshape(new_size) text_embeds = outputs[-4] return (text_embeds, image_embeds, outputs) def image_embedder( self, pixel_values: torch.FloatTensor ) -> Tuple[torch.FloatTensor]: # Get OwlViTModel vision embeddings (same as CLIP) vision_outputs = self.owlvit.vision_model(pixel_values=pixel_values, return_dict=True) # Apply post_layernorm to last_hidden_state, return non-projected output last_hidden_state = vision_outputs[0] image_embeds = self.owlvit.vision_model.post_layernorm(last_hidden_state) # Resize class token new_size = tuple(np.array(image_embeds.shape) - np.array((0, 1, 0))) class_token_out = torch.broadcast_to(image_embeds[:, :1, :], new_size) # Merge image embedding with class tokens image_embeds = image_embeds[:, 1:, :] * class_token_out image_embeds = self.layer_norm(image_embeds) # Resize to [batch_size, num_patches, num_patches, hidden_size] new_size = ( image_embeds.shape[0], int(np.sqrt(image_embeds.shape[1])), int(np.sqrt(image_embeds.shape[1])), image_embeds.shape[-1], ) image_embeds = image_embeds.reshape(new_size) return (image_embeds, vision_outputs) def normalize_grid_corner_coordinates(self, feature_map: torch.FloatTensor): # Computes normalized xy corner coordinates from feature_map. if not feature_map.ndim == 4: raise ValueError("Expected input shape is [batch_size, num_patches, num_patches, hidden_dim]") device = feature_map.device num_patches = feature_map.shape[1] box_coordinates = np.stack(np.meshgrid(np.arange(1, num_patches + 1), np.arange(1, num_patches + 1)), axis=-1).astype(np.float32) box_coordinates /= np.array([num_patches, num_patches], np.float32) # Flatten (h, w, 2) -> (h*w, 2) box_coordinates = box_coordinates.reshape(box_coordinates.shape[0] * box_coordinates.shape[1], box_coordinates.shape[2]) box_coordinates = torch.from_numpy(box_coordinates).to(device) return box_coordinates def compute_box_bias(self, feature_map: torch.FloatTensor) -> torch.FloatTensor: # The box center is biased to its position on the feature grid box_coordinates = self.normalize_grid_corner_coordinates(feature_map) box_coordinates = torch.clip(box_coordinates, 0.0, 1.0) # Unnormalize xy box_coord_bias = torch.log(box_coordinates + 1e-4) - torch.log1p(-box_coordinates + 1e-4) # The box size is biased to the patch size box_size = torch.full_like(box_coord_bias, 1.0 / feature_map.shape[-2]) box_size_bias = torch.log(box_size + 1e-4) - torch.log1p(-box_size + 1e-4) # Compute box bias box_bias = torch.cat([box_coord_bias, box_size_bias], dim=-1) return box_bias def box_predictor( self, image_feats: torch.FloatTensor, feature_map: torch.FloatTensor, ) -> torch.FloatTensor: """ Args: image_feats: Features extracted from the image, returned by the `image_text_embedder` method. feature_map: A spatial re-arrangement of image_features, also returned by the `image_text_embedder` method. Returns: pred_boxes: List of predicted boxes (cxcywh normalized to 0, 1) nested within a dictionary. """ # Bounding box detection head [batch_size, num_boxes, 4]. pred_boxes = self.box_head(image_feats) # Compute the location of each token on the grid and use it to compute a bias for the bbox prediction pred_boxes += self.compute_box_bias(feature_map) pred_boxes = self.sigmoid(pred_boxes) return pred_boxes def class_predictor( self, image_feats: torch.FloatTensor, query_embeds: Optional[torch.FloatTensor] = None, query_mask: Optional[torch.Tensor] = None, ) -> Tuple[torch.FloatTensor]: """ Args: image_feats: Features extracted from the `image_text_embedder`. query_embeds: Text query embeddings. query_mask: Must be provided with query_embeddings. A mask indicating which query embeddings are valid. """ (pred_logits, image_class_embeds) = self.class_head(image_feats, query_embeds, query_mask) return (pred_logits, image_class_embeds) def _get_text_query_mask(self, text_inputs, text_embeds, batch_size: int): # Embed images and text queries input_ids = text_inputs["input_ids"] # Reshape from [batch_size * max_text_queries, hidden_dim] -> [batch_size, max_text_queries, hidden_dim] max_text_queries = input_ids.shape[0] // batch_size text_embeds = text_embeds.reshape(batch_size, max_text_queries, text_embeds.shape[-1]) # If first token is 0, then this is a padded query [batch_size, num_queries]. input_ids = input_ids.reshape(batch_size, max_text_queries, input_ids.shape[-1]) query_mask = input_ids[..., 0] > 0 return query_mask, text_embeds def forward(self, image_inputs, text_inputs_parts, text_embeds, targets: dict = None): # Store outputs for computing losses loss_dict = {} if not isinstance(image_inputs, torch.Tensor): feature_map, _ = self.image_embedder(pixel_values = image_inputs['pixel_values']) else: feature_map = image_inputs batch_size, num_patches, num_patches, hidden_dim = feature_map.shape image_feats = torch.reshape(feature_map, (batch_size, num_patches * num_patches, hidden_dim)) if self.logits_from_teacher: teacher_boxes_logits = torch.stack([target["logits"] for target in targets], dim=0).to(self.device) topk_scores, topk_idxs = torch.topk(teacher_boxes_logits, k=1, dim=1) else: text_embeds_parts = self.owlvit.get_text_features(**text_inputs_parts) # # Embed images and text queries query_mask, text_embeds_parts = self._get_text_query_mask(text_inputs_parts, text_embeds_parts, batch_size) # Predict object classes [batch_size, num_patches, num_queries+1] pred_logits_parts, class_embeds = self.class_predictor(image_feats, text_embeds_parts, query_mask) # Predict object boxes pred_boxes = self.box_predictor(image_feats, feature_map) # Get the top-1 predictions scores = self.sigmoid(pred_logits_parts) topk_scores, topk_idxs = torch.topk(scores, k=1, dim=1) mapping_indices = [(selected_indices, torch.tensor(list(range(self.num_parts))).to(self.device)) for selected_indices in topk_idxs.squeeze(1)] # get the selected_indexs for mapping_indices selected_idxs = torch.stack([item[0].cpu() for item in mapping_indices]) loss_dict["pred_boxes"] = torch.gather(pred_boxes.cpu(), 1, selected_idxs.unsqueeze(-1).expand(*selected_idxs.shape, 4)) if targets is not None: # ---------------------------------------------------------------------------------------- # Computing box + class + symmetric losses for box selection # ---------------------------------------------------------------------------------------- outputs_loss = {} outputs_loss["logits"] = pred_logits_parts outputs_loss["pred_boxes"] = pred_boxes # Compute box + class losses loss_dict = self.criterion(outputs_loss, targets, mapping_indices) # Compute symmetric loss to get rid of the teacher model logits_per_image = torch.softmax(pred_logits_parts, dim=1) logits_per_text = torch.softmax(pred_logits_parts, dim=-1) # For getting rid of the teacher model if self.weight_dict["loss_sym_box_label"] > 0: sym_loss_box_label = self.loss_symmetric(logits_per_image, logits_per_text, teacher_boxes_logits) loss_dict["loss_sym_box_label"] = sym_loss_box_label # ---------------------------------------------------------------------------------------- # Predict image-level classes (batch_size, num_patches, num_queries) image_text_logits, pred_logits, part_logits = self.cls_head(image_feats, text_embeds, topk_idxs) if self.weight_dict["loss_xclip"] > 0: targets_cls = torch.tensor([target["targets_cls"] for target in targets]).unsqueeze(1).to(self.device) if self.network_type == "classification": one_hot = torch.zeros_like(pred_logits).scatter(1, targets_cls, 1).to(self.device) cls_loss = self.ce_loss(pred_logits, one_hot) loss_dict["loss_xclip"] = cls_loss else: # TODO: Need a linear classifier for this approach # Compute symmetric loss for part-descriptor contrastive learning logits_per_image = torch.softmax(image_text_logits, dim=0) logits_per_text = torch.softmax(image_text_logits, dim=-1) sym_loss = self.loss_symmetric(logits_per_image, logits_per_text, targets_cls) loss_dict["loss_xclip"] = sym_loss return pred_logits, part_logits, loss_dict def loss_symmetric(self, text_logits: torch.Tensor, image_logits: torch.Tensor, targets: torch.Tensor, box_labels: torch.Tensor = None) -> torch.Tensor: # text/image logits (batch_size*num_boxes, num_classes*num_descs): The logits that softmax over text descriptors or boxes # targets (batch_size, 1): The ground truth label of box-text pair for classification OR # targets (batch_size, all_boxes, num_parts): The ground truth label of box-text pair for box selection # box_labels (batch_size, num_boxes), 0 for no box, 1 for box assert text_logits.shape == image_logits.shape # For image classification if image_logits.shape != targets.shape: batch_size = targets.shape[0] # get the matching labels (bs * 12, num_classes * num_parts) default_box_labels = torch.kron(torch.ones(batch_size, self.num_classes), torch.eye(self.num_parts)).to(self.device) if box_labels is None: box_labels = default_box_labels.clone() else: # (batch_size, num_boxes) -> (bs * num_boxes, num_classes * num_parts) box_labels = box_labels.view(-1, 1) * default_box_labels # Create one-hot encoding of targets; matching_labels shape: (bs * 12, num_classes * num_parts) target_one_hot = torch.zeros(batch_size, self.num_classes).to(self.device).scatter(1, targets.view(-1, 1), 1) target_one_hot = torch.kron(target_one_hot, torch.ones(self.num_parts, self.num_parts).to(self.device)) matching_labels = target_one_hot * box_labels else: # For box selection: matching_labels shape: (bs, 576, num_parts) values, indices = torch.max(targets, dim=1) matching_labels = torch.zeros_like(targets).scatter(1, indices.unsqueeze(1), 1) loss_i = F.binary_cross_entropy_with_logits(image_logits, matching_labels, reduction='mean') loss_t = F.binary_cross_entropy_with_logits(text_logits, matching_labels, reduction='mean') sym_loss = (loss_i + loss_t).mean() return sym_loss class DetrLoss(nn.Module): """ This class computes the losses for DetrForObjectDetection/DetrForSegmentation. The process happens in two steps: 1) we compute hungarian assignment between ground truth boxes and the outputs of the model 2) we supervise each pair of matched ground-truth / prediction (supervise class and box). A note on the `num_classes` argument (copied from original repo in detr.py): "the naming of the `num_classes` parameter of the criterion is somewhat misleading. It indeed corresponds to `max_obj_id` + 1, where `max_obj_id` is the maximum id for a class in your dataset. For example, COCO has a `max_obj_id` of 90, so we pass `num_classes` to be 91. As another example, for a dataset that has a single class with `id` 1, you should pass `num_classes` to be 2 (`max_obj_id` + 1). For more details on this, check the following discussion https://github.com/facebookresearch/detr/issues/108#issuecomment-650269223" Args: matcher (`DetrHungarianMatcher`): Module able to compute a matching between targets and proposals. num_parts (`int`): Number of object categories, omitting the special no-object category. eos_coef (`float`): Relative classification weight applied to the no-object category. losses (`List[str]`): List of all the losses to be applied. See `get_loss` for a list of all available losses. """ def __init__(self, matcher, num_parts, eos_coef, losses): super().__init__() self.matcher = matcher self.num_parts = num_parts self.eos_coef = eos_coef self.losses = losses # empty_weight = torch.ones(self.num_parts + 1) empty_weight = torch.ones(self.num_parts) empty_weight[-1] = self.eos_coef self.register_buffer("empty_weight", empty_weight) # removed logging parameter, which was part of the original implementation def loss_labels(self, outputs, targets, indices, num_boxes): """ Classification loss (NLL) targets dicts must contain the key "class_labels" containing a tensor of dim [nb_target_boxes] """ if "logits" not in outputs: raise KeyError("No logits were found in the outputs") source_logits = outputs["logits"] idx = self._get_source_permutation_idx(indices) # target_classes_o = torch.cat([t["class_labels"][J] for t, (_, J) in zip(targets, indices)]) # target_classes = torch.full(source_logits.shape[:2], self.num_parts, dtype=torch.int64, device=source_logits.device) # target_classes[idx] = target_classes_o source_logits = source_logits[idx].view(len(indices), -1, self.num_parts) target_classes = torch.stack([t["class_labels"][J] for t, (_, J) in zip(targets, indices)], dim=0) loss_ce = nn.functional.cross_entropy(source_logits.transpose(1, 2), target_classes, self.empty_weight) losses = {"loss_ce": loss_ce} return losses @torch.no_grad() def loss_cardinality(self, outputs, targets, indices, num_boxes): """ Compute the cardinality error, i.e. the absolute error in the number of predicted non-empty boxes. This is not really a loss, it is intended for logging purposes only. It doesn't propagate gradients. """ logits = outputs["logits"] device = logits.device target_lengths = torch.as_tensor([len(v["class_labels"]) for v in targets], device=device) # Count the number of predictions that are NOT "no-object" (which is the last class) card_pred = (logits.argmax(-1) != logits.shape[-1] - 1).sum(1) card_err = nn.functional.l1_loss(card_pred.float(), target_lengths.float()) losses = {"cardinality_error": card_err} return losses def loss_boxes(self, outputs, targets, indices, num_boxes): """ Compute the losses related to the bounding boxes, the L1 regression loss and the GIoU loss. Targets dicts must contain the key "boxes" containing a tensor of dim [nb_target_boxes, 4]. The target boxes are expected in format (center_x, center_y, w, h), normalized by the image size. """ if "pred_boxes" not in outputs: raise KeyError("No predicted boxes found in outputs") idx = self._get_source_permutation_idx(indices) source_boxes = outputs["pred_boxes"][idx] target_boxes = torch.cat([t["boxes"][i] for t, (_, i) in zip(targets, indices)], dim=0) losses = {} loss_bbox = nn.functional.l1_loss(source_boxes, target_boxes, reduction="none") losses["loss_bbox"] = loss_bbox.sum() / num_boxes loss_giou = 1 - torch.diag(generalized_box_iou(center_to_corners_format(source_boxes), center_to_corners_format(target_boxes))) losses["loss_giou"] = loss_giou.sum() / num_boxes return losses def loss_masks(self, outputs, targets, indices, num_boxes): """ Compute the losses related to the masks: the focal loss and the dice loss. Targets dicts must contain the key "masks" containing a tensor of dim [nb_target_boxes, h, w]. """ if "pred_masks" not in outputs: raise KeyError("No predicted masks found in outputs") source_idx = self._get_source_permutation_idx(indices) target_idx = self._get_target_permutation_idx(indices) source_masks = outputs["pred_masks"] source_masks = source_masks[source_idx] masks = [t["masks"] for t in targets] # TODO use valid to mask invalid areas due to padding in loss target_masks, valid = nested_tensor_from_tensor_list(masks).decompose() target_masks = target_masks.to(source_masks) target_masks = target_masks[target_idx] # upsample predictions to the target size source_masks = nn.functional.interpolate( source_masks[:, None], size=target_masks.shape[-2:], mode="bilinear", align_corners=False ) source_masks = source_masks[:, 0].flatten(1) target_masks = target_masks.flatten(1) target_masks = target_masks.view(source_masks.shape) losses = { "loss_mask": sigmoid_focal_loss(source_masks, target_masks, num_boxes), "loss_dice": dice_loss(source_masks, target_masks, num_boxes), } return losses def _get_source_permutation_idx(self, indices): # permute predictions following indices batch_idx = torch.cat([torch.full_like(source, i) for i, (source, _) in enumerate(indices)]) source_idx = torch.cat([source for (source, _) in indices]) return batch_idx, source_idx def _get_target_permutation_idx(self, indices): # permute targets following indices batch_idx = torch.cat([torch.full_like(target, i) for i, (_, target) in enumerate(indices)]) target_idx = torch.cat([target for (_, target) in indices]) return batch_idx, target_idx def get_loss(self, loss, outputs, targets, indices, num_boxes): loss_map = { "labels": self.loss_labels, "cardinality": self.loss_cardinality, "boxes": self.loss_boxes, "masks": self.loss_masks, } if loss not in loss_map: raise ValueError(f"Loss {loss} not supported") return loss_map[loss](outputs, targets, indices, num_boxes) def forward(self, outputs, targets, indices): """ This performs the loss computation. Args: outputs (`dict`, *optional*): Dictionary of tensors, see the output specification of the model for the format. targets (`List[dict]`, *optional*): List of dicts, such that `len(targets) == batch_size`. The expected keys in each dict depends on the losses applied, see each loss' doc. """ outputs_without_aux = {k: v for k, v in outputs.items() if k != "auxiliary_outputs"} # ThangPM: Do NOT use bipartite matching --> Use the boxes selected by argmax for computing symmetric loss # Retrieve the matching between the outputs of the last layer and the targets # indices = self.matcher(outputs_without_aux, targets) # Compute the average number of target boxes across all nodes, for normalization purposes num_boxes = sum(len(t["class_labels"]) for t in targets) num_boxes = torch.as_tensor([num_boxes], dtype=torch.float, device=next(iter(outputs.values())).device) # (Niels): comment out function below, distributed training to be added # if is_dist_avail_and_initialized(): # torch.distributed.all_reduce(num_boxes) # (Niels) in original implementation, num_boxes is divided by get_world_size() num_boxes = torch.clamp(num_boxes, min=1).item() # Compute all the requested losses losses = {} for loss in self.losses: losses.update(self.get_loss(loss, outputs, targets, indices, num_boxes)) # In case of auxiliary losses, we repeat this process with the output of each intermediate layer. if "auxiliary_outputs" in outputs: for i, auxiliary_outputs in enumerate(outputs["auxiliary_outputs"]): # indices = self.matcher(auxiliary_outputs, targets) for loss in self.losses: if loss == "masks": # Intermediate masks losses are too costly to compute, we ignore them. continue l_dict = self.get_loss(loss, auxiliary_outputs, targets, indices, num_boxes) l_dict = {k + f"_{i}": v for k, v in l_dict.items()} losses.update(l_dict) return losses