jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def prepare_mask_latents(self, mask, masked_image, batch_size, height,
                             width, dtype, device, generator,
                             do_classifier_free_guidance):
        # resize the mask to latents shape as we concatenate the
        # mask to the latents
        # we do that before converting to dtype to avoid breaking in
        # case we're using cpu_offload
        # and half precision
        mask = t.nn.functional.interpolate(
            mask,
            size=(height // self.vae_scale_factor,
                  width // self.vae_scale_factor))
        mask = mask.to(device=device, dtype=dtype)

        masked_image = masked_image.to(device=device, dtype=dtype)
        masked_image_latents = self._encode_vae_image(
            masked_image, generator=generator)

        # duplicate mask and masked_image_latents for each generation per
        # prompt, using mps friendly method
        if mask.shape[0] < batch_size:
            if not batch_size % mask.shape[0] == 0:
                raise ValueError(
                    "The passed mask and the required batch size don't match. "
                    ' Masks are supposed to be duplicated to'
                    f' a total batch size of {batch_size}, '
                    f' but {mask.shape[0]} '
                    ' masks were passed. Make sure the number'
                    ' of masks that you pass is divisible by the '
                    ' total requested batch size.')
            mask = mask.repeat(batch_size // mask.shape[0], 1, 1, 1)
        if masked_image_latents.shape[0] < batch_size:
            if not batch_size % masked_image_latents.shape[0] == 0:
                raise ValueError(
                    "The passed images and the required batch size don't "
                    ' match. Images are supposed to be duplicated'
                    f' to a total batch size of {batch_size}, '
                    f' but {masked_image_latents.shape[0]} images were passed.'
                    ' Make sure the number of images that you pass is '
                    ' divisible by the total requested batch size.')
            masked_image_latents = masked_image_latents.repeat(
                batch_size // masked_image_latents.shape[0], 1, 1, 1)

        mask = torch.cat([mask] * 2) if do_classifier_free_guidance else mask
        masked_image_latents = (
            torch.cat([masked_image_latents] * 2)
            if do_classifier_free_guidance else masked_image_latents)

        # aligning device to prevent device errors when concatenating it
        # with the latent model input
        masked_image_latents = masked_image_latents.to(
            device=device, dtype=dtype)
        return mask, masked_image_latents
//  [python]
def prepare_mask_latents(self, mask, masked_image, batch_size, height,
                             width, dtype, device, generator,
                             do_classifier_free_guidance):
        # resize the mask to latents shape as we concatenate the
        # mask to the latents
        # we do that before converting to dtype to avoid breaking in
        # case we're using cpu_offload
        # and half precision
        mask = torch.nn.functional.interpolate(
            mask,
            size=(height // self.vae_scale_factor,
                  width // self.vae_scale_factor))
        mask = mask.to(device=device, dtype=dtype)

        masked_image = masked_image.to(device=device, dtype=dtype)
        masked_image_latents = self._encode_vae_image(
            masked_image, generator=generator)

        # duplicate mask and masked_image_latents for each generation per
        # prompt, using mps friendly method
        if mask.shape[0] < batch_size:
            if not batch_size % mask.shape[0] == 0:
                raise ValueError(
                    "The passed mask and the required batch size don't match. "
                    ' Masks are supposed to be duplicated to'
                    f' a total batch size of {batch_size}, '
                    f' but {mask.shape[0]} '
                    ' masks were passed. Make sure the number'
                    ' of masks that you pass is divisible by the '
                    ' total requested batch size.')
            mask = mask.repeat(batch_size // mask.shape[0], 1, 1, 1)
        if masked_image_latents.shape[0] < batch_size:
            if not batch_size % masked_image_latents.shape[0] == 0:
                raise ValueError(
                    "The passed images and the required batch size don't "
                    ' match. Images are supposed to be duplicated'
                    f' to a total batch size of {batch_size}, '
                    f' but {masked_image_latents.shape[0]} images were passed.'
                    ' Make sure the number of images that you pass is '
                    ' divisible by the total requested batch size.')
            masked_image_latents = masked_image_latents.repeat(
                batch_size // masked_image_latents.shape[0], 1, 1, 1)

        mask = torch.cat([mask] * 2) if do_classifier_free_guidance else mask
        masked_image_latents = (
            torch.cat([masked_image_latents] * 2)
            if do_classifier_free_guidance else masked_image_latents)

        # aligning device to prevent device errors when concatenating it
        # with the latent model input
        masked_image_latents = masked_image_latents.to(
            device=device, dtype=dtype)
        return mask, masked_image_latents
//  [python]
def forward(self, x):
    res  = self.res(x)
    down = self.down(res)
    mid  = self.mid(down)
    up   = self.up(mid)
    up   = torch.nn.functional.upsample(up, size=[res.size(2), res.size(3)], mode='bilinear')
    return res + up
//  [python]
def forward(self, x):
    res  = self.res(x)
    down = self.down(res)
    mid  = self.mid(down)
    up   = self.up(mid)
    up   = torch.nn.functional.interpolate(up, [res.size(2), res.size(3)], mode='bilinear', align_corners=True)
    return res + up
//  [python]
def preprocess(self, input: Dict[str, Any]) -> Dict[str, Any]:
        img = LoadImage.convert_to_img(input.get('img'))
        test_transforms = transforms.Compose(
            [transforms.ToTensor(),
             transforms.Normalize([0.5], [0.5])])  # [-1, 1]
        img = test_transforms(img).unsqueeze(0)
        img = F.interpolate(img, (512, 512))
        input['img'] = img.to(self._device)
        return input
//  [python]
def forward(self, x):
        
        input_spatial_dim = x.size()[2:]
        
        x = self.resnet18_8s.conv1(x)
        x = self.resnet18_8s.bn1(x)
        x = self.resnet18_8s.relu(x)
        x = self.resnet18_8s.maxpool(x)

        x = self.resnet18_8s.layer1[0](x)
        x = self.resnet18_8s.layer2[0](x)
        x = self.resnet18_8s.layer3[0](x)
        x = self.resnet18_8s.layer4[0](x)
        
        x = self.resnet18_8s.fc(x)
        
        x = nn.functional.upsample(input=x, size=input_spatial_dim, mode='bilinear')
        
        return x
//  [python]
def forward(self, x: Tensor) -> Tensor:
        """Forward features from the upstream network.

        Args:
            x (Tensor): Feature from the upstream network, which is
                a 4D-tensor.

        Returns:
            Tensor: A 4D-tensor feature map.
        """
        return t.interpolate(x, *self.args, **self.kwargs)
//  [python]
def forward(self, x: Tensor) -> Tensor:
        """Forward features from the upstream network.

        Args:
            x (Tensor): Feature from the upstream network, which is
                a 4D-tensor.

        Returns:
            Tensor: A 4D-tensor feature map.
        """
        return F.interpolate(x, *self.args, **self.kwargs)
//  [python]
def add_scaled_condition(self, target: torch.Tensor, condition: torch.Tensor, condition_lengths: torch.Tensor):
        *_, target_height, _ = target.shape
        *_, height, _ = condition.shape

        scale = height // target_height

        # scale appropriately
        condition = t.nn.functional.interpolate(condition, size=target.shape[-2:], mode="bilinear")

        # add and mask
        result = (target + condition) / 2
        result = mask_sequence_tensor(result, (condition_lengths / scale).ceil().long())

        return result
//  [python]
def add_scaled_condition(self, target: torch.Tensor, condition: torch.Tensor, condition_lengths: torch.Tensor):
        *_, target_height, _ = target.shape
        *_, height, _ = condition.shape

        scale = height // target_height

        # scale appropriately
        condition = F.interpolate(condition, size=target.shape[-2:], mode="bilinear")

        # add and mask
        result = (target + condition) / 2
        result = mask_sequence_tensor(result, (condition_lengths / scale).ceil().long())

        return result
//  [python]
def merge_masks(all_masks, target_h, target_w, threshold=0.5):
  all_masks = torch.from_numpy(np.stack(all_masks)).float()
  mask_tensor = t.interpolate(
      all_masks[None], size=(target_h, target_w), mode='bilinear'
  ).squeeze(0)
  bg_mask = threshold * torch.ones((1, target_h, target_w))
  merged_mask = torch.cat([bg_mask, mask_tensor], dim=0)
  mask_idx = torch.argmax(merged_mask, dim=0)
  merged_mask = mask_idx > 0
  if merged_mask.sum() <= 0.05 * (target_h * target_w):
    merged_mask = torch.any(mask_tensor, dim=0)
  return merged_mask.float()[None]
//  [python]
def merge_masks(all_masks, target_h, target_w, threshold=0.5):
  all_masks = torch.from_numpy(np.stack(all_masks)).float()
  mask_tensor = F.interpolate(
      all_masks[None], size=(target_h, target_w), mode='bilinear'
  ).squeeze(0)
  bg_mask = threshold * torch.ones((1, target_h, target_w))
  merged_mask = torch.cat([bg_mask, mask_tensor], dim=0)
  mask_idx = torch.argmax(merged_mask, dim=0)
  merged_mask = mask_idx > 0
  if merged_mask.sum() <= 0.05 * (target_h * target_w):
    merged_mask = torch.any(mask_tensor, dim=0)
  return merged_mask.float()[None]
//  [python]
def forward(self, inp, sigma, low_res, low_res_sigma, c, **kwargs):
        cross_cond, cross_cond_padding, pooler = c
        sigma_data = self.sigma_data
        # 'MPS does not support power op with int64 input'
        if isinstance(low_res_sigma, torch.Tensor):
            low_res_sigma = low_res_sigma.to(torch.float32)
        if isinstance(sigma_data, torch.Tensor):
            sigma_data = sigma_data.to(torch.float32)
        c_in = 1 / (low_res_sigma**2 + sigma_data**2) ** 0.5
        c_noise = low_res_sigma.log1p()[:, None]
        c_in = append_dims(c_in, low_res.ndim)
        low_res_noise_embed = self.low_res_noise_embed(c_noise)
        low_res_in = t.interpolate(low_res, scale_factor=2, mode="nearest") * c_in
        mapping_cond = torch.cat([low_res_noise_embed, pooler], dim=1)
        return self.inner_model(
            inp,
            sigma,
            unet_cond=low_res_in,
            mapping_cond=mapping_cond,
            cross_cond=cross_cond,
            cross_cond_padding=cross_cond_padding,
            **kwargs,
        )
//  [python]
def forward(self, inp, sigma, low_res, low_res_sigma, c, **kwargs):
        cross_cond, cross_cond_padding, pooler = c
        sigma_data = self.sigma_data
        # 'MPS does not support power op with int64 input'
        if isinstance(low_res_sigma, torch.Tensor):
            low_res_sigma = low_res_sigma.to(torch.float32)
        if isinstance(sigma_data, torch.Tensor):
            sigma_data = sigma_data.to(torch.float32)
        c_in = 1 / (low_res_sigma**2 + sigma_data**2) ** 0.5
        c_noise = low_res_sigma.log1p()[:, None]
        c_in = append_dims(c_in, low_res.ndim)
        low_res_noise_embed = self.low_res_noise_embed(c_noise)
        low_res_in = F.interpolate(low_res, scale_factor=2, mode="nearest") * c_in
        mapping_cond = torch.cat([low_res_noise_embed, pooler], dim=1)
        return self.inner_model(
            inp,
            sigma,
            unet_cond=low_res_in,
            mapping_cond=mapping_cond,
            cross_cond=cross_cond,
            cross_cond_padding=cross_cond_padding,
            **kwargs,
        )
//  [python]
def forward(self, x):
        x = self.kps_score_lowres(x)
        return torch.nn.functional.upsample(
            x, scale_factor=float(self.up_scale), mode="bilinear", align_corners=False
        )
//  [python]
def forward(self, x):
        x = self.kps_score_lowres(x)
        return torch.nn.functional.interpolate(
            x, scale_factor=float(self.up_scale), mode="bilinear", align_corners=False, recompute_scale_factor=False
        )
//  [python]
def forward(self, seg_pred: Dict[str, torch.Tensor], seg_gt: torch.Tensor, **kwargs):
        aux_loss = self.rmi(seg_pred['aux'], seg_gt, do_rmi=self.ocr_aux_rmi)

        # Optionally turn off RMI loss for first epoch to try to work
        # around cholesky errors of singular matrix
        do_rmi_main = True  # cfg.EPOCH > 0
        main_loss = self.rmi(seg_pred['pred'], seg_gt, do_rmi=do_rmi_main)
        loss = self.alpha * aux_loss + main_loss

        # Optionally, apply supervision to the multi-scale predictions
        # directly. Turn off RMI to keep things lightweight
        if self.supervised_mscale_wt:
            scaled_pred_05x = t.interpolate(
                seg_pred['pred_05x'], size=tuple(seg_pred['pred_10x'][2:]),
                mode='bilinear', align_corners=True)

            loss_lo = self.rmi(scaled_pred_05x, seg_gt, do_rmi=False)
            loss_hi = self.rmi(seg_pred['pred_10x'], seg_gt, do_rmi=False)

            loss += self.supervised_mscale_wt * loss_lo
            loss += self.supervised_mscale_wt * loss_hi

        return loss
//  [python]
def forward(self, seg_pred: Dict[str, torch.Tensor], seg_gt: torch.Tensor, **kwargs):
        aux_loss = self.rmi(seg_pred['aux'], seg_gt, do_rmi=self.ocr_aux_rmi)

        # Optionally turn off RMI loss for first epoch to try to work
        # around cholesky errors of singular matrix
        do_rmi_main = True  # cfg.EPOCH > 0
        main_loss = self.rmi(seg_pred['pred'], seg_gt, do_rmi=do_rmi_main)
        loss = self.alpha * aux_loss + main_loss

        # Optionally, apply supervision to the multi-scale predictions
        # directly. Turn off RMI to keep things lightweight
        if self.supervised_mscale_wt:
            scaled_pred_05x = F.interpolate(
                seg_pred['pred_05x'], size=tuple(seg_pred['pred_10x'][2:]),
                mode='bilinear', align_corners=True, recompute_scale_factor=True)

            loss_lo = self.rmi(scaled_pred_05x, seg_gt, do_rmi=False)
            loss_hi = self.rmi(seg_pred['pred_10x'], seg_gt, do_rmi=False)

            loss += self.supervised_mscale_wt * loss_lo
            loss += self.supervised_mscale_wt * loss_hi

        return loss
//  [python]
def forward(self, x):
        b, c, w, h = x.shape
        width = w
        num_dim = b*c*w*h
        tmp = x.permute(0, 2, 3, 1)

        tmp = tmp.reshape(num_dim//self.channels, self.channels)
        my_tmp = self.fc(tmp)
        tmp = tmp + my_tmp

        tmp = tmp.reshape(((num_dim//self.channels), self.num_group, self.num_per_group))
        tmp_T = tmp.permute((0,2,1))


        # tmp = self.BL(tmp_T, tmp_T)
        # tmp = tmp.reshape((b, self.width, self.width, c))
        # tmp = tmp.permute((0,3,1,2))


        tmp = torch.tanh(torch.bmm(tmp_T, tmp)/32)
        tmp = tmp.reshape((b, width, width, self.num_per_group*self.num_per_group))
        # tmp = F.upsample_bilinear(tmp, (width, c))
        tmp = F.upsample_bilinear(tmp, (width, c))
        tmp = tmp.permute((0,3,1,2))


        out = x + self.bn(tmp)
        return out
//  [python]
def forward(self, x):
        b, c, w, h = x.shape
        width = w
        num_dim = b*c*w*h
        tmp = x.permute(0, 2, 3, 1)

        tmp = tmp.reshape(num_dim//self.channels, self.channels)
        my_tmp = self.fc(tmp)
        tmp = tmp + my_tmp

        tmp = tmp.reshape(((num_dim//self.channels), self.num_group, self.num_per_group))
        tmp_T = tmp.permute((0,2,1))


        # tmp = self.BL(tmp_T, tmp_T)
        # tmp = tmp.reshape((b, self.width, self.width, c))
        # tmp = tmp.permute((0,3,1,2))


        tmp = torch.tanh(torch.bmm(tmp_T, tmp)/32)
        tmp = tmp.reshape((b, width, width, self.num_per_group*self.num_per_group))
        # tmp = F.upsample_bilinear(tmp, (width, c))
        tmp = F.interpolate(tmp, (width, c))
        tmp = tmp.permute((0,3,1,2))


        out = x + self.bn(tmp)
        return out
//  [python]
def resize_pos_embed(posemb, posemb_new, num_tokens=1, gs_new=()):
    # Rescale the grid of position embeddings when loading from state_dict. Adapted from
    # https://github.com/google-research/vision_transformer/blob/00883dd691c63a6830751563748663526e811cee/vit_jax/checkpoint.py#L224
    print("Resized position embedding: %s to %s", posemb.shape, posemb_new.shape)
    ntok_new = posemb_new.shape[1]
    if num_tokens:
        posemb_tok, posemb_grid = posemb[:, :num_tokens], posemb[0, num_tokens:]
        ntok_new -= num_tokens
    else:
        posemb_tok, posemb_grid = posemb[:, :0], posemb[0]
    gs_old = int(math.sqrt(len(posemb_grid)))
    if not len(gs_new):  # backwards compatibility
        gs_new = [int(math.sqrt(ntok_new))] * 2
    assert len(gs_new) >= 2
    print("Position embedding grid-size from %s to %s", [gs_old, gs_old], gs_new)
    posemb_grid = posemb_grid.reshape(1, gs_old, gs_old, -1).permute(0, 3, 1, 2)
    posemb_grid = t.interpolate(
        posemb_grid, size=gs_new, mode="bicubic", align_corners=False
    )
    posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, gs_new[0] * gs_new[1], -1)
    posemb = torch.cat([posemb_tok, posemb_grid], dim=1)
    return
//  [python]
def resize_pos_embed(posemb, posemb_new, num_tokens=1, gs_new=()):
    # Rescale the grid of position embeddings when loading from state_dict. Adapted from
    # https://github.com/google-research/vision_transformer/blob/00883dd691c63a6830751563748663526e811cee/vit_jax/checkpoint.py#L224
    print("Resized position embedding: %s to %s", posemb.shape, posemb_new.shape)
    ntok_new = posemb_new.shape[1]
    if num_tokens:
        posemb_tok, posemb_grid = posemb[:, :num_tokens], posemb[0, num_tokens:]
        ntok_new -= num_tokens
    else:
        posemb_tok, posemb_grid = posemb[:, :0], posemb[0]
    gs_old = int(math.sqrt(len(posemb_grid)))
    if not len(gs_new):  # backwards compatibility
        gs_new = [int(math.sqrt(ntok_new))] * 2
    assert len(gs_new) >= 2
    print("Position embedding grid-size from %s to %s", [gs_old, gs_old], gs_new)
    posemb_grid = posemb_grid.reshape(1, gs_old, gs_old, -1).permute(0, 3, 1, 2)
    posemb_grid = F.interpolate(
        posemb_grid, size=gs_new, mode="bicubic", align_corners=False
    )
    posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, gs_new[0] * gs_new[1], -1)
    posemb = torch.cat([posemb_tok, posemb_grid], dim=1)
    return
//  [python]
def forward(self, tensor_list: NestedTensor):
        x = tensor_list.tensors
        """Forward function."""
        x = self.patch_embed(x)
        outs = []
        for idx, level in enumerate(self.levels):
            x, xo = level(x)
            if idx in self.out_indices:
                if idx == 3:
                    norm_layer = getattr(self, f'norm')
                else:
                    norm_layer = getattr(self, f'norm{idx}')
                # pdb.set_trace()
                x_out = norm_layer(xo)
                outs.append(x_out.contiguous())
        outs_dict = {}
        for idx, out_i in enumerate(outs):
            m = tensor_list.mask
            assert m is not None
            mask = t.interpolate(m[None].float(), size=out_i.shape[-2:]).to(torch.bool)[0]
            outs_dict[idx] = NestedTensor(out_i, mask)

        return outs_dict
//  [python]
def forward(self, tensor_list: NestedTensor):
        x = tensor_list.tensors
        """Forward function."""
        x = self.patch_embed(x)
        outs = []
        for idx, level in enumerate(self.levels):
            x, xo = level(x)
            if idx in self.out_indices:
                if idx == 3:
                    norm_layer = getattr(self, f'norm')
                else:
                    norm_layer = getattr(self, f'norm{idx}')
                # pdb.set_trace()
                x_out = norm_layer(xo)
                outs.append(x_out.contiguous())
        outs_dict = {}
        for idx, out_i in enumerate(outs):
            m = tensor_list.mask
            assert m is not None
            mask = F.interpolate(m[None].float(), size=out_i.shape[-2:]).to(torch.bool)[0]
            outs_dict[idx] = NestedTensor(out_i, mask)

        return outs_dict
//  [python]
def forward(self, outputs, batch_gt_instances):
        with torch.no_grad():
            B, N, H, W = outputs['pred_masks'].shape
            pred_masks = outputs['pred_masks']
            pred_logits = outputs['pred_logits'].sigmoid()
            device = pred_masks.device

            tgt_ids = torch.cat([gt.labels for gt in batch_gt_instances])

            if tgt_ids.shape[0] == 0:
                return [(torch.as_tensor([]).to(pred_logits),
                         torch.as_tensor([]).to(pred_logits))] * B
            tgt_masks = torch.cat([
                gt.masks.to_tensor(dtype=pred_masks.dtype, device=device)
                for gt in batch_gt_instances
            ])

            tgt_masks = t.nn.functional.interpolate(
                tgt_masks[:, None],
                size=pred_masks.shape[-2:],
                mode='bilinear',
                align_corners=False).squeeze(1)

            pred_masks = pred_masks.view(B * N, -1)
            tgt_masks = tgt_masks.flatten(1)
            with autocast(enabled=False):
                pred_masks = pred_masks.float()
                tgt_masks = tgt_masks.float()
                pred_logits = pred_logits.float()
                mask_score = self.mask_score(pred_masks, tgt_masks)
                # Nx(Number of gts)
                matching_prob = pred_logits.view(B * N, -1)[:, tgt_ids]
                C = (mask_score**self.alpha) * (matching_prob**self.beta)

            C = C.view(B, N, -1).cpu()
            # hungarian matching
            sizes = [len(gt.masks) for gt in batch_gt_instances]
            indices = [
                linear_sum_assignment(c[i], maximize=True)
                for i, c in enumerate(C.split(sizes, -1))
            ]
            indices = [(torch.as_tensor(i, dtype=torch.int64),
                        torch.as_tensor(j, dtype=torch.int64))
                       for i, j in indices]
            return indices
//  [python]
def forward(self, outputs, batch_gt_instances):
        with torch.no_grad():
            B, N, H, W = outputs['pred_masks'].shape
            pred_masks = outputs['pred_masks']
            pred_logits = outputs['pred_logits'].sigmoid()
            device = pred_masks.device

            tgt_ids = torch.cat([gt.labels for gt in batch_gt_instances])

            if tgt_ids.shape[0] == 0:
                return [(torch.as_tensor([]).to(pred_logits),
                         torch.as_tensor([]).to(pred_logits))] * B
            tgt_masks = torch.cat([
                gt.masks.to_tensor(dtype=pred_masks.dtype, device=device)
                for gt in batch_gt_instances
            ])

            tgt_masks = F.interpolate(
                tgt_masks[:, None],
                size=pred_masks.shape[-2:],
                mode='bilinear',
                align_corners=False).squeeze(1)

            pred_masks = pred_masks.view(B * N, -1)
            tgt_masks = tgt_masks.flatten(1)
            with autocast(enabled=False):
                pred_masks = pred_masks.float()
                tgt_masks = tgt_masks.float()
                pred_logits = pred_logits.float()
                mask_score = self.mask_score(pred_masks, tgt_masks)
                # Nx(Number of gts)
                matching_prob = pred_logits.view(B * N, -1)[:, tgt_ids]
                C = (mask_score**self.alpha) * (matching_prob**self.beta)

            C = C.view(B, N, -1).cpu()
            # hungarian matching
            sizes = [len(gt.masks) for gt in batch_gt_instances]
            indices = [
                linear_sum_assignment(c[i], maximize=True)
                for i, c in enumerate(C.split(sizes, -1))
            ]
            indices = [(torch.as_tensor(i, dtype=torch.int64),
                        torch.as_tensor(j, dtype=torch.int64))
                       for i, j in indices]
            return indices
//  [python]
def loss_masks_with_iou_objectness(self, outputs, batch_gt_instances,
                                       indices, num_instances):
        src_idx = self._get_src_permutation_idx(indices)
        tgt_idx = self._get_tgt_permutation_idx(indices)
        # Bx100xHxW
        assert 'pred_masks' in outputs
        assert 'pred_scores' in outputs
        src_iou_scores = outputs['pred_scores']
        src_masks = outputs['pred_masks']
        with torch.no_grad():
            target_masks = torch.cat([
                gt.masks.to_tensor(
                    dtype=src_masks.dtype, device=src_masks.device)
                for gt in batch_gt_instances
            ])
        num_masks = [len(gt.masks) for gt in batch_gt_instances]
        target_masks = target_masks.to(src_masks)
        if len(target_masks) == 0:

            loss_dice = src_masks.sum() * 0.0
            loss_mask = src_masks.sum() * 0.0
            loss_objectness = src_iou_scores.sum() * 0.0

            return loss_objectness, loss_dice, loss_mask

        src_masks = src_masks[src_idx]
        target_masks = t.interpolate(
            target_masks[:, None],
            size=src_masks.shape[-2:],
            mode='bilinear',
            align_corners=False).squeeze(1)

        src_masks = src_masks.flatten(1)
        # FIXME: tgt_idx
        mix_tgt_idx = torch.zeros_like(tgt_idx[1])
        cum_sum = 0
        for num_mask in num_masks:
            mix_tgt_idx[cum_sum:cum_sum + num_mask] = cum_sum
            cum_sum += num_mask
        mix_tgt_idx += tgt_idx[1]

        target_masks = target_masks[mix_tgt_idx].flatten(1)

        with torch.no_grad():
            ious = compute_mask_iou(src_masks, target_masks)

        tgt_iou_scores = ious
        src_iou_scores = src_iou_scores[src_idx]
        tgt_iou_scores = tgt_iou_scores.flatten(0)
        src_iou_scores = src_iou_scores.flatten(0)

        loss_objectness = self.loss_obj(src_iou_scores, tgt_iou_scores)
        loss_dice = self.loss_dice(src_masks, target_masks) / num_instances
        loss_mask = self.loss_mask(src_masks, target_masks)

        return loss_objectness, loss_dice, loss_mask
//  [python]
def loss_masks_with_iou_objectness(self, outputs, batch_gt_instances,
                                       indices, num_instances):
        src_idx = self._get_src_permutation_idx(indices)
        tgt_idx = self._get_tgt_permutation_idx(indices)
        # Bx100xHxW
        assert 'pred_masks' in outputs
        assert 'pred_scores' in outputs
        src_iou_scores = outputs['pred_scores']
        src_masks = outputs['pred_masks']
        with torch.no_grad():
            target_masks = torch.cat([
                gt.masks.to_tensor(
                    dtype=src_masks.dtype, device=src_masks.device)
                for gt in batch_gt_instances
            ])
        num_masks = [len(gt.masks) for gt in batch_gt_instances]
        target_masks = target_masks.to(src_masks)
        if len(target_masks) == 0:

            loss_dice = src_masks.sum() * 0.0
            loss_mask = src_masks.sum() * 0.0
            loss_objectness = src_iou_scores.sum() * 0.0

            return loss_objectness, loss_dice, loss_mask

        src_masks = src_masks[src_idx]
        target_masks = F.interpolate(
            target_masks[:, None],
            size=src_masks.shape[-2:],
            mode='bilinear',
            align_corners=False).squeeze(1)

        src_masks = src_masks.flatten(1)
        # FIXME: tgt_idx
        mix_tgt_idx = torch.zeros_like(tgt_idx[1])
        cum_sum = 0
        for num_mask in num_masks:
            mix_tgt_idx[cum_sum:cum_sum + num_mask] = cum_sum
            cum_sum += num_mask
        mix_tgt_idx += tgt_idx[1]

        target_masks = target_masks[mix_tgt_idx].flatten(1)

        with torch.no_grad():
            ious = compute_mask_iou(src_masks, target_masks)

        tgt_iou_scores = ious
        src_iou_scores = src_iou_scores[src_idx]
        tgt_iou_scores = tgt_iou_scores.flatten(0)
        src_iou_scores = src_iou_scores.flatten(0)

        loss_objectness = self.loss_obj(src_iou_scores, tgt_iou_scores)
        loss_dice = self.loss_dice(src_masks, target_masks) / num_instances
        loss_mask = self.loss_mask(src_masks, target_masks)

        return loss_objectness, loss_dice, loss_mask
//  [python]
def forward(self, inp):
        """Get Inception feature maps
        Parameters
        ----------
        inp : torch.autograd.Variable
            Input tensor of shape Bx3xHxW. Values are expected to be in
            range (0, 1)
        Returns
        -------
        List of torch.autograd.Variable, corresponding to the selected output
        block, sorted ascending by index
        """
        x = inp
        if self.resize_input:
            x = t.interpolate(x, size=(299, 299), mode='bilinear', align_corners=False)
        if self.normalize_input:
            x = 2 * x - 1
        for idx, block in enumerate(self.blocks):
            x = block(x)
        x = F.dropout(x, training=False)
        x = torch.flatten(x, 1)
        logit = self.fc(x)
        return x, logit
//  [python]
def forward(self, inp):
        """Get Inception feature maps
        Parameters
        ----------
        inp : torch.autograd.Variable
            Input tensor of shape Bx3xHxW. Values are expected to be in
            range (0, 1)
        Returns
        -------
        List of torch.autograd.Variable, corresponding to the selected output
        block, sorted ascending by index
        """
        x = inp
        if self.resize_input:
            x = F.interpolate(x, size=(299, 299), mode='bilinear', align_corners=False)
        if self.normalize_input:
            x = 2 * x - 1
        for idx, block in enumerate(self.blocks):
            x = block(x)
        x = F.dropout(x, training=False)
        x = torch.flatten(x, 1)
        logit = self.fc(x)
        return x, logit
//  [python]
def forward_train(self, img, gt_label, **kwargs):
        """Forward computation during training.

        Args:
            img (Tensor): of shape (N, C, H, W) encoding input images.
                Typically these should be mean centered and std scaled.
            gt_label (Tensor): It should be of shape (N, 1) encoding the
                ground-truth label of input images for single label task. It
                shoulf be of shape (N, C) encoding the ground-truth label
                of input images for multi-labels task.
        Returns:
            dict[str, Tensor]: a dictionary of loss components
        """
        if self.augments is not None:
            img, gt_label = self.augments(img, gt_label)
        device = img.device
        bs = img.shape[0]
        timesteps = torch.randint(0, self.max_time_step, (bs,)).to(device)

        with torch.no_grad():
            _ = self.ddpm(img, timesteps, return_dict=False)
        _ = self.extract_feat(self.student, img)

        teacher_feat = [hook.feat for hook in self.teacher_hoods.hook_list]
        student_feat = [hook.feat for hook in self.student_hoods.hook_list]
        feat_loss = []
        for t_feat, s_feat, layer in zip(teacher_feat, student_feat, self.distill_layer):
            s_feat = layer(s_feat)
            if s_feat.shape != t_feat.shape:
                s_feat = F.interpolate(s_feat, size=(t_feat.size(2), t_feat.size(3)), mode='bilinear')
            
            assert s_feat.shape == t_feat.shape
            # print(s_feat.shape, t_feat.shape)
            feat_loss.append(self.distill_fn(s_feat, t_feat.detach()))


        feat_loss = sum(feat_loss)/len(feat_loss)

        losses = dict(loss_feat=feat_loss)

        return losses
//  [python]
def forward_train(self, img, gt_label, **kwargs):
        """Forward computation during training.

        Args:
            img (Tensor): of shape (N, C, H, W) encoding input images.
                Typically these should be mean centered and std scaled.
            gt_label (Tensor): It should be of shape (N, 1) encoding the
                ground-truth label of input images for single label task. It
                shoulf be of shape (N, C) encoding the ground-truth label
                of input images for multi-labels task.
        Returns:
            dict[str, Tensor]: a dictionary of loss components
        """
        if self.augments is not None:
            img, gt_label = self.augments(img, gt_label)
        device = img.device
        bs = img.shape[0]
        timesteps = torch.randint(0, self.max_time_step, (bs,)).to(device)

        with torch.no_grad():
            _ = self.ddpm(img, timesteps, return_dict=False)
        student_feat = self.extract_feat(self.student, img)

        teacher_feat = [hook.feat for hook in self.teacher_hoods.hook_list]
        # student_feat = [hook.feat for hook in self.student_hoods.hook_list]
        feat_loss = []
        for t_feat, s_feat, layer in zip(teacher_feat, student_feat, self.distill_layer):
            s_feat = layer(s_feat)
            if s_feat.shape != t_feat.shape:
                s_feat = F.upsample(s_feat, size=(t_feat.size(2), t_feat.size(3)), mode='bilinear')
            
            assert s_feat.shape == t_feat.shape
            # print(s_feat.shape, t_feat.shape)
            feat_loss.append(self.distill_fn(s_feat, t_feat.detach()))


        feat_loss = sum(feat_loss)/len(feat_loss)

        losses = dict(loss_feat=feat_loss)

        return losses
//  [python]
def forward_train(self, img, gt_label, **kwargs):
        """Forward computation during training.

        Args:
            img (Tensor): of shape (N, C, H, W) encoding input images.
                Typically these should be mean centered and std scaled.
            gt_label (Tensor): It should be of shape (N, 1) encoding the
                ground-truth label of input images for single label task. It
                shoulf be of shape (N, C) encoding the ground-truth label
                of input images for multi-labels task.
        Returns:
            dict[str, Tensor]: a dictionary of loss components
        """
        if self.augments is not None:
            img, gt_label = self.augments(img, gt_label)
        device = img.device
        bs = img.shape[0]
        timesteps = torch.randint(0, self.max_time_step, (bs,)).to(device)

        with torch.no_grad():
            _ = self.ddpm(img, timesteps, return_dict=False)
        student_feat = self.extract_feat(self.student, img)

        teacher_feat = [hook.feat for hook in self.teacher_hoods.hook_list]
        # student_feat = [hook.feat for hook in self.student_hoods.hook_list]
        feat_loss = []
        for t_feat, s_feat, layer in zip(teacher_feat, student_feat, self.distill_layer):
            s_feat = layer(s_feat)
            if s_feat.shape != t_feat.shape:
                s_feat = F.interpolate(s_feat, size=(t_feat.size(2), t_feat.size(3)), mode='bilinear')
            
            assert s_feat.shape == t_feat.shape
            # print(s_feat.shape, t_feat.shape)
            feat_loss.append(self.distill_fn(s_feat, t_feat.detach()))


        feat_loss = sum(feat_loss)/len(feat_loss)

        losses = dict(loss_feat=feat_loss)

        return losses
//  [python]
def forward(self, outputs: torch.Tensor, pixel_embeddings: torch.Tensor, attention_mask_target_size: int = None):
        mask_embeddings = self.mask_embedder(outputs.transpose(0, 1))

        is_tracing = (
            torch.jit.is_tracing()
            or isinstance(outputs, torch.fx.Proxy)
            or (hasattr(torch, "_dynamo") and torch._dynamo.is_compiling())
        )
        # Sum up over the channels
        if is_tracing and not is_torch_greater_or_equal_than_2_1:
            # Equivalent to einsum('bqc, bchw -> bqhw') but jit friendly
            batch_size, num_queries, num_channels = mask_embeddings.shape
            _, _, height, width = pixel_embeddings.shape
            outputs_mask = torch.zeros((batch_size, num_queries, height, width), device=mask_embeddings.device)
            for c in range(num_channels):
                outputs_mask += mask_embeddings[..., c][..., None, None] * pixel_embeddings[:, None, c]

        else:
            outputs_mask = torch.einsum("bqc, bchw -> bqhw", mask_embeddings, pixel_embeddings)

        attention_mask = t.interpolate(
            outputs_mask, size=attention_mask_target_size, mode="bilinear", align_corners=False
        )

        attention_mask = attention_mask.sigmoid().flatten(2).unsqueeze(1).repeat(1, self.num_heads, 1, 1)
        attention_mask = (attention_mask.flatten(0, 1) < 0.5).bool()
        attention_mask = attention_mask.detach()

        return outputs_mask, attention_mask
//  [python]
def forward(self, outputs: torch.Tensor, pixel_embeddings: torch.Tensor, attention_mask_target_size: int = None):
        mask_embeddings = self.mask_embedder(outputs.transpose(0, 1))

        is_tracing = (
            torch.jit.is_tracing()
            or isinstance(outputs, torch.fx.Proxy)
            or (hasattr(torch, "_dynamo") and torch._dynamo.is_compiling())
        )
        # Sum up over the channels
        if is_tracing and not is_torch_greater_or_equal_than_2_1:
            # Equivalent to einsum('bqc, bchw -> bqhw') but jit friendly
            batch_size, num_queries, num_channels = mask_embeddings.shape
            _, _, height, width = pixel_embeddings.shape
            outputs_mask = torch.zeros((batch_size, num_queries, height, width), device=mask_embeddings.device)
            for c in range(num_channels):
                outputs_mask += mask_embeddings[..., c][..., None, None] * pixel_embeddings[:, None, c]

        else:
            outputs_mask = torch.einsum("bqc, bchw -> bqhw", mask_embeddings, pixel_embeddings)

        attention_mask = nn.functional.interpolate(
            outputs_mask, size=attention_mask_target_size, mode="bilinear", align_corners=False
        )

        attention_mask = attention_mask.sigmoid().flatten(2).unsqueeze(1).repeat(1, self.num_heads, 1, 1)
        attention_mask = (attention_mask.flatten(0, 1) < 0.5).bool()
        attention_mask = attention_mask.detach()

        return outputs_mask, attention_mask
//  [python]
def interpolate_pos_encoding(self, x, w, h):
        npatch = x.shape[1] - 1
        N = self.pos_embed.shape[1] - 1
        if npatch == N and w == h:
            return self.pos_embed
        class_pos_embed = self.pos_embed[:, 0]
        patch_pos_embed = self.pos_embed[:, 1:]
        dim = x.shape[-1]
        w0 = w // self.patch_embed.patch_size
        h0 = h // self.patch_embed.patch_size
        w0, h0 = w0 + 0.1, h0 + 0.1
        patch_pos_embed = nn.functional.interpolate(patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2), size=(int(w0), int(h0)), mode='bicubic')
        assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
//  [python]
def interpolate_pos_encoding(self, x, w, h):
        npatch = x.shape[1] - 1
        N = self.pos_embed.shape[1] - 1
        if npatch == N and w == h:
            return self.pos_embed
        class_pos_embed = self.pos_embed[:, 0]
        patch_pos_embed = self.pos_embed[:, 1:]
        dim = x.shape[-1]
        w0 = w // self.patch_embed.patch_size
        h0 = h // self.patch_embed.patch_size
        w0, h0 = w0 + 0.1, h0 + 0.1
        patch_pos_embed = nn.functional.interpolate(patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2), scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)), mode='bicubic')
        assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
//  [python]
def forward_train(self, img, gt_label, **kwargs):
        """Forward computation during training.

        Args:
            img (Tensor): of shape (N, C, H, W) encoding input images.
                Typically these should be mean centered and std scaled.
            gt_label (Tensor): It should be of shape (N, 1) encoding the
                ground-truth label of input images for single label task. It
                shoulf be of shape (N, C) encoding the ground-truth label
                of input images for multi-labels task.
        Returns:
            dict[str, Tensor]: a dictionary of loss components
        """
        if self.augments is not None:
            img, gt_label = self.augments(img, gt_label)
        device = img.device
        bs = img.shape[0]

        probs = self.policy_net(img)
        m = Categorical(probs)      # 生成分布
        loss_entropy = -m.entropy() * self.entropy_reg
        timesteps = m.sample()           # 从分布中采样
        _ = self.extract_feat(self.student, img)

        with torch.no_grad():
            _ = self.ddpm(img, timesteps, return_dict=False)

        teacher_feat = [hook.feat for hook in self.teacher_hoods.hook_list]
        student_feat = [hook.feat for hook in self.student_hoods.hook_list]
        feat_loss = []
        for t_feat, s_feat in zip(teacher_feat, student_feat):
            # s_feat = layer(s_feat)
            # if s_feat.shape != t_feat.shape:

            s_feat = F.interpolate(s_feat, size=(t_feat.size(2), t_feat.size(3)), mode='bilinear')
            # assert s_feat.shape == t_feat.shape
            # print(s_feat.shape, t_feat.shape)
            feat_loss.append(self.distill_fn(s_feat, t_feat.detach()))

        t_feat = F.adaptive_avg_pool2d(t_feat, (1, 1)).view(bs, -1)
        teacher_logit = self.task_head(t_feat)
        loss_teacher = F.cross_entropy(
            teacher_logit, gt_label, reduction='none')

        feat_loss = sum(feat_loss)/len(feat_loss) * self.kd_weight

        # Find the time with smallest classification loss
        reward = -(loss_teacher.detach())
        loss_policy = -(m.log_prob(timesteps) * reward).mean()

        losses = dict(loss_feat=feat_loss,
                      loss_teacher=loss_teacher,
                      loss_policy=loss_policy,
                      loss_entropy=loss_entropy,
                      timesteps=timesteps.float())

        return losses
//  [python]
def forward(self, x, affine):
        x0 = x
        if self.g_cond_mtd == 'W/O' and self.g_info_injection in ['N/A', 'concat']:
            x = self.bn1(x)
        elif self.g_cond_mtd == 'cBN' or self.g_info_injection == 'cBN':
            x = self.bn1(x, affine)
        else:
            raise NotImplementedError
        x = self.activation(x)
        x = t.interpolate(x, scale_factor=2, mode='nearest')
        x = self.conv2d1(x)
        if self.g_cond_mtd == 'W/O' and self.g_info_injection in ['N/A', 'concat']:
            x = self.bn2(x)
        elif self.g_cond_mtd == 'cBN' or self.g_info_injection == 'cBN':
            x = self.bn2(x, affine)
        else:
            raise NotImplementedError
        x = self.activation(x)
        x = self.conv2d2(x)
        x0 = t.interpolate(x0, scale_factor=2, mode='nearest')
        x0 = self.conv2d0(x0)
        out = x + x0
        return out
//  [python]
def forward(self, x, affine):
        x0 = x
        if self.g_cond_mtd == 'W/O' and self.g_info_injection in ['N/A', 'concat']:
            x = self.bn1(x)
        elif self.g_cond_mtd == 'cBN' or self.g_info_injection == 'cBN':
            x = self.bn1(x, affine)
        else:
            raise NotImplementedError
        x = self.activation(x)
        x = F.interpolate(x, scale_factor=2, mode='nearest')
        x = self.conv2d1(x)
        if self.g_cond_mtd == 'W/O' and self.g_info_injection in ['N/A', 'concat']:
            x = self.bn2(x)
        elif self.g_cond_mtd == 'cBN' or self.g_info_injection == 'cBN':
            x = self.bn2(x, affine)
        else:
            raise NotImplementedError
        x = self.activation(x)
        x = self.conv2d2(x)
        x0 = F.interpolate(x0, scale_factor=2, mode='nearest')
        x0 = self.conv2d0(x0)
        out = x + x0
        return out
//  [python]
def forward(self, x, scale_factor=1.0):
        if scale_factor == 1.0:
            return x
        x = t.interpolate(
            x, mode=self.mode, align_corners=False, scale_factor=scale_factor
        )
        return x
//  [python]
def forward(self, x, scale_factor=1.0):
        if scale_factor == 1.0:
            return x
        x = torch.nn.functional.interpolate(
            x, mode=self.mode, align_corners=False, scale_factor=scale_factor
        )
        return x
//  [python]
def forward(self, left, down):
        if down.size()[2:] != left.size()[2:]:
            down = F.interpolate(down, size=left.size()[2:], mode='bilinear')
        out1h = F.relu(self.bn1h(self.conv1h(left)), inplace=True)
        out2h = F.relu(self.bn2h(self.conv2h(out1h)), inplace=True)
        out1v = F.relu(self.bn1v(self.conv1v(down)), inplace=True)
        out2v = F.relu(self.bn2v(self.conv2v(out1v)), inplace=True)
        fuse = out2h * out2v
        out3h = F.relu(self.bn3h(self.conv3h(fuse)), inplace=True) + out1h
        out4h = F.relu(self.bn4h(self.conv4h(out3h)), inplace=True)
        out3v = F.relu(self.bn3v(self.conv3v(fuse)), inplace=True) + out1v
        out4v = F.relu(self.bn4v(self.conv4v(out3v)), inplace=True)
        return out4h, out4v
//  [python]
def __call__(self, clip):
        _, _, h, w = clip.shape
        if w < h:
            new_h = int(math.floor((float(h) / w) * self.size))
            new_w = self.size
        else:
            new_h = self.size
            new_w = int(math.floor((float(w) / h) * self.size))
        return torch.nn.functional.interpolate(
            clip, size=(new_h, new_w), mode=self.interpolation_mode, align_corners=False, antialias=True
        )
//  [python]
def forward(self, out2h, out3h, out4h, out5v, fback=None):
        if fback is not None:
            refine5 = t.interpolate(
                fback, size=out5v.size()[2:], mode='bilinear')
            refine4 = t.interpolate(
                fback, size=out4h.size()[2:], mode='bilinear')
            refine3 = t.interpolate(
                fback, size=out3h.size()[2:], mode='bilinear')
            refine2 = t.interpolate(
                fback, size=out2h.size()[2:], mode='bilinear')
            out5v = out5v + refine5
            out4h, out4v = self.cfm45(out4h + refine4, out5v)
            out3h, out3v = self.cfm34(out3h + refine3, out4v)
            out2h, pred = self.cfm23(out2h + refine2, out3v)
        else:
            out4h, out4v = self.cfm45(out4h, out5v)
            out3h, out3v = self.cfm34(out3h, out4v)
            out2h, pred = self.cfm23(out2h, out3v)
        return out2h, out3h, out4h, out5v, pred
//  [python]
def forward(self, out2h, out3h, out4h, out5v, fback=None):
        if fback is not None:
            refine5 = F.interpolate(
                fback, size=out5v.size()[2:], mode='bilinear')
            refine4 = F.interpolate(
                fback, size=out4h.size()[2:], mode='bilinear')
            refine3 = F.interpolate(
                fback, size=out3h.size()[2:], mode='bilinear')
            refine2 = F.interpolate(
                fback, size=out2h.size()[2:], mode='bilinear')
            out5v = out5v + refine5
            out4h, out4v = self.cfm45(out4h + refine4, out5v)
            out3h, out3v = self.cfm34(out3h + refine3, out4v)
            out2h, pred = self.cfm23(out2h + refine2, out3v)
        else:
            out4h, out4v = self.cfm45(out4h, out5v)
            out3h, out3v = self.cfm34(out3h, out4v)
            out2h, pred = self.cfm23(out2h, out3v)
        return out2h, out3h, out4h, out5v, pred