setup_mmcv_shim.py

"""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')