Create modeling_fcrnn.py

modeling_fcrnn.py

@@ -0,0 +1,1926 @@
+"""
+ coding=utf-8
+ Copyright 2022, Ontocord, LLC
+ Copyright 2018, Antonio Mendoza Hao Tan, Mohit Bansal
+ Adapted From Facebook Inc, Detectron2 && Huggingface Co.
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.import copy
+ """
+import itertools
+import math
+import os
+from abc import ABCMeta, abstractmethod
+from collections import OrderedDict, namedtuple
+from typing import Dict, List, Tuple
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn.modules.batchnorm import BatchNorm2d
+from torchvision.ops import RoIPool
+from torchvision.ops.boxes import batched_nms, nms
+
+from .utils import WEIGHTS_NAME, Config, cached_path, hf_bucket_url, is_remote_url, load_checkpoint
+
+
+# other:
+def norm_box(boxes, raw_sizes):
+    if not isinstance(boxes, torch.Tensor):
+        normalized_boxes = boxes.copy()
+    else:
+        normalized_boxes = boxes.clone()
+    normalized_boxes[:, :, (0, 2)] /= raw_sizes[:, 1].view(-1, 1, 1)
+    normalized_boxes[:, :, (1, 3)] /= raw_sizes[:, 0].view(-1, 1, 1)
+    return normalized_boxes
+
+
+def pad_list_tensors(
+    list_tensors,
+    preds_per_image,
+    max_detections=None,
+    return_tensors=None,
+    padding=None,
+    pad_value=0,
+    location=None,
+):
+    """
+    location will always be cpu for np tensors
+    """
+    if location is None:
+        location = "cpu"
+    assert return_tensors in {"pt", "np", None}
+    assert padding in {"max_detections", "max_batch", None}
+    new = []
+    if padding is None:
+        if return_tensors is None:
+            return list_tensors
+        elif return_tensors == "pt":
+            if not isinstance(list_tensors, torch.Tensor):
+                return torch.stack(list_tensors).to(location)
+            else:
+                return list_tensors.to(location)
+        else:
+            if not isinstance(list_tensors, list):
+                return np.array(list_tensors.to(location))
+            else:
+                return list_tensors.to(location)
+    if padding == "max_detections":
+        assert max_detections is not None, "specify max number of detections per batch"
+    elif padding == "max_batch":
+        max_detections = max(preds_per_image)
+    for i in range(len(list_tensors)):
+        too_small = False
+        tensor_i = list_tensors.pop(0)
+        if tensor_i.ndim < 2:
+            too_small = True
+            tensor_i = tensor_i.unsqueeze(-1)
+        assert isinstance(tensor_i, torch.Tensor)
+        tensor_i = nn.functional.pad(
+            input=tensor_i,
+            pad=(0, 0, 0, max_detections - preds_per_image[i]),
+            mode="constant",
+            value=pad_value,
+        )
+        if too_small:
+            tensor_i = tensor_i.squeeze(-1)
+        if return_tensors is None:
+            if location == "cpu":
+                tensor_i = tensor_i.cpu()
+            tensor_i = tensor_i.tolist()
+        if return_tensors == "np":
+            if location == "cpu":
+                tensor_i = tensor_i.cpu()
+            tensor_i = tensor_i.numpy()
+        else:
+            if location == "cpu":
+                tensor_i = tensor_i.cpu()
+        new.append(tensor_i)
+    if return_tensors == "np":
+        return np.stack(new, axis=0)
+    elif return_tensors == "pt" and not isinstance(new, torch.Tensor):
+        return torch.stack(new, dim=0)
+    else:
+        return list_tensors
+
+
+def do_nms(boxes, scores, image_shape, score_thresh, nms_thresh, mind, maxd):
+    scores = scores[:, :-1]
+    num_bbox_reg_classes = boxes.shape[1] // 4
+    # Convert to Boxes to use the `clip` function ...
+    boxes = boxes.reshape(-1, 4)
+    _clip_box(boxes, image_shape)
+    boxes = boxes.view(-1, num_bbox_reg_classes, 4)  # R x C x 4
+
+    # Select max scores
+    max_scores, max_classes = scores.max(1)  # R x C --> R
+    num_objs = boxes.size(0)
+    boxes = boxes.view(-1, 4)
+    idxs = torch.arange(num_objs).to(boxes.device) * num_bbox_reg_classes + max_classes
+    max_boxes = boxes[idxs]  # Select max boxes according to the max scores.
+
+    # Apply NMS
+    keep = nms(max_boxes, max_scores, nms_thresh)
+    keep = keep[:maxd]
+    if keep.shape[-1] >= mind and keep.shape[-1] <= maxd:
+        max_boxes, max_scores = max_boxes[keep], max_scores[keep]
+        classes = max_classes[keep]
+        return max_boxes, max_scores, classes, keep
+    else:
+        return None
+
+
+# Helper Functions
+def _clip_box(tensor, box_size: Tuple[int, int]):
+    assert torch.isfinite(tensor).all(), "Box tensor contains infinite or NaN!"
+    h, w = box_size
+    tensor[:, 0].clamp_(min=0, max=w)
+    tensor[:, 1].clamp_(min=0, max=h)
+    tensor[:, 2].clamp_(min=0, max=w)
+    tensor[:, 3].clamp_(min=0, max=h)
+
+
+def _nonempty_boxes(box, threshold: float = 0.0) -> torch.Tensor:
+    widths = box[:, 2] - box[:, 0]
+    heights = box[:, 3] - box[:, 1]
+    keep = (widths > threshold) & (heights > threshold)
+    return keep
+
+
+def get_norm(norm, out_channels):
+    if isinstance(norm, str):
+        if len(norm) == 0:
+            return None
+        norm = {
+            "BN": BatchNorm2d,
+            "GN": lambda channels: nn.GroupNorm(32, channels),
+            "nnSyncBN": nn.SyncBatchNorm,  # keep for debugging
+            "": lambda x: x,
+        }[norm]
+    return norm(out_channels)
+
+
+def _create_grid_offsets(size: List[int], stride: int, offset: float, device):
+
+    grid_height, grid_width = size
+    shifts_x = torch.arange(
+        offset * stride,
+        grid_width * stride,
+        step=stride,
+        dtype=torch.float32,
+        device=device,
+    )
+    shifts_y = torch.arange(
+        offset * stride,
+        grid_height * stride,
+        step=stride,
+        dtype=torch.float32,
+        device=device,
+    )
+
+    shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x)
+    shift_x = shift_x.reshape(-1)
+    shift_y = shift_y.reshape(-1)
+    return shift_x, shift_y
+
+
+def build_backbone(cfg):
+    input_shape = ShapeSpec(channels=len(cfg.MODEL.PIXEL_MEAN))
+    norm = cfg.RESNETS.NORM
+    stem = BasicStem(
+        in_channels=input_shape.channels,
+        out_channels=cfg.RESNETS.STEM_OUT_CHANNELS,
+        norm=norm,
+        caffe_maxpool=cfg.MODEL.MAX_POOL,
+    )
+    freeze_at = cfg.BACKBONE.FREEZE_AT
+
+    if freeze_at >= 1:
+        for p in stem.parameters():
+            p.requires_grad = False
+
+    out_features = cfg.RESNETS.OUT_FEATURES
+    depth = cfg.RESNETS.DEPTH
+    num_groups = cfg.RESNETS.NUM_GROUPS
+    width_per_group = cfg.RESNETS.WIDTH_PER_GROUP
+    bottleneck_channels = num_groups * width_per_group
+    in_channels = cfg.RESNETS.STEM_OUT_CHANNELS
+    out_channels = cfg.RESNETS.RES2_OUT_CHANNELS
+    stride_in_1x1 = cfg.RESNETS.STRIDE_IN_1X1
+    res5_dilation = cfg.RESNETS.RES5_DILATION
+    assert res5_dilation in {1, 2}, "res5_dilation cannot be {}.".format(res5_dilation)
+
+    num_blocks_per_stage = {50: [3, 4, 6, 3], 101: [3, 4, 23, 3], 152: [3, 8, 36, 3]}[depth]
+
+    stages = []
+    out_stage_idx = [{"res2": 2, "res3": 3, "res4": 4, "res5": 5}[f] for f in out_features]
+    max_stage_idx = max(out_stage_idx)
+    for idx, stage_idx in enumerate(range(2, max_stage_idx + 1)):
+        dilation = res5_dilation if stage_idx == 5 else 1
+        first_stride = 1 if idx == 0 or (stage_idx == 5 and dilation == 2) else 2
+        stage_kargs = {
+            "num_blocks": num_blocks_per_stage[idx],
+            "first_stride": first_stride,
+            "in_channels": in_channels,
+            "bottleneck_channels": bottleneck_channels,
+            "out_channels": out_channels,
+            "num_groups": num_groups,
+            "norm": norm,
+            "stride_in_1x1": stride_in_1x1,
+            "dilation": dilation,
+        }
+
+        stage_kargs["block_class"] = BottleneckBlock
+        blocks = ResNet.make_stage(**stage_kargs)
+        in_channels = out_channels
+        out_channels *= 2
+        bottleneck_channels *= 2
+
+        if freeze_at >= stage_idx:
+            for block in blocks:
+                block.freeze()
+        stages.append(blocks)
+
+    return ResNet(stem, stages, out_features=out_features)
+
+
+def find_top_rpn_proposals(
+    proposals,
+    pred_objectness_logits,
+    images,
+    image_sizes,
+    nms_thresh,
+    pre_nms_topk,
+    post_nms_topk,
+    min_box_side_len,
+    training,
+):
+    """Args:
+        proposals (list[Tensor]): (L, N, Hi*Wi*A, 4).
+        pred_objectness_logits: tensors of length L.
+        nms_thresh (float): IoU threshold to use for NMS
+        pre_nms_topk (int): before nms
+        post_nms_topk (int): after nms
+        min_box_side_len (float): minimum proposal box side
+        training (bool): True if proposals are to be used in training,
+    Returns:
+        results (List[Dict]): stores post_nms_topk object proposals for image i.
+    """
+    num_images = len(images)
+    device = proposals[0].device
+
+    # 1. Select top-k anchor for every level and every image
+    topk_scores = []  # #lvl Tensor, each of shape N x topk
+    topk_proposals = []
+    level_ids = []  # #lvl Tensor, each of shape (topk,)
+    batch_idx = torch.arange(num_images, device=device)
+    for level_id, proposals_i, logits_i in zip(itertools.count(), proposals, pred_objectness_logits):
+        Hi_Wi_A = logits_i.shape[1]
+        num_proposals_i = min(pre_nms_topk, Hi_Wi_A)
+
+        # sort is faster than topk (https://github.com/pytorch/pytorch/issues/22812)
+        # topk_scores_i, topk_idx = logits_i.topk(num_proposals_i, dim=1)
+        logits_i, idx = logits_i.sort(descending=True, dim=1)
+        topk_scores_i = logits_i[batch_idx, :num_proposals_i]
+        topk_idx = idx[batch_idx, :num_proposals_i]
+
+        # each is N x topk
+        topk_proposals_i = proposals_i[batch_idx[:, None], topk_idx]  # N x topk x 4
+
+        topk_proposals.append(topk_proposals_i)
+        topk_scores.append(topk_scores_i)
+        level_ids.append(torch.full((num_proposals_i,), level_id, dtype=torch.int64, device=device))
+
+    # 2. Concat all levels together
+    topk_scores = torch.cat(topk_scores, dim=1)
+    topk_proposals = torch.cat(topk_proposals, dim=1)
+    level_ids = torch.cat(level_ids, dim=0)
+
+    # if I change to batched_nms, I wonder if this will make a difference
+    # 3. For each image, run a per-level NMS, and choose topk results.
+    results = []
+    for n, image_size in enumerate(image_sizes):
+        boxes = topk_proposals[n]
+        scores_per_img = topk_scores[n]
+        # I will have to take a look at the boxes clip method
+        _clip_box(boxes, image_size)
+        # filter empty boxes
+        keep = _nonempty_boxes(boxes, threshold=min_box_side_len)
+        lvl = level_ids
+        if keep.sum().item() != len(boxes):
+            boxes, scores_per_img, lvl = (
+                boxes[keep],
+                scores_per_img[keep],
+                level_ids[keep],
+            )
+
+        keep = batched_nms(boxes, scores_per_img, lvl, nms_thresh)
+        keep = keep[:post_nms_topk]
+
+        res = (boxes[keep], scores_per_img[keep])
+        results.append(res)
+
+    # I wonder if it would be possible for me to pad all these things.
+    return results
+
+
+def subsample_labels(labels, num_samples, positive_fraction, bg_label):
+    """
+    Returns:
+        pos_idx, neg_idx (Tensor):
+            1D vector of indices. The total length of both is `num_samples` or fewer.
+    """
+    positive = torch.nonzero((labels != -1) & (labels != bg_label)).squeeze(1)
+    negative = torch.nonzero(labels == bg_label).squeeze(1)
+
+    num_pos = int(num_samples * positive_fraction)
+    # protect against not enough positive examples
+    num_pos = min(positive.numel(), num_pos)
+    num_neg = num_samples - num_pos
+    # protect against not enough negative examples
+    num_neg = min(negative.numel(), num_neg)
+
+    # randomly select positive and negative examples
+    perm1 = torch.randperm(positive.numel(), device=positive.device)[:num_pos]
+    perm2 = torch.randperm(negative.numel(), device=negative.device)[:num_neg]
+
+    pos_idx = positive[perm1]
+    neg_idx = negative[perm2]
+    return pos_idx, neg_idx
+
+
+def add_ground_truth_to_proposals(gt_boxes, proposals):
+    raise NotImplementedError()
+
+
+def add_ground_truth_to_proposals_single_image(gt_boxes, proposals):
+    raise NotImplementedError()
+
+
+def _fmt_box_list(box_tensor, batch_index: int):
+    repeated_index = torch.full(
+        (len(box_tensor), 1),
+        batch_index,
+        dtype=box_tensor.dtype,
+        device=box_tensor.device,
+    )
+    return torch.cat((repeated_index, box_tensor), dim=1)
+
+
+def convert_boxes_to_pooler_format(box_lists: List[torch.Tensor]):
+    pooler_fmt_boxes = torch.cat(
+        [_fmt_box_list(box_list, i) for i, box_list in enumerate(box_lists)],
+        dim=0,
+    )
+    return pooler_fmt_boxes
+
+
+def assign_boxes_to_levels(
+    box_lists: List[torch.Tensor],
+    min_level: int,
+    max_level: int,
+    canonical_box_size: int,
+    canonical_level: int,
+):
+
+    box_sizes = torch.sqrt(torch.cat([boxes.area() for boxes in box_lists]))
+    # Eqn.(1) in FPN paper
+    level_assignments = torch.floor(canonical_level + torch.log2(box_sizes / canonical_box_size + 1e-8))
+    # clamp level to (min, max), in case the box size is too large or too small
+    # for the available feature maps
+    level_assignments = torch.clamp(level_assignments, min=min_level, max=max_level)
+    return level_assignments.to(torch.int64) - min_level
+
+
+# Helper Classes
+class _NewEmptyTensorOp(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, new_shape):
+        ctx.shape = x.shape
+        return x.new_empty(new_shape)
+
+    @staticmethod
+    def backward(ctx, grad):
+        shape = ctx.shape
+        return _NewEmptyTensorOp.apply(grad, shape), None
+
+
+class ShapeSpec(namedtuple("_ShapeSpec", ["channels", "height", "width", "stride"])):
+    def __new__(cls, *, channels=None, height=None, width=None, stride=None):
+        return super().__new__(cls, channels, height, width, stride)
+
+
+class Box2BoxTransform(object):
+    """
+    This R-CNN transformation scales the box's width and height
+    by exp(dw), exp(dh) and shifts a box's center by the offset
+    (dx * width, dy * height).
+    """
+
+    def __init__(self, weights: Tuple[float, float, float, float], scale_clamp: float = None):
+        """
+        Args:
+            weights (4-element tuple): Scaling factors that are applied to the
+                (dx, dy, dw, dh) deltas. In Fast R-CNN, these were originally set
+                such that the deltas have unit variance; now they are treated as
+                hyperparameters of the system.
+            scale_clamp (float): When predicting deltas, the predicted box scaling
+                factors (dw and dh) are clamped such that they are <= scale_clamp.
+        """
+        self.weights = weights
+        if scale_clamp is not None:
+            self.scale_clamp = scale_clamp
+        else:
+            """
+            Value for clamping large dw and dh predictions.
+            The heuristic is that we clamp such that dw and dh are no larger
+            than what would transform a 16px box into a 1000px box
+            (based on a small anchor, 16px, and a typical image size, 1000px).
+            """
+            self.scale_clamp = math.log(1000.0 / 16)
+
+    def get_deltas(self, src_boxes, target_boxes):
+        """
+        Get box regression transformation deltas (dx, dy, dw, dh) that can be used
+        to transform the `src_boxes` into the `target_boxes`. That is, the relation
+        ``target_boxes == self.apply_deltas(deltas, src_boxes)`` is true (unless
+        any delta is too large and is clamped).
+        Args:
+            src_boxes (Tensor): source boxes, e.g., object proposals
+            target_boxes (Tensor): target of the transformation, e.g., ground-truth
+                boxes.
+        """
+        assert isinstance(src_boxes, torch.Tensor), type(src_boxes)
+        assert isinstance(target_boxes, torch.Tensor), type(target_boxes)
+
+        src_widths = src_boxes[:, 2] - src_boxes[:, 0]
+        src_heights = src_boxes[:, 3] - src_boxes[:, 1]
+        src_ctr_x = src_boxes[:, 0] + 0.5 * src_widths
+        src_ctr_y = src_boxes[:, 1] + 0.5 * src_heights
+
+        target_widths = target_boxes[:, 2] - target_boxes[:, 0]
+        target_heights = target_boxes[:, 3] - target_boxes[:, 1]
+        target_ctr_x = target_boxes[:, 0] + 0.5 * target_widths
+        target_ctr_y = target_boxes[:, 1] + 0.5 * target_heights
+
+        wx, wy, ww, wh = self.weights
+        dx = wx * (target_ctr_x - src_ctr_x) / src_widths
+        dy = wy * (target_ctr_y - src_ctr_y) / src_heights
+        dw = ww * torch.log(target_widths / src_widths)
+        dh = wh * torch.log(target_heights / src_heights)
+
+        deltas = torch.stack((dx, dy, dw, dh), dim=1)
+        assert (src_widths > 0).all().item(), "Input boxes to Box2BoxTransform are not valid!"
+        return deltas
+
+    def apply_deltas(self, deltas, boxes):
+        """
+        Apply transformation `deltas` (dx, dy, dw, dh) to `boxes`.
+        Args:
+            deltas (Tensor): transformation deltas of shape (N, k*4), where k >= 1.
+                deltas[i] represents k potentially different class-specific
+                box transformations for the single box boxes[i].
+            boxes (Tensor): boxes to transform, of shape (N, 4)
+        """
+        boxes = boxes.to(deltas.dtype)
+
+        widths = boxes[:, 2] - boxes[:, 0]
+        heights = boxes[:, 3] - boxes[:, 1]
+        ctr_x = boxes[:, 0] + 0.5 * widths
+        ctr_y = boxes[:, 1] + 0.5 * heights
+
+        wx, wy, ww, wh = self.weights
+        dx = deltas[:, 0::4] / wx
+        dy = deltas[:, 1::4] / wy
+        dw = deltas[:, 2::4] / ww
+        dh = deltas[:, 3::4] / wh
+
+        # Prevent sending too large values into torch.exp()
+        dw = torch.clamp(dw, max=self.scale_clamp)
+        dh = torch.clamp(dh, max=self.scale_clamp)
+
+        pred_ctr_x = dx * widths[:, None] + ctr_x[:, None]
+        pred_ctr_y = dy * heights[:, None] + ctr_y[:, None]
+        pred_w = torch.exp(dw) * widths[:, None]
+        pred_h = torch.exp(dh) * heights[:, None]
+
+        pred_boxes = torch.zeros_like(deltas)
+        pred_boxes[:, 0::4] = pred_ctr_x - 0.5 * pred_w  # x1
+        pred_boxes[:, 1::4] = pred_ctr_y - 0.5 * pred_h  # y1
+        pred_boxes[:, 2::4] = pred_ctr_x + 0.5 * pred_w  # x2
+        pred_boxes[:, 3::4] = pred_ctr_y + 0.5 * pred_h  # y2
+        return pred_boxes
+
+
+class Matcher(object):
+    """
+    This class assigns to each predicted "element" (e.g., a box) a ground-truth
+    element. Each predicted element will have exactly zero or one matches; each
+    ground-truth element may be matched to zero or more predicted elements.
+    The matching is determined by the MxN match_quality_matrix, that characterizes
+    how well each (ground-truth, prediction)-pair match each other. For example,
+    if the elements are boxes, this matrix may contain box intersection-over-union
+    overlap values.
+    The matcher returns (a) a vector of length N containing the index of the
+    ground-truth element m in [0, M) that matches to prediction n in [0, N).
+    (b) a vector of length N containing the labels for each prediction.
+    """
+
+    def __init__(
+        self,
+        thresholds: List[float],
+        labels: List[int],
+        allow_low_quality_matches: bool = False,
+    ):
+        """
+        Args:
+            thresholds (list): a list of thresholds used to stratify predictions
+                into levels.
+            labels (list): a list of values to label predictions belonging at
+                each level. A label can be one of {-1, 0, 1} signifying
+                {ignore, negative class, positive class}, respectively.
+            allow_low_quality_matches (bool): if True, produce additional matches or predictions with maximum match quality lower than high_threshold.
+                For example, thresholds = [0.3, 0.5] labels = [0, -1, 1] All predictions with iou < 0.3 will be marked with 0 and
+                thus will be considered as false positives while training. All predictions with 0.3 <= iou < 0.5 will be marked with -1 and
+                thus will be ignored. All predictions with 0.5 <= iou will be marked with 1 and thus will be considered as true positives.
+        """
+        thresholds = thresholds[:]
+        assert thresholds[0] > 0
+        thresholds.insert(0, -float("inf"))
+        thresholds.append(float("inf"))
+        assert all([low <= high for (low, high) in zip(thresholds[:-1], thresholds[1:])])
+        assert all([label_i in [-1, 0, 1] for label_i in labels])
+        assert len(labels) == len(thresholds) - 1
+        self.thresholds = thresholds
+        self.labels = labels
+        self.allow_low_quality_matches = allow_low_quality_matches
+
+    def __call__(self, match_quality_matrix):
+        """
+        Args:
+            match_quality_matrix (Tensor[float]): an MxN tensor, containing the pairwise quality between M ground-truth elements and N predicted
+                elements. All elements must be >= 0 (due to the us of `torch.nonzero` for selecting indices in :meth:`set_low_quality_matches_`).
+        Returns:
+            matches (Tensor[int64]): a vector of length N, where matches[i] is a matched ground-truth index in [0, M)
+            match_labels (Tensor[int8]): a vector of length N, where pred_labels[i] indicates true or false positive or ignored
+        """
+        assert match_quality_matrix.dim() == 2
+        if match_quality_matrix.numel() == 0:
+            default_matches = match_quality_matrix.new_full((match_quality_matrix.size(1),), 0, dtype=torch.int64)
+            # When no gt boxes exist, we define IOU = 0 and therefore set labels
+            # to `self.labels[0]`, which usually defaults to background class 0
+            # To choose to ignore instead,
+            # can make labels=[-1,0,-1,1] + set appropriate thresholds
+            default_match_labels = match_quality_matrix.new_full(
+                (match_quality_matrix.size(1),), self.labels[0], dtype=torch.int8
+            )
+            return default_matches, default_match_labels
+
+        assert torch.all(match_quality_matrix >= 0)
+
+        # match_quality_matrix is M (gt) x N (predicted)
+        # Max over gt elements (dim 0) to find best gt candidate for each prediction
+        matched_vals, matches = match_quality_matrix.max(dim=0)
+
+        match_labels = matches.new_full(matches.size(), 1, dtype=torch.int8)
+
+        for (l, low, high) in zip(self.labels, self.thresholds[:-1], self.thresholds[1:]):
+            low_high = (matched_vals >= low) & (matched_vals < high)
+            match_labels[low_high] = l
+
+        if self.allow_low_quality_matches:
+            self.set_low_quality_matches_(match_labels, match_quality_matrix)
+
+        return matches, match_labels
+
+    def set_low_quality_matches_(self, match_labels, match_quality_matrix):
+        """
+        Produce additional matches for predictions that have only low-quality matches.
+        Specifically, for each ground-truth G find the set of predictions that have
+        maximum overlap with it (including ties); for each prediction in that set, if
+        it is unmatched, then match it to the ground-truth G.
+        This function implements the RPN assignment case (i)
+        in Sec. 3.1.2 of Faster R-CNN.
+        """
+        # For each gt, find the prediction with which it has highest quality
+        highest_quality_foreach_gt, _ = match_quality_matrix.max(dim=1)
+        # Find the highest quality match available, even if it is low, including ties.
+        # Note that the matches qualities must be positive due to the use of
+        # `torch.nonzero`.
+        of_quality_inds = match_quality_matrix == highest_quality_foreach_gt[:, None]
+        if of_quality_inds.dim() == 0:
+            (_, pred_inds_with_highest_quality) = of_quality_inds.unsqueeze(0).nonzero().unbind(1)
+        else:
+            (_, pred_inds_with_highest_quality) = of_quality_inds.nonzero().unbind(1)
+        match_labels[pred_inds_with_highest_quality] = 1
+
+
+class RPNOutputs(object):
+    def __init__(
+        self,
+        box2box_transform,
+        anchor_matcher,
+        batch_size_per_image,
+        positive_fraction,
+        images,
+        pred_objectness_logits,
+        pred_anchor_deltas,
+        anchors,
+        boundary_threshold=0,
+        gt_boxes=None,
+        smooth_l1_beta=0.0,
+    ):
+        """
+        Args:
+            box2box_transform (Box2BoxTransform): :class:`Box2BoxTransform` instance for anchor-proposal transformations.
+            anchor_matcher (Matcher): :class:`Matcher` instance for matching anchors to ground-truth boxes; used to determine training labels.
+            batch_size_per_image (int): number of proposals to sample when training
+            positive_fraction (float): target fraction of sampled proposals that should be positive
+            images (ImageList): :class:`ImageList` instance representing N input images
+            pred_objectness_logits (list[Tensor]): A list of L elements. Element i is a tensor of shape (N, A, Hi, W)
+            pred_anchor_deltas (list[Tensor]): A list of L elements. Element i is a tensor of shape (N, A*4, Hi, Wi)
+            anchors (list[torch.Tensor]): nested list of boxes. anchors[i][j] at (n, l) stores anchor array for feature map l
+            boundary_threshold (int): if >= 0, then anchors that extend beyond the image boundary by more than boundary_thresh are not used in training.
+            gt_boxes (list[Boxes], optional): A list of N elements.
+            smooth_l1_beta (float): The transition point between L1 and L2 lossn. When set to 0, the loss becomes L1. When +inf, it is ignored
+        """
+        self.box2box_transform = box2box_transform
+        self.anchor_matcher = anchor_matcher
+        self.batch_size_per_image = batch_size_per_image
+        self.positive_fraction = positive_fraction
+        self.pred_objectness_logits = pred_objectness_logits
+        self.pred_anchor_deltas = pred_anchor_deltas
+
+        self.anchors = anchors
+        self.gt_boxes = gt_boxes
+        self.num_feature_maps = len(pred_objectness_logits)
+        self.num_images = len(images)
+        self.boundary_threshold = boundary_threshold
+        self.smooth_l1_beta = smooth_l1_beta
+
+    def _get_ground_truth(self):
+        raise NotImplementedError()
+
+    def predict_proposals(self):
+        # pred_anchor_deltas: (L, N, ? Hi, Wi)
+        # anchors:(N, L, -1, B)
+        # here we loop over specific feature map, NOT images
+        proposals = []
+        anchors = self.anchors.transpose(0, 1)
+        for anchors_i, pred_anchor_deltas_i in zip(anchors, self.pred_anchor_deltas):
+            B = anchors_i.size(-1)
+            N, _, Hi, Wi = pred_anchor_deltas_i.shape
+            anchors_i = anchors_i.flatten(start_dim=0, end_dim=1)
+            pred_anchor_deltas_i = pred_anchor_deltas_i.view(N, -1, B, Hi, Wi).permute(0, 3, 4, 1, 2).reshape(-1, B)
+            proposals_i = self.box2box_transform.apply_deltas(pred_anchor_deltas_i, anchors_i)
+            # Append feature map proposals with shape (N, Hi*Wi*A, B)
+            proposals.append(proposals_i.view(N, -1, B))
+        proposals = torch.stack(proposals)
+        return proposals
+
+    def predict_objectness_logits(self):
+        """
+        Returns:
+            pred_objectness_logits (list[Tensor]) -> (N, Hi*Wi*A).
+        """
+        pred_objectness_logits = [
+            # Reshape: (N, A, Hi, Wi) -> (N, Hi, Wi, A) -> (N, Hi*Wi*A)
+            score.permute(0, 2, 3, 1).reshape(self.num_images, -1)
+            for score in self.pred_objectness_logits
+        ]
+        return pred_objectness_logits
+
+
+# Main Classes
+class Conv2d(nn.Conv2d):
+    def __init__(self, *args, **kwargs):
+        norm = kwargs.pop("norm", None)
+        activation = kwargs.pop("activation", None)
+        super().__init__(*args, **kwargs)
+
+        self.norm = norm
+        self.activation = activation
+
+    def forward(self, x):
+        if x.numel() == 0 and self.training:
+            assert not isinstance(self.norm, nn.SyncBatchNorm)
+        if x.numel() == 0:
+            assert not isinstance(self.norm, nn.GroupNorm)
+            output_shape = [
+                (i + 2 * p - (di * (k - 1) + 1)) // s + 1
+                for i, p, di, k, s in zip(
+                    x.shape[-2:],
+                    self.padding,
+                    self.dilation,
+                    self.kernel_size,
+                    self.stride,
+                )
+            ]
+            output_shape = [x.shape[0], self.weight.shape[0]] + output_shape
+            empty = _NewEmptyTensorOp.apply(x, output_shape)
+            if self.training:
+                _dummy = sum(x.view(-1)[0] for x in self.parameters()) * 0.0
+                return empty + _dummy
+            else:
+                return empty
+
+        x = super().forward(x)
+        if self.norm is not None:
+            x = self.norm(x)
+        if self.activation is not None:
+            x = self.activation(x)
+        return x
+
+
+class LastLevelMaxPool(nn.Module):
+    """
+    This module is used in the original FPN to generate a downsampled P6 feature from P5.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.num_levels = 1
+        self.in_feature = "p5"
+
+    def forward(self, x):
+        return [nn.functional.max_pool2d(x, kernel_size=1, stride=2, padding=0)]
+
+
+class LastLevelP6P7(nn.Module):
+    """
+    This module is used in RetinaNet to generate extra layers, P6 and P7 from C5 feature.
+    """
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.num_levels = 2
+        self.in_feature = "res5"
+        self.p6 = nn.Conv2d(in_channels, out_channels, 3, 2, 1)
+        self.p7 = nn.Conv2d(out_channels, out_channels, 3, 2, 1)
+
+    def forward(self, c5):
+        p6 = self.p6(c5)
+        p7 = self.p7(nn.functional.relu(p6))
+        return [p6, p7]
+
+
+class BasicStem(nn.Module):
+    def __init__(self, in_channels=3, out_channels=64, norm="BN", caffe_maxpool=False):
+        super().__init__()
+        self.conv1 = Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size=7,
+            stride=2,
+            padding=3,
+            bias=False,
+            norm=get_norm(norm, out_channels),
+        )
+        self.caffe_maxpool = caffe_maxpool
+        # use pad 1 instead of pad zero
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = nn.functional.relu_(x)
+        if self.caffe_maxpool:
+            x = nn.functional.max_pool2d(x, kernel_size=3, stride=2, padding=0, ceil_mode=True)
+        else:
+            x = nn.functional.max_pool2d(x, kernel_size=3, stride=2, padding=1)
+        return x
+
+    @property
+    def out_channels(self):
+        return self.conv1.out_channels
+
+    @property
+    def stride(self):
+        return 4  # = stride 2 conv -> stride 2 max pool
+
+
+class ResNetBlockBase(nn.Module):
+    def __init__(self, in_channels, out_channels, stride):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.stride = stride
+
+    def freeze(self):
+        for p in self.parameters():
+            p.requires_grad = False
+        return self
+
+
+class BottleneckBlock(ResNetBlockBase):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        bottleneck_channels,
+        stride=1,
+        num_groups=1,
+        norm="BN",
+        stride_in_1x1=False,
+        dilation=1,
+    ):
+        super().__init__(in_channels, out_channels, stride)
+
+        if in_channels != out_channels:
+            self.shortcut = Conv2d(
+                in_channels,
+                out_channels,
+                kernel_size=1,
+                stride=stride,
+                bias=False,
+                norm=get_norm(norm, out_channels),
+            )
+        else:
+            self.shortcut = None
+
+        # The original MSRA ResNet models have stride in the first 1x1 conv
+        # The subsequent fb.torch.resnet and Caffe2 ResNe[X]t implementations have
+        # stride in the 3x3 conv
+        stride_1x1, stride_3x3 = (stride, 1) if stride_in_1x1 else (1, stride)
+
+        self.conv1 = Conv2d(
+            in_channels,
+            bottleneck_channels,
+            kernel_size=1,
+            stride=stride_1x1,
+            bias=False,
+            norm=get_norm(norm, bottleneck_channels),
+        )
+
+        self.conv2 = Conv2d(
+            bottleneck_channels,
+            bottleneck_channels,
+            kernel_size=3,
+            stride=stride_3x3,
+            padding=1 * dilation,
+            bias=False,
+            groups=num_groups,
+            dilation=dilation,
+            norm=get_norm(norm, bottleneck_channels),
+        )
+
+        self.conv3 = Conv2d(
+            bottleneck_channels,
+            out_channels,
+            kernel_size=1,
+            bias=False,
+            norm=get_norm(norm, out_channels),
+        )
+
+    def forward(self, x):
+        out = self.conv1(x)
+        out = nn.functional.relu_(out)
+
+        out = self.conv2(out)
+        out = nn.functional.relu_(out)
+
+        out = self.conv3(out)
+
+        if self.shortcut is not None:
+            shortcut = self.shortcut(x)
+        else:
+            shortcut = x
+
+        out += shortcut
+        out = nn.functional.relu_(out)
+        return out
+
+
+class Backbone(nn.Module, metaclass=ABCMeta):
+    def __init__(self):
+        super().__init__()
+
+    @abstractmethod
+    def forward(self):
+        pass
+
+    @property
+    def size_divisibility(self):
+        """
+        Some backbones require the input height and width to be divisible by a specific integer. This is
+        typically true for encoder / decoder type networks with lateral connection (e.g., FPN) for which feature maps need to match
+        dimension in the "bottom up" and "top down" paths. Set to 0 if no specific input size divisibility is required.
+        """
+        return 0
+
+    def output_shape(self):
+        return {
+            name: ShapeSpec(
+                channels=self._out_feature_channels[name],
+                stride=self._out_feature_strides[name],
+            )
+            for name in self._out_features
+        }
+
+    @property
+    def out_features(self):
+        """deprecated"""
+        return self._out_features
+
+    @property
+    def out_feature_strides(self):
+        """deprecated"""
+        return {f: self._out_feature_strides[f] for f in self._out_features}
+
+    @property
+    def out_feature_channels(self):
+        """deprecated"""
+        return {f: self._out_feature_channels[f] for f in self._out_features}
+
+
+class ResNet(Backbone):
+    def __init__(self, stem, stages, num_classes=None, out_features=None):
+        """
+        Args:
+            stem (nn.Module): a stem module
+            stages (list[list[ResNetBlock]]): several (typically 4) stages, each contains multiple :class:`ResNetBlockBase`.
+            num_classes (None or int): if None, will not perform classification.
+            out_features (list[str]): name of the layers whose outputs should be returned in forward. Can be anything in:
+            "stem", "linear", or "res2" ... If None, will return the output of the last layer.
+        """
+        super(ResNet, self).__init__()
+        self.stem = stem
+        self.num_classes = num_classes
+
+        current_stride = self.stem.stride
+        self._out_feature_strides = {"stem": current_stride}
+        self._out_feature_channels = {"stem": self.stem.out_channels}
+
+        self.stages_and_names = []
+        for i, blocks in enumerate(stages):
+            for block in blocks:
+                assert isinstance(block, ResNetBlockBase), block
+                curr_channels = block.out_channels
+            stage = nn.Sequential(*blocks)
+            name = "res" + str(i + 2)
+            self.add_module(name, stage)
+            self.stages_and_names.append((stage, name))
+            self._out_feature_strides[name] = current_stride = int(
+                current_stride * np.prod([k.stride for k in blocks])
+            )
+            self._out_feature_channels[name] = blocks[-1].out_channels
+
+        if num_classes is not None:
+            self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+            self.linear = nn.Linear(curr_channels, num_classes)
+
+            # Sec 5.1 in "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour":
+            # "The 1000-way fully-connected layer is initialized by
+            # drawing weights from a zero-mean Gaussian with std of 0.01."
+            nn.init.normal_(self.linear.weight, stddev=0.01)
+            name = "linear"
+
+        if out_features is None:
+            out_features = [name]
+        self._out_features = out_features
+        assert len(self._out_features)
+        children = [x[0] for x in self.named_children()]
+        for out_feature in self._out_features:
+            assert out_feature in children, "Available children: {}".format(", ".join(children))
+
+    def forward(self, x):
+        outputs = {}
+        x = self.stem(x)
+        if "stem" in self._out_features:
+            outputs["stem"] = x
+        for stage, name in self.stages_and_names:
+            x = stage(x)
+            if name in self._out_features:
+                outputs[name] = x
+        if self.num_classes is not None:
+            x = self.avgpool(x)
+            x = self.linear(x)
+            if "linear" in self._out_features:
+                outputs["linear"] = x
+        return outputs
+
+    def output_shape(self):
+        return {
+            name: ShapeSpec(
+                channels=self._out_feature_channels[name],
+                stride=self._out_feature_strides[name],
+            )
+            for name in self._out_features
+        }
+
+    @staticmethod
+    def make_stage(
+        block_class,
+        num_blocks,
+        first_stride=None,
+        *,
+        in_channels,
+        out_channels,
+        **kwargs,
+    ):
+        """
+        Usually, layers that produce the same feature map spatial size
+        are defined as one "stage".
+        Under such definition, stride_per_block[1:] should all be 1.
+        """
+        if first_stride is not None:
+            assert "stride" not in kwargs and "stride_per_block" not in kwargs
+            kwargs["stride_per_block"] = [first_stride] + [1] * (num_blocks - 1)
+        blocks = []
+        for i in range(num_blocks):
+            curr_kwargs = {}
+            for k, v in kwargs.items():
+                if k.endswith("_per_block"):
+                    assert len(v) == num_blocks, (
+                        f"Argument '{k}' of make_stage should have the " f"same length as num_blocks={num_blocks}."
+                    )
+                    newk = k[: -len("_per_block")]
+                    assert newk not in kwargs, f"Cannot call make_stage with both {k} and {newk}!"
+                    curr_kwargs[newk] = v[i]
+                else:
+                    curr_kwargs[k] = v
+
+            blocks.append(block_class(in_channels=in_channels, out_channels=out_channels, **curr_kwargs))
+            in_channels = out_channels
+
+        return blocks
+
+
+class ROIPooler(nn.Module):
+    """
+    Region of interest feature map pooler that supports pooling from one or more
+    feature maps.
+    """
+
+    def __init__(
+        self,
+        output_size,
+        scales,
+        sampling_ratio,
+        canonical_box_size=224,
+        canonical_level=4,
+    ):
+        super().__init__()
+        # assumption that stride is a power of 2.
+        min_level = -math.log2(scales[0])
+        max_level = -math.log2(scales[-1])
+
+        # a bunch of testing
+        assert math.isclose(min_level, int(min_level)) and math.isclose(max_level, int(max_level))
+        assert len(scales) == max_level - min_level + 1, "not pyramid"
+        assert 0 < min_level and min_level <= max_level
+        if isinstance(output_size, int):
+            output_size = (output_size, output_size)
+        assert len(output_size) == 2 and isinstance(output_size[0], int) and isinstance(output_size[1], int)
+        if len(scales) > 1:
+            assert min_level <= canonical_level and canonical_level <= max_level
+        assert canonical_box_size > 0
+
+        self.output_size = output_size
+        self.min_level = int(min_level)
+        self.max_level = int(max_level)
+        self.level_poolers = nn.ModuleList(RoIPool(output_size, spatial_scale=scale) for scale in scales)
+        self.canonical_level = canonical_level
+        self.canonical_box_size = canonical_box_size
+
+    def forward(self, feature_maps, boxes):
+        """
+        Args:
+            feature_maps: List[torch.Tensor(N,C,W,H)]
+            box_lists: list[torch.Tensor])
+        Returns:
+            A tensor of shape(N*B, Channels, output_size, output_size)
+        """
+        x = [v for v in feature_maps.values()]
+        num_level_assignments = len(self.level_poolers)
+        assert len(x) == num_level_assignments and len(boxes) == x[0].size(0)
+
+        pooler_fmt_boxes = convert_boxes_to_pooler_format(boxes)
+
+        if num_level_assignments == 1:
+            return self.level_poolers[0](x[0], pooler_fmt_boxes)
+
+        level_assignments = assign_boxes_to_levels(
+            boxes,
+            self.min_level,
+            self.max_level,
+            self.canonical_box_size,
+            self.canonical_level,
+        )
+
+        num_boxes = len(pooler_fmt_boxes)
+        num_channels = x[0].shape[1]
+        output_size = self.output_size[0]
+
+        dtype, device = x[0].dtype, x[0].device
+        output = torch.zeros(
+            (num_boxes, num_channels, output_size, output_size),
+            dtype=dtype,
+            device=device,
+        )
+
+        for level, (x_level, pooler) in enumerate(zip(x, self.level_poolers)):
+            inds = torch.nonzero(level_assignments == level).squeeze(1)
+            pooler_fmt_boxes_level = pooler_fmt_boxes[inds]
+            output[inds] = pooler(x_level, pooler_fmt_boxes_level)
+
+        return output
+
+
+class ROIOutputs(object):
+    def __init__(self, cfg, training=False):
+        self.smooth_l1_beta = cfg.ROI_BOX_HEAD.SMOOTH_L1_BETA
+        self.box2box_transform = Box2BoxTransform(weights=cfg.ROI_BOX_HEAD.BBOX_REG_WEIGHTS)
+        self.training = training
+        self.score_thresh = cfg.ROI_HEADS.SCORE_THRESH_TEST
+        self.min_detections = cfg.MIN_DETECTIONS
+        self.max_detections = cfg.MAX_DETECTIONS
+
+        nms_thresh = cfg.ROI_HEADS.NMS_THRESH_TEST
+        if not isinstance(nms_thresh, list):
+            nms_thresh = [nms_thresh]
+        self.nms_thresh = nms_thresh
+
+    def _predict_boxes(self, proposals, box_deltas, preds_per_image):
+        num_pred = box_deltas.size(0)
+        B = proposals[0].size(-1)
+        K = box_deltas.size(-1) // B
+        box_deltas = box_deltas.view(num_pred * K, B)
+        proposals = torch.cat(proposals, dim=0).unsqueeze(-2).expand(num_pred, K, B)
+        proposals = proposals.reshape(-1, B)
+        boxes = self.box2box_transform.apply_deltas(box_deltas, proposals)
+        return boxes.view(num_pred, K * B).split(preds_per_image, dim=0)
+
+    def _predict_objs(self, obj_logits, preds_per_image):
+        probs = nn.functional.softmax(obj_logits, dim=-1)
+        probs = probs.split(preds_per_image, dim=0)
+        return probs
+
+    def _predict_attrs(self, attr_logits, preds_per_image):
+        attr_logits = attr_logits[..., :-1].softmax(-1)
+        attr_probs, attrs = attr_logits.max(-1)
+        return attr_probs.split(preds_per_image, dim=0), attrs.split(preds_per_image, dim=0)
+
+    @torch.no_grad()
+    def inference(
+        self,
+        obj_logits,
+        attr_logits,
+        box_deltas,
+        pred_boxes,
+        features,
+        sizes,
+        scales=None,
+    ):
+        # only the pred boxes is the
+        preds_per_image = [p.size(0) for p in pred_boxes]
+        boxes_all = self._predict_boxes(pred_boxes, box_deltas, preds_per_image)
+        obj_scores_all = self._predict_objs(obj_logits, preds_per_image)  # list of length N
+        attr_probs_all, attrs_all = self._predict_attrs(attr_logits, preds_per_image)
+        features = features.split(preds_per_image, dim=0)
+
+        # fun for each image too, also I can experiment and do multiple images
+        final_results = []
+        zipped = zip(boxes_all, obj_scores_all, attr_probs_all, attrs_all, sizes)
+        for i, (boxes, obj_scores, attr_probs, attrs, size) in enumerate(zipped):
+            for nms_t in self.nms_thresh:
+                outputs = do_nms(
+                    boxes,
+                    obj_scores,
+                    size,
+                    self.score_thresh,
+                    nms_t,
+                    self.min_detections,
+                    self.max_detections,
+                )
+                if outputs is not None:
+                    max_boxes, max_scores, classes, ids = outputs
+                    break
+
+            if scales is not None:
+                scale_yx = scales[i]
+                max_boxes[:, 0::2] *= scale_yx[1]
+                max_boxes[:, 1::2] *= scale_yx[0]
+
+            final_results.append(
+                (
+                    max_boxes,
+                    classes,
+                    max_scores,
+                    attrs[ids],
+                    attr_probs[ids],
+                    features[i][ids],
+                )
+            )
+        boxes, classes, class_probs, attrs, attr_probs, roi_features = map(list, zip(*final_results))
+        return boxes, classes, class_probs, attrs, attr_probs, roi_features
+
+    def training(self, obj_logits, attr_logits, box_deltas, pred_boxes, features, sizes):
+        pass
+
+    def __call__(
+        self,
+        obj_logits,
+        attr_logits,
+        box_deltas,
+        pred_boxes,
+        features,
+        sizes,
+        scales=None,
+    ):
+        if self.training:
+            raise NotImplementedError()
+        return self.inference(
+            obj_logits,
+            attr_logits,
+            box_deltas,
+            pred_boxes,
+            features,
+            sizes,
+            scales=scales,
+        )
+
+
+class Res5ROIHeads(nn.Module):
+    """
+    ROIHeads perform all per-region computation in an R-CNN.
+    It contains logic of cropping the regions, extract per-region features
+    (by the res-5 block in this case), and make per-region predictions.
+    """
+
+    def __init__(self, cfg, input_shape):
+        super().__init__()
+        self.batch_size_per_image = cfg.RPN.BATCH_SIZE_PER_IMAGE
+        self.positive_sample_fraction = cfg.ROI_HEADS.POSITIVE_FRACTION
+        self.in_features = cfg.ROI_HEADS.IN_FEATURES
+        self.num_classes = cfg.ROI_HEADS.NUM_CLASSES
+        self.proposal_append_gt = cfg.ROI_HEADS.PROPOSAL_APPEND_GT
+        self.feature_strides = {k: v.stride for k, v in input_shape.items()}
+        self.feature_channels = {k: v.channels for k, v in input_shape.items()}
+        self.cls_agnostic_bbox_reg = cfg.ROI_BOX_HEAD.CLS_AGNOSTIC_BBOX_REG
+        self.stage_channel_factor = 2 ** 3  # res5 is 8x res2
+        self.out_channels = cfg.RESNETS.RES2_OUT_CHANNELS * self.stage_channel_factor
+
+        # self.proposal_matcher = Matcher(
+        #     cfg.ROI_HEADS.IOU_THRESHOLDS,
+        #     cfg.ROI_HEADS.IOU_LABELS,
+        #     allow_low_quality_matches=False,
+        # )
+
+        pooler_resolution = cfg.ROI_BOX_HEAD.POOLER_RESOLUTION
+        pooler_scales = (1.0 / self.feature_strides[self.in_features[0]],)
+        sampling_ratio = cfg.ROI_BOX_HEAD.POOLER_SAMPLING_RATIO
+        res5_halve = cfg.ROI_BOX_HEAD.RES5HALVE
+        use_attr = cfg.ROI_BOX_HEAD.ATTR
+        num_attrs = cfg.ROI_BOX_HEAD.NUM_ATTRS
+
+        self.pooler = ROIPooler(
+            output_size=pooler_resolution,
+            scales=pooler_scales,
+            sampling_ratio=sampling_ratio,
+        )
+
+        self.res5 = self._build_res5_block(cfg)
+        if not res5_halve:
+            """
+            Modifications for VG in RoI heads:
+            1. Change the stride of conv1 and shortcut in Res5.Block1 from 2 to 1
+            2. Modifying all conv2 with (padding: 1 --> 2) and (dilation: 1 --> 2)
+            """
+            self.res5[0].conv1.stride = (1, 1)
+            self.res5[0].shortcut.stride = (1, 1)
+            for i in range(3):
+                self.res5[i].conv2.padding = (2, 2)
+                self.res5[i].conv2.dilation = (2, 2)
+
+        self.box_predictor = FastRCNNOutputLayers(
+            self.out_channels,
+            self.num_classes,
+            self.cls_agnostic_bbox_reg,
+            use_attr=use_attr,
+            num_attrs=num_attrs,
+        )
+
+    def _build_res5_block(self, cfg):
+        stage_channel_factor = self.stage_channel_factor  # res5 is 8x res2
+        num_groups = cfg.RESNETS.NUM_GROUPS
+        width_per_group = cfg.RESNETS.WIDTH_PER_GROUP
+        bottleneck_channels = num_groups * width_per_group * stage_channel_factor
+        out_channels = self.out_channels
+        stride_in_1x1 = cfg.RESNETS.STRIDE_IN_1X1
+        norm = cfg.RESNETS.NORM
+
+        blocks = ResNet.make_stage(
+            BottleneckBlock,
+            3,
+            first_stride=2,
+            in_channels=out_channels // 2,
+            bottleneck_channels=bottleneck_channels,
+            out_channels=out_channels,
+            num_groups=num_groups,
+            norm=norm,
+            stride_in_1x1=stride_in_1x1,
+        )
+        return nn.Sequential(*blocks)
+
+    def _shared_roi_transform(self, features, boxes):
+        x = self.pooler(features, boxes)
+        return self.res5(x)
+
+    def forward(self, features, proposal_boxes, gt_boxes=None):
+        if self.training:
+            """
+            see https://github.com/airsplay/py-bottom-up-attention/\
+                    blob/master/detectron2/modeling/roi_heads/roi_heads.py
+            """
+            raise NotImplementedError()
+
+        assert not proposal_boxes[0].requires_grad
+        box_features = self._shared_roi_transform(features, proposal_boxes)
+        feature_pooled = box_features.mean(dim=[2, 3])  # pooled to 1x1
+        obj_logits, attr_logits, pred_proposal_deltas = self.box_predictor(feature_pooled)
+        return obj_logits, attr_logits, pred_proposal_deltas, feature_pooled
+
+
+class AnchorGenerator(nn.Module):
+    """
+    For a set of image sizes and feature maps, computes a set of anchors.
+    """
+
+    def __init__(self, cfg, input_shape: List[ShapeSpec]):
+        super().__init__()
+        sizes = cfg.ANCHOR_GENERATOR.SIZES
+        aspect_ratios = cfg.ANCHOR_GENERATOR.ASPECT_RATIOS
+        self.strides = [x.stride for x in input_shape]
+        self.offset = cfg.ANCHOR_GENERATOR.OFFSET
+        assert 0.0 <= self.offset < 1.0, self.offset
+
+        """
+        sizes (list[list[int]]): sizes[i] is the list of anchor sizes for feat map i
+            1. given in absolute lengths in units of the input image;
+            2. they do not dynamically scale if the input image size changes.
+        aspect_ratios (list[list[float]])
+        strides (list[int]): stride of each input feature.
+        """
+
+        self.num_features = len(self.strides)
+        self.cell_anchors = nn.ParameterList(self._calculate_anchors(sizes, aspect_ratios))
+        self._spacial_feat_dim = 4
+
+    def _calculate_anchors(self, sizes, aspect_ratios):
+        # If one size (or aspect ratio) is specified and there are multiple feature
+        # maps, then we "broadcast" anchors of that single size (or aspect ratio)
+        if len(sizes) == 1:
+            sizes *= self.num_features
+        if len(aspect_ratios) == 1:
+            aspect_ratios *= self.num_features
+        assert self.num_features == len(sizes)
+        assert self.num_features == len(aspect_ratios)
+
+        cell_anchors = [self.generate_cell_anchors(s, a).float() for s, a in zip(sizes, aspect_ratios)]
+
+        return cell_anchors
+
+    @property
+    def box_dim(self):
+        return self._spacial_feat_dim
+
+    @property
+    def num_cell_anchors(self):
+        """
+        Returns:
+            list[int]: Each int is the number of anchors at every pixel location, on that feature map.
+        """
+        return [len(cell_anchors) for cell_anchors in self.cell_anchors]
+
+    def grid_anchors(self, grid_sizes):
+        anchors = []
+        for (size, stride, base_anchors) in zip(grid_sizes, self.strides, self.cell_anchors):
+            shift_x, shift_y = _create_grid_offsets(size, stride, self.offset, base_anchors.device)
+            shifts = torch.stack((shift_x, shift_y, shift_x, shift_y), dim=1)
+
+            anchors.append((shifts.view(-1, 1, 4) + base_anchors.view(1, -1, 4)).reshape(-1, 4))
+
+        return anchors
+
+    def generate_cell_anchors(self, sizes=(32, 64, 128, 256, 512), aspect_ratios=(0.5, 1, 2)):
+        """
+        anchors are continuous geometric rectangles
+        centered on one feature map point sample.
+        We can later build the set of anchors
+        for the entire feature map by tiling these tensors
+        """
+
+        anchors = []
+        for size in sizes:
+            area = size ** 2.0
+            for aspect_ratio in aspect_ratios:
+                w = math.sqrt(area / aspect_ratio)
+                h = aspect_ratio * w
+                x0, y0, x1, y1 = -w / 2.0, -h / 2.0, w / 2.0, h / 2.0
+                anchors.append([x0, y0, x1, y1])
+        return nn.Parameter(torch.tensor(anchors))
+
+    def forward(self, features):
+        """
+        Args:
+            features List[torch.Tensor]: list of feature maps on which to generate anchors.
+        Returns:
+            torch.Tensor: a list of #image elements.
+        """
+        num_images = features[0].size(0)
+        grid_sizes = [feature_map.shape[-2:] for feature_map in features]
+        anchors_over_all_feature_maps = self.grid_anchors(grid_sizes)
+        anchors_over_all_feature_maps = torch.stack(anchors_over_all_feature_maps)
+        return anchors_over_all_feature_maps.unsqueeze(0).repeat_interleave(num_images, dim=0)
+
+
+class RPNHead(nn.Module):
+    """
+    RPN classification and regression heads. Uses a 3x3 conv to produce a shared
+    hidden state from which one 1x1 conv predicts objectness logits for each anchor
+    and a second 1x1 conv predicts bounding-box deltas specifying how to deform
+    each anchor into an object proposal.
+    """
+
+    def __init__(self, cfg, input_shape: List[ShapeSpec]):
+        super().__init__()
+
+        # Standard RPN is shared across levels:
+        in_channels = [s.channels for s in input_shape]
+        assert len(set(in_channels)) == 1, "Each level must have the same channel!"
+        in_channels = in_channels[0]
+
+        anchor_generator = AnchorGenerator(cfg, input_shape)
+        num_cell_anchors = anchor_generator.num_cell_anchors
+        box_dim = anchor_generator.box_dim
+        assert len(set(num_cell_anchors)) == 1, "Each level must have the same number of cell anchors"
+        num_cell_anchors = num_cell_anchors[0]
+
+        if cfg.PROPOSAL_GENERATOR.HIDDEN_CHANNELS == -1:
+            hid_channels = in_channels
+        else:
+            hid_channels = cfg.PROPOSAL_GENERATOR.HIDDEN_CHANNELS
+            # Modifications for VG in RPN (modeling/proposal_generator/rpn.py)
+            # Use hidden dim  instead fo the same dim as Res4 (in_channels)
+
+        # 3x3 conv for the hidden representation
+        self.conv = nn.Conv2d(in_channels, hid_channels, kernel_size=3, stride=1, padding=1)
+        # 1x1 conv for predicting objectness logits
+        self.objectness_logits = nn.Conv2d(hid_channels, num_cell_anchors, kernel_size=1, stride=1)
+        # 1x1 conv for predicting box2box transform deltas
+        self.anchor_deltas = nn.Conv2d(hid_channels, num_cell_anchors * box_dim, kernel_size=1, stride=1)
+
+        for layer in [self.conv, self.objectness_logits, self.anchor_deltas]:
+            nn.init.normal_(layer.weight, std=0.01)
+            nn.init.constant_(layer.bias, 0)
+
+    def forward(self, features):
+        """
+        Args:
+            features (list[Tensor]): list of feature maps
+        """
+        pred_objectness_logits = []
+        pred_anchor_deltas = []
+        for x in features:
+            t = nn.functional.relu(self.conv(x))
+            pred_objectness_logits.append(self.objectness_logits(t))
+            pred_anchor_deltas.append(self.anchor_deltas(t))
+        return pred_objectness_logits, pred_anchor_deltas
+
+
+class RPN(nn.Module):
+    """
+    Region Proposal Network, introduced by the Faster R-CNN paper.
+    """
+
+    def __init__(self, cfg, input_shape: Dict[str, ShapeSpec]):
+        super().__init__()
+
+        self.min_box_side_len = cfg.PROPOSAL_GENERATOR.MIN_SIZE
+        self.in_features = cfg.RPN.IN_FEATURES
+        self.nms_thresh = cfg.RPN.NMS_THRESH
+        self.batch_size_per_image = cfg.RPN.BATCH_SIZE_PER_IMAGE
+        self.positive_fraction = cfg.RPN.POSITIVE_FRACTION
+        self.smooth_l1_beta = cfg.RPN.SMOOTH_L1_BETA
+        self.loss_weight = cfg.RPN.LOSS_WEIGHT
+
+        self.pre_nms_topk = {
+            True: cfg.RPN.PRE_NMS_TOPK_TRAIN,
+            False: cfg.RPN.PRE_NMS_TOPK_TEST,
+        }
+        self.post_nms_topk = {
+            True: cfg.RPN.POST_NMS_TOPK_TRAIN,
+            False: cfg.RPN.POST_NMS_TOPK_TEST,
+        }
+        self.boundary_threshold = cfg.RPN.BOUNDARY_THRESH
+
+        self.anchor_generator = AnchorGenerator(cfg, [input_shape[f] for f in self.in_features])
+        self.box2box_transform = Box2BoxTransform(weights=cfg.RPN.BBOX_REG_WEIGHTS)
+        self.anchor_matcher = Matcher(
+            cfg.RPN.IOU_THRESHOLDS,
+            cfg.RPN.IOU_LABELS,
+            allow_low_quality_matches=True,
+        )
+        self.rpn_head = RPNHead(cfg, [input_shape[f] for f in self.in_features])
+
+    def training(self, images, image_shapes, features, gt_boxes):
+        pass
+
+    def inference(self, outputs, images, image_shapes, features, gt_boxes=None):
+        outputs = find_top_rpn_proposals(
+            outputs.predict_proposals(),
+            outputs.predict_objectness_logits(),
+            images,
+            image_shapes,
+            self.nms_thresh,
+            self.pre_nms_topk[self.training],
+            self.post_nms_topk[self.training],
+            self.min_box_side_len,
+            self.training,
+        )
+
+        results = []
+        for img in outputs:
+            im_boxes, img_box_logits = img
+            img_box_logits, inds = img_box_logits.sort(descending=True)
+            im_boxes = im_boxes[inds]
+            results.append((im_boxes, img_box_logits))
+
+        (proposal_boxes, logits) = tuple(map(list, zip(*results)))
+        return proposal_boxes, logits
+
+    def forward(self, images, image_shapes, features, gt_boxes=None):
+        """
+        Args:
+            images (torch.Tensor): input images of length `N`
+            features (dict[str: Tensor])
+            gt_instances
+        """
+        # features is dict, key = block level, v = feature_map
+        features = [features[f] for f in self.in_features]
+        pred_objectness_logits, pred_anchor_deltas = self.rpn_head(features)
+        anchors = self.anchor_generator(features)
+        outputs = RPNOutputs(
+            self.box2box_transform,
+            self.anchor_matcher,
+            self.batch_size_per_image,
+            self.positive_fraction,
+            images,
+            pred_objectness_logits,
+            pred_anchor_deltas,
+            anchors,
+            self.boundary_threshold,
+            gt_boxes,
+            self.smooth_l1_beta,
+        )
+        # For RPN-only models, the proposals are the final output
+
+        if self.training:
+            raise NotImplementedError()
+            return self.training(outputs, images, image_shapes, features, gt_boxes)
+        else:
+            return self.inference(outputs, images, image_shapes, features, gt_boxes)
+
+
+class FastRCNNOutputLayers(nn.Module):
+    """
+    Two linear layers for predicting Fast R-CNN outputs:
+      (1) proposal-to-detection box regression deltas
+      (2) classification scores
+    """
+
+    def __init__(
+        self,
+        input_size,
+        num_classes,
+        cls_agnostic_bbox_reg,
+        box_dim=4,
+        use_attr=False,
+        num_attrs=-1,
+    ):
+        """
+        Args:
+            input_size (int): channels, or (channels, height, width)
+            num_classes (int)
+            cls_agnostic_bbox_reg (bool)
+            box_dim (int)
+        """
+        super().__init__()
+
+        if not isinstance(input_size, int):
+            input_size = np.prod(input_size)
+
+        # (do + 1 for background class)
+        self.cls_score = nn.Linear(input_size, num_classes + 1)
+        num_bbox_reg_classes = 1 if cls_agnostic_bbox_reg else num_classes
+        self.bbox_pred = nn.Linear(input_size, num_bbox_reg_classes * box_dim)
+
+        self.use_attr = use_attr
+        if use_attr:
+            """
+            Modifications for VG in RoI heads
+            Embedding: {num_classes + 1} --> {input_size // 8}
+            Linear: {input_size + input_size // 8} --> {input_size // 4}
+            Linear: {input_size // 4} --> {num_attrs + 1}
+            """
+            self.cls_embedding = nn.Embedding(num_classes + 1, input_size // 8)
+            self.fc_attr = nn.Linear(input_size + input_size // 8, input_size // 4)
+            self.attr_score = nn.Linear(input_size // 4, num_attrs + 1)
+
+        nn.init.normal_(self.cls_score.weight, std=0.01)
+        nn.init.normal_(self.bbox_pred.weight, std=0.001)
+        for item in [self.cls_score, self.bbox_pred]:
+            nn.init.constant_(item.bias, 0)
+
+    def forward(self, roi_features):
+        if roi_features.dim() > 2:
+            roi_features = torch.flatten(roi_features, start_dim=1)
+        scores = self.cls_score(roi_features)
+        proposal_deltas = self.bbox_pred(roi_features)
+        if self.use_attr:
+            _, max_class = scores.max(-1)  # [b, c] --> [b]
+            cls_emb = self.cls_embedding(max_class)  # [b] --> [b, 256]
+            roi_features = torch.cat([roi_features, cls_emb], -1)  # [b, 2048] + [b, 256] --> [b, 2304]
+            roi_features = self.fc_attr(roi_features)
+            roi_features = nn.functional.relu(roi_features)
+            attr_scores = self.attr_score(roi_features)
+            return scores, attr_scores, proposal_deltas
+        else:
+            return scores, proposal_deltas
+
+
+class GeneralizedRCNN(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+
+        self.backbone = build_backbone(cfg)
+        self.proposal_generator = RPN(cfg, self.backbone.output_shape())
+        self.roi_heads = Res5ROIHeads(cfg, self.backbone.output_shape())
+        self.roi_outputs = ROIOutputs(cfg)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        config = kwargs.pop("config", None)
+        state_dict = kwargs.pop("state_dict", None)
+        cache_dir = kwargs.pop("cache_dir", None)
+        from_tf = kwargs.pop("from_tf", False)
+        force_download = kwargs.pop("force_download", False)
+        resume_download = kwargs.pop("resume_download", False)
+        proxies = kwargs.pop("proxies", None)
+        local_files_only = kwargs.pop("local_files_only", False)
+        use_cdn = kwargs.pop("use_cdn", True)
+
+        # Load config if we don't provide a configuration
+        if not isinstance(config, Config):
+            config_path = config if config is not None else pretrained_model_name_or_path
+            # try:
+            config = Config.from_pretrained(
+                config_path,
+                cache_dir=cache_dir,
+                force_download=force_download,
+                resume_download=resume_download,
+                proxies=proxies,
+                local_files_only=local_files_only,
+            )
+
+        # Load model
+        if pretrained_model_name_or_path is not None:
+            if os.path.isdir(pretrained_model_name_or_path):
+                if os.path.isfile(os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)):
+                    # Load from a PyTorch checkpoint
+                    archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)
+                else:
+                    raise EnvironmentError(
+                        "Error no file named {} found in directory {} ".format(
+                            WEIGHTS_NAME,
+                            pretrained_model_name_or_path,
+                        )
+                    )
+            elif os.path.isfile(pretrained_model_name_or_path) or is_remote_url(pretrained_model_name_or_path):
+                archive_file = pretrained_model_name_or_path
+            elif os.path.isfile(pretrained_model_name_or_path + ".index"):
+                assert (
+                    from_tf
+                ), "We found a TensorFlow checkpoint at {}, please set from_tf to True to load from this checkpoint".format(
+                    pretrained_model_name_or_path + ".index"
+                )
+                archive_file = pretrained_model_name_or_path + ".index"
+            else:
+                archive_file = hf_bucket_url(
+                    pretrained_model_name_or_path,
+                    filename=WEIGHTS_NAME,
+                    use_cdn=use_cdn,
+                )
+
+            try:
+                # Load from URL or cache if already cached
+                resolved_archive_file = cached_path(
+                    archive_file,
+                    cache_dir=cache_dir,
+                    force_download=force_download,
+                    proxies=proxies,
+                    resume_download=resume_download,
+                    local_files_only=local_files_only,
+                )
+                if resolved_archive_file is None:
+                    raise EnvironmentError
+            except EnvironmentError:
+                msg = f"Can't load weights for '{pretrained_model_name_or_path}'."
+                raise EnvironmentError(msg)
+
+            if resolved_archive_file == archive_file:
+                print("loading weights file {}".format(archive_file))
+            else:
+                print("loading weights file {} from cache at {}".format(archive_file, resolved_archive_file))
+        else:
+            resolved_archive_file = None
+
+        # Instantiate model.
+        model = cls(config)
+
+        if state_dict is None:
+            try:
+                try:
+                    state_dict = torch.load(resolved_archive_file, map_location="cpu")
+                except Exception:
+                    state_dict = load_checkpoint(resolved_archive_file)
+
+            except Exception:
+                raise OSError(
+                    "Unable to load weights from pytorch checkpoint file. "
+                    "If you tried to load a PyTorch model from a TF 2.0 checkpoint, please set from_tf=True. "
+                )
+
+        missing_keys = []
+        unexpected_keys = []
+        error_msgs = []
+
+        # Convert old format to new format if needed from a PyTorch state_dict
+        old_keys = []
+        new_keys = []
+        for key in state_dict.keys():
+            new_key = None
+            if "gamma" in key:
+                new_key = key.replace("gamma", "weight")
+            if "beta" in key:
+                new_key = key.replace("beta", "bias")
+            if new_key:
+                old_keys.append(key)
+                new_keys.append(new_key)
+        for old_key, new_key in zip(old_keys, new_keys):
+            state_dict[new_key] = state_dict.pop(old_key)
+
+        # copy state_dict so _load_from_state_dict can modify it
+        metadata = getattr(state_dict, "_metadata", None)
+        state_dict = state_dict.copy()
+        if metadata is not None:
+            state_dict._metadata = metadata
+
+        model_to_load = model
+        model_to_load.load_state_dict(state_dict)
+
+        if model.__class__.__name__ != model_to_load.__class__.__name__:
+            base_model_state_dict = model_to_load.state_dict().keys()
+            head_model_state_dict_without_base_prefix = [
+                key.split(cls.base_model_prefix + ".")[-1] for key in model.state_dict().keys()
+            ]
+            missing_keys.extend(head_model_state_dict_without_base_prefix - base_model_state_dict)
+
+        if len(unexpected_keys) > 0:
+            print(
+                f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when "
+                f"initializing {model.__class__.__name__}: {unexpected_keys}\n"
+                f"- This IS expected if you are initializing {model.__class__.__name__} from the checkpoint of a model trained on another task "
+                f"or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n"
+                f"- This IS NOT expected if you are initializing {model.__class__.__name__} from the checkpoint of a model that you expect "
+                f"to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model)."
+            )
+        else:
+            print(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n")
+        if len(missing_keys) > 0:
+            print(
+                f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at {pretrained_model_name_or_path} "
+                f"and are newly initialized: {missing_keys}\n"
+                f"You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference."
+            )
+        else:
+            print(
+                f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at {pretrained_model_name_or_path}.\n"
+                f"If your task is similar to the task the model of the checkpoint was trained on, "
+                f"you can already use {model.__class__.__name__} for predictions without further training."
+            )
+        if len(error_msgs) > 0:
+            raise RuntimeError(
+                "Error(s) in loading state_dict for {}:\n\t{}".format(
+                    model.__class__.__name__, "\n\t".join(error_msgs)
+                )
+            )
+        # Set model in evaluation mode to deactivate DropOut modules by default
+        model.eval()
+        return model
+
+    def forward(
+        self,
+        images,
+        image_shapes,
+        gt_boxes=None,
+        proposals=None,
+        scales_yx=None,
+        **kwargs,
+    ):
+        """
+        kwargs:
+            max_detections (int), return_tensors {"np", "pt", None}, padding {None,
+            "max_detections"}, pad_value (int), location = {"cuda", "cpu"}
+        """
+        data = next(self.parameters()).data
+        with torch.no_grad():
+         if self.training:
+             print ("warning. you are attempting to train the frcnn model which is not supportd. switching to eval mode")
+             self.eval()
+             for param in self.parameters():
+               param.requires_grad_(False)
+         #print (image_shapes.dtype)
+         return self.inference(
+             images=images.to(dtype=data.dtype, device=data.device),
+             image_shapes=image_shapes.to(device=data.device),
+             gt_boxes=gt_boxes.to(dtype=data.dtype, device=data.device) if gt_boxes is not None else None,
+             proposals=proposals.to(dtype=data.dtype, device=data.device) if proposals is not None else None,
+             scales_yx=scales_yx.to(dtype=data.dtype, device=data.device) if scales_yx is not None else None,
+             **kwargs,
+         )
+
+    @torch.no_grad()
+    def inference(
+        self,
+        images,
+        image_shapes,
+        gt_boxes=None,
+        proposals=None,
+        scales_yx=None,
+        **kwargs,
+    ):
+        # run images through backbone
+        original_sizes = image_shapes * scales_yx
+        features = self.backbone(images)
+
+        # generate proposals if none are available
+        if proposals is None:
+            proposal_boxes, _ = self.proposal_generator(images, image_shapes, features, gt_boxes)
+        else:
+            assert proposals is not None
+
+        # pool object features from either gt_boxes, or from proposals
+        obj_logits, attr_logits, box_deltas, feature_pooled = self.roi_heads(features, proposal_boxes, gt_boxes)
+
+        # prepare FRCNN Outputs and select top proposals
+        boxes, classes, class_probs, attrs, attr_probs, roi_features = self.roi_outputs(
+            obj_logits=obj_logits,
+            attr_logits=attr_logits,
+            box_deltas=box_deltas,
+            pred_boxes=proposal_boxes,
+            features=feature_pooled,
+            sizes=image_shapes,
+            scales=scales_yx,
+        )
+
+        # will we pad???
+        subset_kwargs = {
+            "max_detections": kwargs.get("max_detections", None),
+            "return_tensors": kwargs.get("return_tensors", None),
+            "pad_value": kwargs.get("pad_value", 0),
+            "padding": kwargs.get("padding", None),
+        }
+        preds_per_image = torch.tensor([p.size(0) for p in boxes])
+        boxes = pad_list_tensors(boxes, preds_per_image, **subset_kwargs)
+        classes = pad_list_tensors(classes, preds_per_image, **subset_kwargs)
+        class_probs = pad_list_tensors(class_probs, preds_per_image, **subset_kwargs)
+        attrs = pad_list_tensors(attrs, preds_per_image, **subset_kwargs)
+        attr_probs = pad_list_tensors(attr_probs, preds_per_image, **subset_kwargs)
+        roi_features = pad_list_tensors(roi_features, preds_per_image, **subset_kwargs)
+        subset_kwargs["padding"] = None
+        preds_per_image = pad_list_tensors(preds_per_image, None, **subset_kwargs)
+        sizes = pad_list_tensors(image_shapes, None, **subset_kwargs)
+        #print (boxes.device, original_sizes.device)
+        normalized_boxes = norm_box(boxes, original_sizes.to(boxes.device))
+        return OrderedDict(
+            {
+                "obj_ids": classes,
+                "obj_probs": class_probs,
+                "attr_ids": attrs,
+                "attr_probs": attr_probs,
+                "boxes": boxes,
+                "sizes": sizes,
+                "preds_per_image": preds_per_image,
+                "roi_features": roi_features,
+                "normalized_boxes": normalized_boxes,
+            }
+        )