modeling_frcnn.py
0164de5c5cd4main

modeling_frcnn.py
View Options

			"""
			coding=utf-8
			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 import functional as F
			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]
			normalized_boxes[:, :, (1, 3)] /= raw_sizes[:, 0]
			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 = F.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, HiWiA, 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, HiWiA, 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, HiWiA).
			"""
			pred_objectness_logits = [
			# Reshape: (N, A, Hi, Wi) -> (N, Hi, Wi, A) -> (N, HiWiA)
			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(torch.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, torch.nn.SyncBatchNorm)
			if x.numel() == 0:
			assert not isinstance(self.norm, torch.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 [F.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(F.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 = F.relu_(x)
			if self.caffe_maxpool:
			x = F.max_pool2d(x, kernel_size=3, stride=2, padding=0, ceil_mode=True)
			else:
			x = F.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 = F.relu_(out)

			out = self.conv2(out)
			out = F.relu_(out)

			out = self.conv3(out)

			if self.shortcut is not None:
			shortcut = self.shortcut(x)
			else:
			shortcut = x

			out += shortcut
			out = F.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 = F.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 = F.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 = F.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.device = torch.device(cfg.MODEL.DEVICE)
			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)
			self.to(self.device)

			@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"}
			"""
			if self.training:
			raise NotImplementedError()
			return self.inference(
			images=images,
			image_shapes=image_shapes,
			gt_boxes=gt_boxes,
			proposals=proposals,
			scales_yx=scales_yx,
			**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)
			normalized_boxes = norm_box(boxes, original_sizes)
			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,
			}
			)

R11484/source_code/transformers/examples/research_projects/lxmert/modeling_frcnn.py0164de5c5cd4main

modeling_frcnn.pyView Options

R11484/source_code/transformers/examples/research_projects/lxmert/modeling_frcnn.py
0164de5c5cd4main

modeling_frcnn.py
View Options