"""Create mmcv.ops compatibility shim using torchvision ops.""" import os import mmcv ops_dir = os.path.join(os.path.dirname(mmcv.__file__), 'ops') os.makedirs(ops_dir, exist_ok=True) init_code = ''' import torch import torchvision.ops as tv_ops class NMSop(torch.autograd.Function): @staticmethod def forward(ctx, bboxes, scores, iou_threshold, offset, score_threshold, max_num): if score_threshold > 0: valid_mask = scores > score_threshold bboxes_f, scores_f = bboxes[valid_mask], scores[valid_mask] valid_inds = torch.nonzero(valid_mask, as_tuple=False).squeeze(dim=1) else: bboxes_f, scores_f = bboxes, scores valid_inds = None if bboxes_f.numel() == 0: return torch.zeros(0, dtype=torch.long, device=bboxes.device) inds = tv_ops.nms(bboxes_f, scores_f, iou_threshold) if max_num > 0: inds = inds[:max_num] if valid_inds is not None: inds = valid_inds[inds] return inds def batched_nms(boxes, scores, idxs, nms_cfg, class_agnostic=False): nms_cfg_ = nms_cfg.copy() class_agnostic = nms_cfg_.pop("class_agnostic", class_agnostic) iou_thr = nms_cfg_.get("iou_threshold", nms_cfg_.get("iou_thr", 0.5)) if class_agnostic: boxes_for_nms = boxes else: if boxes.numel() == 0: return boxes.new_zeros((0, 5)), torch.zeros(0, dtype=torch.long, device=boxes.device) max_coordinate = boxes.max() offsets = idxs.to(boxes) * (max_coordinate + 1) boxes_for_nms = boxes + offsets[:, None] if boxes_for_nms.numel() == 0: return boxes.new_zeros((0, 5)), torch.zeros(0, dtype=torch.long, device=boxes.device) keep = tv_ops.nms(boxes_for_nms, scores, iou_thr) max_num = nms_cfg_.get("max_num", -1) if max_num > 0 and len(keep) > max_num: keep = keep[:max_num] dets = torch.cat([boxes[keep], scores[keep].unsqueeze(1)], dim=1) return dets, keep def nms(boxes, scores, iou_threshold, offset=0, score_threshold=0, max_num=-1): inds = NMSop.apply(boxes, scores, iou_threshold, offset, score_threshold, max_num) dets = torch.cat([boxes[inds], scores[inds].unsqueeze(1)], dim=1) return dets, inds def roi_align(input, rois, output_size, spatial_scale=1.0, sampling_ratio=-1, pool_mode="avg", aligned=True): return tv_ops.roi_align(input, rois, output_size, spatial_scale=spatial_scale, sampling_ratio=sampling_ratio if sampling_ratio > 0 else 2, aligned=aligned) class RoIAlign(torch.nn.Module): def __init__(self, output_size, spatial_scale=1.0, sampling_ratio=-1, pool_mode="avg", aligned=True, use_torchvision=False): super().__init__() if isinstance(output_size, int): output_size = (output_size, output_size) self.output_size = output_size self.spatial_scale = spatial_scale self.sampling_ratio = sampling_ratio self.aligned = aligned def forward(self, input, rois): return roi_align(input, rois, self.output_size, self.spatial_scale, self.sampling_ratio, aligned=self.aligned) class RoIPool(torch.nn.Module): def __init__(self, output_size, spatial_scale=1.0): super().__init__() if isinstance(output_size, int): output_size = (output_size, output_size) self.output_size = output_size self.spatial_scale = spatial_scale def forward(self, input, rois): return tv_ops.roi_pool(input, rois, self.output_size, self.spatial_scale) class DeformConv2d(torch.nn.Module): def __init__(self, *args, **kwargs): super().__init__() raise NotImplementedError("DeformConv2d not available in CPU shim") class ModulatedDeformConv2d(torch.nn.Module): def __init__(self, *args, **kwargs): super().__init__() raise NotImplementedError("ModulatedDeformConv2d not available in CPU shim") DeformConv2dPack = DeformConv2d ModulatedDeformConv2dPack = ModulatedDeformConv2d class MaskedConv2d(torch.nn.Conv2d): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) class CornerPool(torch.nn.Module): def __init__(self, *args, **kwargs): super().__init__() raise NotImplementedError("CornerPool not available in CPU shim") def point_sample(*args, **kwargs): raise NotImplementedError("point_sample not available in CPU shim") def rel_roi_point_to_rel_img_point(*args, **kwargs): raise NotImplementedError("rel_roi_point_to_rel_img_point not available in CPU shim") def nms_match(*args, **kwargs): raise NotImplementedError("nms_match not available in CPU shim") def sigmoid_focal_loss(input, target, gamma=2.0, alpha=0.25, weight=None, reduction="mean"): p = torch.sigmoid(input) ce_loss = torch.nn.functional.binary_cross_entropy_with_logits(input, target, reduction="none") p_t = p * target + (1 - p) * (1 - target) loss = ce_loss * ((1 - p_t) ** gamma) if alpha >= 0: alpha_t = alpha * target + (1 - alpha) * (1 - target) loss = alpha_t * loss if weight is not None: loss = loss * weight if reduction == "mean": return loss.mean() elif reduction == "sum": return loss.sum() return loss def deform_conv2d(*args, **kwargs): raise NotImplementedError("deform_conv2d not available in CPU shim") def get_onnxruntime_op_path(): return "" # Import multi_scale_deform_attn to trigger ATTENTION registry registration try: from mmcv.ops.multi_scale_deform_attn import MultiScaleDeformableAttention as _MSDA except Exception: pass ''' with open(os.path.join(ops_dir, '__init__.py'), 'w') as f: f.write(init_code) # Create nms submodule nms_dir = os.path.join(ops_dir, 'nms') os.makedirs(nms_dir, exist_ok=True) with open(os.path.join(nms_dir, '__init__.py'), 'w') as f: f.write('from mmcv.ops import NMSop, batched_nms, nms\n') # Create carafe stub carafe_path = os.path.join(ops_dir, 'carafe.py') with open(carafe_path, 'w') as f: f.write('class CARAFEPack: pass\n') # Create roi_align submodule roi_align_dir = os.path.join(ops_dir, 'roi_align') os.makedirs(roi_align_dir, exist_ok=True) with open(os.path.join(roi_align_dir, '__init__.py'), 'w') as f: f.write('from mmcv.ops import roi_align, RoIAlign\n') # Create modulated_deform_conv stub with open(os.path.join(ops_dir, 'modulated_deform_conv.py'), 'w') as f: f.write('from mmcv.ops import ModulatedDeformConv2d, ModulatedDeformConv2dPack\n') # Create merge_cells stub with open(os.path.join(ops_dir, 'merge_cells.py'), 'w') as f: f.write('''class GlobalPoolingCell: def __init__(self, *args, **kwargs): raise NotImplementedError class SumCell: def __init__(self, *args, **kwargs): raise NotImplementedError class ConcatCell: def __init__(self, *args, **kwargs): raise NotImplementedError ''') # Create multi_scale_deform_attn with full CPU implementation with open(os.path.join(ops_dir, 'multi_scale_deform_attn.py'), 'w') as f: f.write('''import math import warnings import torch import torch.nn as nn import torch.nn.functional as F from mmcv.cnn import constant_init, xavier_init from mmcv.cnn.bricks.registry import ATTENTION from mmcv.runner import BaseModule def multi_scale_deformable_attn_pytorch(value, value_spatial_shapes, sampling_locations, attention_weights): """CPU version of multi-scale deformable attention.""" bs, _, num_heads, embed_dims = value.shape _, num_queries, num_heads, num_levels, num_points, _ = sampling_locations.shape value_list = value.split([int(H_ * W_) for H_, W_ in value_spatial_shapes], dim=1) sampling_grids = 2 * sampling_locations - 1 sampling_value_list = [] for level, (H_, W_) in enumerate(value_spatial_shapes): value_l_ = value_list[level].flatten(2).transpose(1, 2).reshape( bs * num_heads, embed_dims, int(H_), int(W_)) sampling_grid_l_ = sampling_grids[:, :, :, level].transpose(1, 2).flatten(0, 1) sampling_value_l_ = F.grid_sample( value_l_, sampling_grid_l_, mode=\'bilinear\', padding_mode=\'zeros\', align_corners=False) sampling_value_list.append(sampling_value_l_) attention_weights = attention_weights.transpose(1, 2).reshape( bs * num_heads, 1, num_queries, num_levels * num_points) output = (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights).sum(-1).view(bs, num_heads * embed_dims, num_queries) return output.transpose(1, 2).contiguous() @ATTENTION.register_module() class MultiScaleDeformableAttention(BaseModule): def __init__(self, embed_dims=256, num_heads=8, num_levels=4, num_points=4, im2col_step=64, dropout=0.1, batch_first=False, norm_cfg=None, init_cfg=None, **kwargs): super().__init__(init_cfg) if embed_dims % num_heads != 0: raise ValueError(f\'embed_dims must be divisible by num_heads, \' f\'but got {embed_dims} and {num_heads}\') self.norm_cfg = norm_cfg self.dropout = nn.Dropout(dropout) self.batch_first = batch_first self.im2col_step = im2col_step self.embed_dims = embed_dims self.num_levels = num_levels self.num_heads = num_heads self.num_points = num_points self.sampling_offsets = nn.Linear(embed_dims, num_heads * num_levels * num_points * 2) self.attention_weights = nn.Linear(embed_dims, num_heads * num_levels * num_points) self.value_proj = nn.Linear(embed_dims, embed_dims) self.output_proj = nn.Linear(embed_dims, embed_dims) self.init_weights() def init_weights(self): constant_init(self.sampling_offsets, 0.) device = next(self.parameters()).device thetas = torch.arange(self.num_heads, dtype=torch.float32, device=device) * (2.0 * math.pi / self.num_heads) grid_init = torch.stack([thetas.cos(), thetas.sin()], -1) grid_init = (grid_init / grid_init.abs().max(-1, keepdim=True)[0]).view( self.num_heads, 1, 1, 2).repeat(1, self.num_levels, self.num_points, 1) for i in range(self.num_points): grid_init[:, :, i, :] *= i + 1 self.sampling_offsets.bias.data = grid_init.view(-1) constant_init(self.attention_weights, val=0., bias=0.) xavier_init(self.value_proj, distribution=\'uniform\', bias=0.) xavier_init(self.output_proj, distribution=\'uniform\', bias=0.) self._is_init = True def forward(self, query, key=None, value=None, identity=None, query_pos=None, key_padding_mask=None, reference_points=None, spatial_shapes=None, level_start_index=None, **kwargs): if value is None: value = query if identity is None: identity = query if query_pos is not None: query = query + query_pos if not self.batch_first: query = query.permute(1, 0, 2) value = value.permute(1, 0, 2) bs, num_query, _ = query.shape bs, num_value, _ = value.shape value = self.value_proj(value) if key_padding_mask is not None: value = value.masked_fill(key_padding_mask[..., None], 0.0) value = value.view(bs, num_value, self.num_heads, -1) sampling_offsets = self.sampling_offsets(query).view( bs, num_query, self.num_heads, self.num_levels, self.num_points, 2) attention_weights = self.attention_weights(query).view( bs, num_query, self.num_heads, self.num_levels * self.num_points) attention_weights = attention_weights.softmax(-1) attention_weights = attention_weights.view( bs, num_query, self.num_heads, self.num_levels, self.num_points) if reference_points.shape[-1] == 2: offset_normalizer = torch.stack( [spatial_shapes[..., 1], spatial_shapes[..., 0]], -1) sampling_locations = reference_points[:, :, None, :, None, :] \\ + sampling_offsets / offset_normalizer[None, None, None, :, None, :] elif reference_points.shape[-1] == 4: sampling_locations = reference_points[:, :, None, :, None, :2] \\ + sampling_offsets / self.num_points \\ * reference_points[:, :, None, :, None, 2:] * 0.5 else: raise ValueError(f\'Last dim of reference_points must be 2 or 4, \' f\'but get {reference_points.shape[-1]} instead.\') output = multi_scale_deformable_attn_pytorch( value, spatial_shapes, sampling_locations, attention_weights) output = self.output_proj(output) if not self.batch_first: output = output.permute(1, 0, 2) return self.dropout(output) + identity ''') print('mmcv.ops shim created successfully')