Pikaliov/DCN_custom_op
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initial commit: DCNv4 custom op mirror setup

Browse Source 
- Add enhanced README with project structure and quick start guide
- Initialize repository with DCNv4 CUDA extension (PyTorch module)
- Include classification, detection, and segmentation subdirectories
- Reference upstream OpenGVLab DCNv4 implementation

Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
This commit is contained in:
Pikaliov
2026-06-11 10:30:44 +03:00
commit 1b3206b6a7
290 changed files with 41632 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
detection/mmdet_custom/models/utils/__init__.py Normal file
View File

@@ -0,0 +1,6 @@
from .query_denoising import build_dn_generator
from .transformer import (DinoTransformer, DinoTransformerDecoder)
from .convModule_norm import ConvModule_Norm
__all__ = ['build_dn_generator', 'DinoTransformer', 'DinoTransformerDecoder']

34
detection/mmdet_custom/models/utils/convModule_norm.py Normal file
View File

@@ -0,0 +1,34 @@
# --------------------------------------------------------
# DCNv4
# Copyright (c) 2024 OpenGVLab
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------
from mmcv.cnn.bricks.conv_module import ConvModule
class ConvModule_Norm(ConvModule):
def __init__(self, in_channels,
out_channels,
kernel, **kwargs):
super().__init__(in_channels, out_channels, kernel, **kwargs)
self.normType = kwargs.get('norm_cfg', {'type':''})
if self.normType is not None:
self.normType = self.normType['type']
def forward(self, x, activate=True, norm=True):
for layer in self.order:
if layer == 'conv':
if self.with_explicit_padding:
x = self.padding_layer(x)
x = self.conv(x)
elif layer == 'norm' and norm and self.with_norm:
if 'LN' in self.normType:
x = x.permute(0, 2, 3, 1)
x = self.norm(x)
x = x.permute(0, 3, 1, 2).contiguous()
else:
x = self.norm(x)
elif layer == 'act' and activate and self.with_activation:
x = self.activate(x)
return x

234
detection/mmdet_custom/models/utils/query_denoising.py Normal file
View File

@@ -0,0 +1,234 @@
# Copyright (c) OpenMMLab. All rights reserved.
import torch
from mmcv.runner import BaseModule
from mmdet.core import bbox_xyxy_to_cxcywh
from mmdet.models.utils.transformer import inverse_sigmoid
class DnQueryGenerator(BaseModule):
def __init__(self,
num_queries,
hidden_dim,
num_classes,
noise_scale=dict(label=0.5, box=0.4),
group_cfg=dict(
dynamic=True, num_groups=None, num_dn_queries=None)):
super(DnQueryGenerator, self).__init__()
self.num_queries = num_queries
self.hidden_dim = hidden_dim
self.num_classes = num_classes
self.label_noise_scale = noise_scale['label']
self.box_noise_scale = noise_scale['box']
self.dynamic_dn_groups = group_cfg.get('dynamic', False)
if self.dynamic_dn_groups:
assert 'num_dn_queries' in group_cfg, \
'num_dn_queries should be set when using ' \
'dynamic dn groups'
self.num_dn = group_cfg['num_dn_queries']
else:
assert 'num_groups' in group_cfg, \
'num_groups should be set when using ' \
'static dn groups'
self.num_dn = group_cfg['num_groups']
assert isinstance(self.num_dn, int) and self.num_dn >= 1, \
f'Expected the num in group_cfg to have type int. ' \
f'Found {type(self.num_dn)} '
def get_num_groups(self, group_queries=None):
"""
Args:
group_queries (int): Number of dn queries in one group.
"""
if self.dynamic_dn_groups:
assert group_queries is not None, \
'group_queries should be provided when using ' \
'dynamic dn groups'
if group_queries == 0:
num_groups = 1
else:
num_groups = self.num_dn // group_queries
else:
num_groups = self.num_dn
if num_groups < 1: # avoid num_groups < 1 in query generator
num_groups = 1
return int(num_groups)
def forward(self,
gt_bboxes,
gt_labels=None,
label_enc=None,
img_metas=None):
"""
Args:
gt_bboxes (List[Tensor]): List of ground truth bboxes
of the image, shape of each (num_gts, 4).
gt_labels (List[Tensor]): List of ground truth labels
of the image, shape of each (num_gts,), if None,
TODO:noisy_label would be None.
Returns:
TODO
"""
# TODO: temp only support for CDN
# TODO: temp assert gt_labels is not None and label_enc is not None
if self.training:
if gt_labels is not None:
assert len(gt_bboxes) == len(gt_labels), \
f'the length of provided gt_labels ' \
f'{len(gt_labels)} should be equal to' \
f' that of gt_bboxes {len(gt_bboxes)}'
assert gt_labels is not None \
and label_enc is not None \
and img_metas is not None # TODO: adjust args
batch_size = len(gt_bboxes)
# convert bbox
gt_bboxes_list = []
for img_meta, bboxes in zip(img_metas, gt_bboxes):
img_h, img_w, _ = img_meta['img_shape']
factor = bboxes.new_tensor([img_w, img_h, img_w,
img_h]).unsqueeze(0)
bboxes_normalized = bbox_xyxy_to_cxcywh(bboxes) / factor
gt_bboxes_list.append(bboxes_normalized)
gt_bboxes = gt_bboxes_list
known = [torch.ones_like(labels) for labels in gt_labels]
known_num = [sum(k) for k in known]
num_groups = self.get_num_groups(int(max(known_num)))
unmask_bbox = unmask_label = torch.cat(known)
labels = torch.cat(gt_labels)
boxes = torch.cat(gt_bboxes)
batch_idx = torch.cat([
torch.full_like(t.long(), i) for i, t in enumerate(gt_labels)
])
known_indice = torch.nonzero(unmask_label + unmask_bbox)
known_indice = known_indice.view(-1)
known_indice = known_indice.repeat(2 * num_groups, 1).view(-1)
known_labels = labels.repeat(2 * num_groups, 1).view(-1)
known_bid = batch_idx.repeat(2 * num_groups, 1).view(-1)
known_bboxs = boxes.repeat(2 * num_groups, 1)
known_labels_expand = known_labels.clone()
known_bbox_expand = known_bboxs.clone()
if self.label_noise_scale > 0:
p = torch.rand_like(known_labels_expand.float())
chosen_indice = torch.nonzero(
p < (self.label_noise_scale * 0.5)).view(-1)
new_label = torch.randint_like(chosen_indice, 0,
self.num_classes)
known_labels_expand.scatter_(0, chosen_indice, new_label)
single_pad = int(max(known_num)) # TODO
pad_size = int(single_pad * 2 * num_groups)
positive_idx = torch.tensor(range(
len(boxes))).long().cuda().unsqueeze(0).repeat(num_groups, 1)
positive_idx += (torch.tensor(range(num_groups)) * len(boxes) *
2).long().cuda().unsqueeze(1)
positive_idx = positive_idx.flatten()
negative_idx = positive_idx + len(boxes)
if self.box_noise_scale > 0:
known_bbox_ = torch.zeros_like(known_bboxs)
known_bbox_[:, : 2] = \
known_bboxs[:, : 2] - known_bboxs[:, 2:] / 2
known_bbox_[:, 2:] = \
known_bboxs[:, :2] + known_bboxs[:, 2:] / 2
diff = torch.zeros_like(known_bboxs)
diff[:, :2] = known_bboxs[:, 2:] / 2
diff[:, 2:] = known_bboxs[:, 2:] / 2
rand_sign = torch.randint_like(
known_bboxs, low=0, high=2, dtype=torch.float32)
rand_sign = rand_sign * 2.0 - 1.0
rand_part = torch.rand_like(known_bboxs)
rand_part[negative_idx] += 1.0
rand_part *= rand_sign
known_bbox_ += \
torch.mul(rand_part, diff).cuda() * self.box_noise_scale
known_bbox_ = known_bbox_.clamp(min=0.0, max=1.0)
known_bbox_expand[:, :2] = \
(known_bbox_[:, :2] + known_bbox_[:, 2:]) / 2
known_bbox_expand[:, 2:] = \
known_bbox_[:, 2:] - known_bbox_[:, :2]
m = known_labels_expand.long().to('cuda')
input_label_embed = label_enc(m)
input_bbox_embed = inverse_sigmoid(known_bbox_expand, eps=1e-3)
padding_label = torch.zeros(pad_size, self.hidden_dim).cuda()
padding_bbox = torch.zeros(pad_size, 4).cuda()
input_query_label = padding_label.repeat(batch_size, 1, 1)
input_query_bbox = padding_bbox.repeat(batch_size, 1, 1)
map_known_indice = torch.tensor([]).to('cuda')
if len(known_num):
map_known_indice = torch.cat(
[torch.tensor(range(num)) for num in known_num])
map_known_indice = torch.cat([
map_known_indice + single_pad * i
for i in range(2 * num_groups)
]).long()
if len(known_bid):
input_query_label[(known_bid.long(),
map_known_indice)] = input_label_embed
input_query_bbox[(known_bid.long(),
map_known_indice)] = input_bbox_embed
tgt_size = pad_size + self.num_queries
attn_mask = torch.ones(tgt_size, tgt_size).to('cuda') < 0
# match query cannot see the reconstruct
attn_mask[pad_size:, :pad_size] = True
# reconstruct cannot see each other
for i in range(num_groups):
if i == 0:
attn_mask[single_pad * 2 * i:single_pad * 2 * (i + 1),
single_pad * 2 * (i + 1):pad_size] = True
if i == num_groups - 1:
attn_mask[single_pad * 2 * i:single_pad * 2 *
(i + 1), :single_pad * i * 2] = True
else:
attn_mask[single_pad * 2 * i:single_pad * 2 * (i + 1),
single_pad * 2 * (i + 1):pad_size] = True
attn_mask[single_pad * 2 * i:single_pad * 2 *
(i + 1), :single_pad * 2 * i] = True
dn_meta = {
'pad_size': pad_size,
'num_dn_group': num_groups,
}
else:
input_query_label = None
input_query_bbox = None
attn_mask = None
dn_meta = None
return input_query_label, input_query_bbox, attn_mask, dn_meta
class CdnQueryGenerator(DnQueryGenerator):
def __init__(self, *args, **kwargs):
super(CdnQueryGenerator, self).__init__(*args, **kwargs)
def build_dn_generator(dn_args):
"""
Args:
dn_args (dict):
Returns:
"""
if dn_args is None:
return None
type = dn_args.pop('type')
if type == 'DnQueryGenerator':
return DnQueryGenerator(**dn_args)
elif type == 'CdnQueryGenerator':
return CdnQueryGenerator(**dn_args)
else:
raise NotImplementedError(f'{type} is not supported yet')

278
detection/mmdet_custom/models/utils/transformer.py Normal file
View File

@@ -0,0 +1,278 @@
import math
import torch
import torch.nn as nn
from mmdet.models.utils.builder import TRANSFORMER
from mmcv.cnn.bricks.registry import (
TRANSFORMER_LAYER_SEQUENCE, FEEDFORWARD_NETWORK, DROPOUT_LAYERS)
from mmdet.models.utils.transformer import (inverse_sigmoid,
DeformableDetrTransformerDecoder,
DeformableDetrTransformer)
def build_MLP(input_dim, hidden_dim, output_dim, num_layers):
# TODO: It can be implemented by add an out_channel arg of
# mmcv.cnn.bricks.transformer.FFN
assert num_layers > 1, \
f'num_layers should be greater than 1 but got {num_layers}'
h = [hidden_dim] * (num_layers - 1)
layers = list()
for n, k in zip([input_dim] + h[:-1], h):
layers.extend((nn.Linear(n, k), nn.ReLU()))
# Note that the relu func of MLP in original DETR repo is set
# 'inplace=False', however the ReLU cfg of FFN in mmdet is set
# 'inplace=True' by default.
layers.append(nn.Linear(hidden_dim, output_dim))
return nn.Sequential(*layers)
@TRANSFORMER_LAYER_SEQUENCE.register_module()
class DinoTransformerDecoder(DeformableDetrTransformerDecoder):
def __init__(self, *args, with_rp_noise=False, **kwargs):
super(DinoTransformerDecoder, self).__init__(*args, **kwargs)
self.with_rp_noise = with_rp_noise
self._init_layers()
def _init_layers(self):
self.ref_point_head = build_MLP(
self.embed_dims * 2,
self.embed_dims,
self.embed_dims,
2)
self.norm = nn.LayerNorm(self.embed_dims)
# @staticmethod
def gen_sineembed_for_position(self, pos_tensor):
# n_query, bs, _ = pos_tensor.size()
# sineembed_tensor = torch.zeros(n_query, bs, 256)
scale = 2 * math.pi
dim_t = torch.arange(
self.embed_dims//2, dtype=torch.float32, device=pos_tensor.device)
dim_t = 10000**(2 * (dim_t // 2) / (self.embed_dims//2))
x_embed = pos_tensor[:, :, 0] * scale
y_embed = pos_tensor[:, :, 1] * scale
pos_x = x_embed[:, :, None] / dim_t
pos_y = y_embed[:, :, None] / dim_t
pos_x = torch.stack((pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()),
dim=3).flatten(2)
pos_y = torch.stack((pos_y[:, :, 0::2].sin(), pos_y[:, :, 1::2].cos()),
dim=3).flatten(2)
if pos_tensor.size(-1) == 2:
pos = torch.cat((pos_y, pos_x), dim=2)
elif pos_tensor.size(-1) == 4:
w_embed = pos_tensor[:, :, 2] * scale
pos_w = w_embed[:, :, None] / dim_t
pos_w = torch.stack(
(pos_w[:, :, 0::2].sin(), pos_w[:, :, 1::2].cos()),
dim=3).flatten(2)
h_embed = pos_tensor[:, :, 3] * scale
pos_h = h_embed[:, :, None] / dim_t
pos_h = torch.stack(
(pos_h[:, :, 0::2].sin(), pos_h[:, :, 1::2].cos()),
dim=3).flatten(2)
pos = torch.cat((pos_y, pos_x, pos_w, pos_h), dim=2)
else:
raise ValueError('Unknown pos_tensor shape(-1):{}'.format(
pos_tensor.size(-1)))
return pos
def forward(self,
query,
*args,
reference_points=None,
valid_ratios=None,
reg_branches=None,
**kwargs):
output = query
intermediate = []
intermediate_reference_points = [reference_points]
for lid, layer in enumerate(self.layers):
if reference_points.shape[-1] == 4:
reference_points_input = \
reference_points[:, :, None] * torch.cat(
[valid_ratios, valid_ratios], -1)[:, None]
else:
assert reference_points.shape[-1] == 2
reference_points_input = \
reference_points[:, :, None] * valid_ratios[:, None]
if self.with_rp_noise and self.training:
device = reference_points.device
b, n, d = reference_points.size()
noise = torch.rand(b, n, d).to(device) * 0.02 - 0.01
reference_points = (reference_points + noise).clamp(0, 1)
query_sine_embed = self.gen_sineembed_for_position(
reference_points_input[:, :, 0, :])
query_pos = self.ref_point_head(query_sine_embed)
query_pos = query_pos.permute(1, 0, 2)
output = layer(
output,
*args,
query_pos=query_pos,
reference_points=reference_points_input,
**kwargs)
output = output.permute(1, 0, 2)
if reg_branches is not None:
tmp = reg_branches[lid](output)
assert reference_points.shape[-1] == 4
new_reference_points = tmp + inverse_sigmoid(
reference_points, eps=1e-3)
new_reference_points = new_reference_points.sigmoid()
reference_points = new_reference_points.detach()
output = output.permute(1, 0, 2)
if self.return_intermediate:
intermediate.append(self.norm(output))
intermediate_reference_points.append(new_reference_points)
# NOTE this is for the "Look Forward Twice" module,
# in the DeformDETR, reference_points was appended.
if self.return_intermediate:
return torch.stack(intermediate), torch.stack(
intermediate_reference_points)
return output, reference_points
@TRANSFORMER.register_module()
class DinoTransformer(DeformableDetrTransformer):
def __init__(self, *args, **kwargs):
super(DinoTransformer, self).__init__(*args, **kwargs)
def init_layers(self):
"""Initialize layers of the DinoTransformer."""
self.level_embeds = nn.Parameter(
torch.Tensor(self.num_feature_levels, self.embed_dims))
self.enc_output = nn.Linear(self.embed_dims, self.embed_dims)
self.enc_output_norm = nn.LayerNorm(self.embed_dims)
self.query_embed = nn.Embedding(self.two_stage_num_proposals,
self.embed_dims)
def init_weights(self):
super().init_weights()
nn.init.normal_(self.query_embed.weight.data)
def forward(self,
mlvl_feats,
mlvl_masks,
query_embed,
mlvl_pos_embeds,
dn_label_query,
dn_bbox_query,
attn_mask,
reg_branches=None,
cls_branches=None,
**kwargs):
assert self.as_two_stage and query_embed is None, \
'as_two_stage must be True for DINO'
feat_flatten = []
mask_flatten = []
lvl_pos_embed_flatten = []
spatial_shapes = []
for lvl, (feat, mask, pos_embed) in enumerate(
zip(mlvl_feats, mlvl_masks, mlvl_pos_embeds)):
bs, c, h, w = feat.shape
spatial_shape = (h, w)
spatial_shapes.append(spatial_shape)
feat = feat.flatten(2).transpose(1, 2)
mask = mask.flatten(1)
pos_embed = pos_embed.flatten(2).transpose(1, 2)
lvl_pos_embed = pos_embed + self.level_embeds[lvl].view(1, 1, -1)
lvl_pos_embed_flatten.append(lvl_pos_embed)
feat_flatten.append(feat)
mask_flatten.append(mask)
feat_flatten = torch.cat(feat_flatten, 1)
mask_flatten = torch.cat(mask_flatten, 1)
lvl_pos_embed_flatten = torch.cat(lvl_pos_embed_flatten, 1)
spatial_shapes = torch.as_tensor(
spatial_shapes, dtype=torch.long, device=feat_flatten.device)
level_start_index = torch.cat((spatial_shapes.new_zeros(
(1, )), spatial_shapes.prod(1).cumsum(0)[:-1]))
valid_ratios = torch.stack(
[self.get_valid_ratio(m) for m in mlvl_masks], 1)
reference_points = self.get_reference_points(
spatial_shapes, valid_ratios, device=feat.device)
feat_flatten = feat_flatten.permute(1, 0, 2) # (H*W, bs, embed_dims)
lvl_pos_embed_flatten = lvl_pos_embed_flatten.permute(
1, 0, 2) # (H*W, bs, embed_dims)
memory = self.encoder(
query=feat_flatten,
key=None,
value=None,
query_pos=lvl_pos_embed_flatten,
query_key_padding_mask=mask_flatten,
spatial_shapes=spatial_shapes,
reference_points=reference_points,
level_start_index=level_start_index,
valid_ratios=valid_ratios,
**kwargs)
memory = memory.permute(1, 0, 2)
bs, _, c = memory.shape
output_memory, output_proposals = self.gen_encoder_output_proposals(
memory, mask_flatten, spatial_shapes)
enc_outputs_class = cls_branches[self.decoder.num_layers](
output_memory)
enc_outputs_coord_unact = reg_branches[self.decoder.num_layers](
output_memory) + output_proposals
cls_out_features = cls_branches[self.decoder.num_layers].out_features
topk = self.two_stage_num_proposals
# NOTE In DeformDETR, enc_outputs_class[..., 0] is used for topk TODO
topk_indices = torch.topk(enc_outputs_class.max(-1)[0], topk, dim=1)[1]
# topk_proposal = torch.gather(
# output_proposals, 1,
# topk_indices.unsqueeze(-1).repeat(1, 1, 4)).sigmoid()
# topk_memory = torch.gather(
# output_memory, 1,
# topk_indices.unsqueeze(-1).repeat(1, 1, self.embed_dims))
topk_score = torch.gather(
enc_outputs_class, 1,
topk_indices.unsqueeze(-1).repeat(1, 1, cls_out_features))
topk_coords_unact = torch.gather(
enc_outputs_coord_unact, 1,
topk_indices.unsqueeze(-1).repeat(1, 1, 4))
topk_anchor = topk_coords_unact.sigmoid()
# NOTE In the original DeformDETR, init_reference_out is obtained
# from detached topk_coords_unact, which is different with DINO. TODO
topk_coords_unact = topk_coords_unact.detach()
query = self.query_embed.weight[:, None, :].repeat(1, bs,
1).transpose(0, 1)
if dn_label_query is not None:
query = torch.cat([dn_label_query, query], dim=1)
if dn_bbox_query is not None:
reference_points = torch.cat([dn_bbox_query, topk_coords_unact],
dim=1)
else:
reference_points = topk_coords_unact
reference_points = reference_points.sigmoid()
# decoder
query = query.permute(1, 0, 2)
memory = memory.permute(1, 0, 2)
inter_states, inter_references = self.decoder(
query=query,
key=None,
value=memory,
attn_masks=attn_mask,
key_padding_mask=mask_flatten,
reference_points=reference_points,
spatial_shapes=spatial_shapes,
level_start_index=level_start_index,
valid_ratios=valid_ratios,
reg_branches=reg_branches,
**kwargs)
inter_references_out = inter_references
return inter_states, inter_references_out, topk_score, topk_anchor