birth

.gitignore

@@ -0,0 +1,25 @@
+
+.idea/
+.DS_Store
+__pycache__/
+classification/convertor/
+segmentation/convertor/
+checkpoint_dir/
+demo/
+detection/work_dirs
+*.pth
+ops_dcnv3/build
+ops_dcnv3/dist
+ops_dcnv3/DCNv3.egg-info
+build/
+dist/
+ckpts/
+ckpts
+data
+data/
+detection/data
+detection/ckpts
+segmentation/data
+segmentation/ckpts
+work_dirs/
+output

DCNv4_op/DCNv4/functions/__init__.py

@@ -0,0 +1,11 @@
+# ------------------------------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------------------
+# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+# ------------------------------------------------------------------------------------------------
+
+# from .ms_flash_deform_attn_func import FlashMSDeformAttnFunction
+from .flash_deform_attn_func import FlashDeformAttnFunction
+from .dcnv4_func import DCNv4Function

DCNv4_op/DCNv4/functions/dcnv4_func.py

@@ -0,0 +1,129 @@
+# --------------------------------------------------------
+# InternImage
+# Copyright (c) 2022 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import torch
+import math
+import torch.nn.functional as F
+from torch.autograd import Function
+from torch.autograd.function import once_differentiable
+from torch.cuda.amp import custom_bwd, custom_fwd
+from .table import TABLE, BWDTABLE
+
+from DCNv4 import ext
+
+def factors(N):
+    res = []
+    for i in range(1, N+1):
+        if N % i == 0:
+            res.append(i)
+    return res
+
+def findspec(B, H, W, G, C):
+    key = f"{B}x{H}x{W}x{G}x{C}"
+    if key in TABLE:
+        return TABLE[key][0], TABLE[key][1]
+
+    d_stride = 8
+    ms = factors(B*H*W)
+    multiplier = 1
+    for m in ms:
+        if m <= 64 and (m * G * C // d_stride) <= 512:
+            multiplier = m
+    n_thread = multiplier * G * C // d_stride
+    key = f"{B}x{H}x{W}x{G}x{C}"
+    TABLE[key] = (d_stride, n_thread)
+    return d_stride, n_thread
+
+def find_spec_bwd(B, H, W, G, C):
+    key = f"{B}x{H}x{W}x{G}x{C}"
+    if key in BWDTABLE:
+        return BWDTABLE[key][0], BWDTABLE[key][1]
+
+    if C >= 64:
+        d_stride = 2
+    else:
+        d_stride = 1
+
+    ms = factors(B*H*W)
+    multiplier = 1
+    for m in ms:
+        if m <= 64 and (m * G * C // d_stride) <= 256:
+            multiplier = m
+    n_thread = multiplier * G * C // d_stride
+    return d_stride, n_thread
+
+class DCNv4Function(Function):
+    @staticmethod
+    @custom_fwd
+    def forward(
+            ctx, input, offset_mask,
+            kernel_h, kernel_w, stride_h, stride_w,
+            pad_h, pad_w, dilation_h, dilation_w,
+            group, group_channels, offset_scale,
+            im2col_step, remove_center):
+
+        forward_d_stride, forward_block_thread = findspec(input.shape[0], input.shape[1], input.shape[2], group, group_channels)
+        backward_d_stride, backward_block_thread = find_spec_bwd(input.shape[0], input.shape[1], input.shape[2], group, group_channels)
+
+        ctx.kernel_h = kernel_h
+        ctx.kernel_w = kernel_w
+        ctx.stride_h = stride_h
+        ctx.stride_w = stride_w
+        ctx.pad_h = pad_h
+        ctx.pad_w = pad_w
+        ctx.dilation_h = dilation_h
+        ctx.dilation_w = dilation_w
+        ctx.group = group
+        ctx.group_channels = group_channels
+        ctx.offset_scale = offset_scale
+        ctx.im2col_step = im2col_step
+        ctx.remove_center = remove_center
+        ctx.backward_d_stride = backward_d_stride
+        ctx.backward_block_thread = backward_block_thread
+
+        args = [
+            input, offset_mask, kernel_h,
+            kernel_w, stride_h, stride_w, pad_h,
+            pad_w, dilation_h, dilation_w, group,
+            group_channels, offset_scale,
+            ctx.im2col_step,
+            remove_center,
+            forward_d_stride,
+            forward_block_thread,
+            False,
+        ]
+
+        output = ext.dcnv4_forward(*args)
+        ctx.save_for_backward(input, offset_mask)
+
+        return output
+
+    @staticmethod
+    @once_differentiable
+    @custom_bwd
+    def backward(ctx, grad_output):
+        input, offset_mask = ctx.saved_tensors
+
+        args = [
+            input, offset_mask, ctx.kernel_h,
+            ctx.kernel_w, ctx.stride_h, ctx.stride_w, ctx.pad_h,
+            ctx.pad_w, ctx.dilation_h, ctx.dilation_w, ctx.group,
+            ctx.group_channels, ctx.offset_scale, ctx.im2col_step,
+            grad_output.contiguous(), ctx.remove_center,
+            ctx.backward_d_stride, ctx.backward_block_thread,
+            False
+        ]
+
+        grad_input, grad_offset_mask = \
+            ext.dcnv4_backward(*args)
+
+        return grad_input, grad_offset_mask, \
+            None, None, None, None, None, None, None,\
+            None, None, None, None, None, None

DCNv4_op/DCNv4/functions/flash_deform_attn_func.py

@@ -0,0 +1,114 @@
+# ------------------------------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------------------
+# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+# ------------------------------------------------------------------------------------------------
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import torch
+import torch.nn.functional as F
+from torch.autograd import Function
+from torch.autograd.function import once_differentiable
+import numpy as np
+
+from DCNv4 import ext
+
+shm_size_dict = {
+    "8.0": 163000,
+    "8.6": 99000,
+    "8.7": 163000,
+    "8.9": 99000,
+    "9.0": 227000,
+    "7.5": 64000,
+    "7.0": 96000,
+}
+
+cuda_capability = f"{torch.cuda.get_device_properties(0).major}.{torch.cuda.get_device_properties(0).minor}"
+
+if cuda_capability not in shm_size_dict:
+    raise NotImplementedError
+
+shm_size_cap = shm_size_dict[cuda_capability]
+
+def factors(N):
+    res = []
+    for i in range(1, N+1):
+        if N % i == 0:
+            res.append(i)
+    return res
+
+def findspec(B, Q, G, C):
+    d_stride = 8
+    ms = factors(B*Q)
+    multiplier = 1
+    for m in ms:
+        if m <= 64 and (m * G * C // d_stride) <= 512:
+            multiplier = m
+    n_thread = multiplier * G * C // d_stride
+    return d_stride, n_thread
+
+def findspec_bwd(B, Q, G, C):
+    if C >= 64:
+        d_stride = 2
+    else:
+        d_stride = 1
+    
+    ms = factors(B*Q)
+    multiplier = 1
+    for m in ms:
+        if m <= 64 and (m * G * C // d_stride) <= 256:
+            multiplier = m
+    n_thread = multiplier * G * C // d_stride
+    return d_stride, n_thread
+
+class FlashDeformAttnFunction(Function):
+    @staticmethod
+    @torch.autocast("cuda", enabled=True, dtype=torch.float16)
+    def forward(
+        ctx, value, value_spatial_shapes, value_level_start_index, 
+        sampling_loc_attn, im2col_step, K=8
+    ):
+
+        ctx.im2col_step = im2col_step
+        ctx.K = K
+        d_stride, blockthread = findspec(value.shape[0], sampling_loc_attn.shape[1], value.shape[2], value.shape[3])
+        d_stride_backward, blockthread_backward = findspec_bwd(value.shape[0], sampling_loc_attn.shape[1], value.shape[2], value.shape[3])
+
+        ctx.d_stride_backward = d_stride_backward
+        ctx.blockthread_backward = blockthread_backward
+
+        output = ext.flash_deform_attn_forward(
+            value,
+            value_spatial_shapes,
+            value_level_start_index,
+            sampling_loc_attn,
+            ctx.im2col_step,
+            K,
+            d_stride,
+            blockthread,
+        )
+        ctx.save_for_backward(value, value_spatial_shapes, value_level_start_index, sampling_loc_attn)
+        return output
+
+    @staticmethod
+    @once_differentiable
+    def backward(ctx, grad_output):
+        value, value_spatial_shapes, value_level_start_index, sampling_loc_attn = ctx.saved_tensors
+        grad_value, grad_sampling_loc_attn = ext.flash_deform_attn_backward(
+            value,
+            value_spatial_shapes,
+            value_level_start_index,
+            sampling_loc_attn,
+            grad_output.contiguous(),
+            ctx.im2col_step,
+            ctx.K,
+            ctx.d_stride_backward,
+            ctx.blockthread_backward,
+        )
+
+        return grad_value, None, None, grad_sampling_loc_attn, None, None

DCNv4_op/DCNv4/modules/__init__.py

@@ -0,0 +1,10 @@
+# ------------------------------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------------------
+# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+# ------------------------------------------------------------------------------------------------
+
+from .flash_deform_attn import FlashDeformAttn
+from .dcnv4 import DCNv4

DCNv4_op/DCNv4/modules/dcnv4.py

@@ -0,0 +1,152 @@
+# --------------------------------------------------------
+# Deformable Convolution v4
+# Copyright (c) 2023 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import math
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn.init import xavier_uniform_, constant_
+from ..functions import DCNv4Function
+
+class CenterFeatureScaleModule(nn.Module):
+    def forward(self,
+                query,
+                center_feature_scale_proj_weight,
+                center_feature_scale_proj_bias):
+        center_feature_scale = F.linear(query,
+                                        weight=center_feature_scale_proj_weight,
+                                        bias=center_feature_scale_proj_bias).sigmoid()
+        return center_feature_scale
+
+class DCNv4(nn.Module):
+    def __init__(
+            self,
+            channels=64,
+            kernel_size=3,
+            stride=1,
+            pad=1,
+            dilation=1,
+            group=4,
+            offset_scale=1.0,
+            dw_kernel_size=None,
+            center_feature_scale=False,
+            remove_center=False,
+            output_bias=True,
+            without_pointwise=False,
+            **kwargs):
+        """
+        DCNv4 Module
+        :param channels
+        :param kernel_size
+        :param stride
+        :param pad
+        :param dilation
+        :param group
+        :param offset_scale
+        :param act_layer
+        :param norm_layer
+        """
+        super().__init__()
+        if channels % group != 0:
+            raise ValueError(
+                f'channels must be divisible by group, but got {channels} and {group}')
+        _d_per_group = channels // group
+
+        # you'd better set _d_per_group to a power of 2 which is more efficient in our CUDA implementation
+        assert _d_per_group % 16 == 0
+
+        self.offset_scale = offset_scale
+        self.channels = channels
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.dilation = dilation
+        self.pad = pad
+        self.group = group
+        self.group_channels = channels // group
+        self.offset_scale = offset_scale
+        self.dw_kernel_size = dw_kernel_size
+        self.center_feature_scale = center_feature_scale
+        self.remove_center = int(remove_center)
+        self.without_pointwise = without_pointwise
+
+        self.K =  group * (kernel_size * kernel_size - self.remove_center)
+        if dw_kernel_size is not None:
+            self.offset_mask_dw = nn.Conv2d(channels, channels, dw_kernel_size, stride=1, padding=(dw_kernel_size - 1) // 2, groups=channels)
+        self.offset_mask = nn.Linear(channels, int(math.ceil((self.K * 3)/8)*8))
+        if not without_pointwise:
+            self.value_proj = nn.Linear(channels, channels)
+            self.output_proj = nn.Linear(channels, channels, bias=output_bias)
+        self._reset_parameters()
+
+        if center_feature_scale:
+            self.center_feature_scale_proj_weight = nn.Parameter(
+                torch.zeros((group, channels), dtype=torch.float))
+            self.center_feature_scale_proj_bias = nn.Parameter(
+                torch.tensor(0.0, dtype=torch.float).view((1,)).repeat(group, ))
+            self.center_feature_scale_module = CenterFeatureScaleModule()
+
+    def _reset_parameters(self):
+        constant_(self.offset_mask.weight.data, 0.)
+        constant_(self.offset_mask.bias.data, 0.)
+        if not self.without_pointwise:
+            xavier_uniform_(self.value_proj.weight.data)
+            constant_(self.value_proj.bias.data, 0.)
+            xavier_uniform_(self.output_proj.weight.data)
+            if self.output_proj.bias is not None:
+                constant_(self.output_proj.bias.data, 0.)
+
+    def forward(self, input, shape=None):
+        """
+        :param query                       (N, H, W, C)
+        :return output                     (N, H, W, C)
+        """
+        N, L, C = input.shape
+        if shape is not None:
+            H, W = shape
+        else:
+            H, W = int(L**0.5), int(L**0.5)
+
+
+        x = input
+        if not self.without_pointwise:
+            x = self.value_proj(x)
+        x = x.reshape(N, H, W, -1)
+        if self.dw_kernel_size is not None:
+            offset_mask_input = self.offset_mask_dw(input.view(N, H, W, C).permute(0, 3, 1, 2))
+            offset_mask_input = offset_mask_input.permute(0, 2, 3, 1).view(N, L, C)
+        else:
+            offset_mask_input = input
+        offset_mask = self.offset_mask(offset_mask_input).reshape(N, H, W, -1)
+
+        x_proj = x
+
+        x = DCNv4Function.apply(
+            x, offset_mask,
+            self.kernel_size, self.kernel_size,
+            self.stride, self.stride,
+            self.pad, self.pad,
+            self.dilation, self.dilation,
+            self.group, self.group_channels,
+            self.offset_scale,
+            256,
+            self.remove_center
+            )
+
+        x = x.view(N, L, -1)
+        if self.center_feature_scale:
+            center_feature_scale = self.center_feature_scale_module(
+                x, self.center_feature_scale_proj_weight, self.center_feature_scale_proj_bias)
+            center_feature_scale = center_feature_scale[..., None].repeat(
+                1, 1, 1, 1, self.channels // self.group).flatten(-2)
+            x = x * (1 - center_feature_scale) + x_proj * center_feature_scale
+        if not self.without_pointwise:
+            x = self.output_proj(x)
+        return x
+

DCNv4_op/DCNv4/modules/flash_deform_attn.py

@@ -0,0 +1,141 @@
+# ------------------------------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------------------
+# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+# ------------------------------------------------------------------------------------------------
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import warnings
+import math
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn.init import xavier_uniform_, constant_
+
+from ..functions import FlashDeformAttnFunction
+
+
+def _is_power_of_2(n):
+    if (not isinstance(n, int)) or (n < 0):
+        raise ValueError("invalid input for _is_power_of_2: {} (type: {})".format(n, type(n)))
+    return (n & (n - 1) == 0) and n != 0
+
+
+class FlashDeformAttn(nn.Module):
+    def __init__(self, d_model=256, n_levels=4, n_heads=8, n_points=4):
+        """
+        Multi-Scale Deformable Attention Module
+        :param d_model      hidden dimension
+        :param n_levels     number of feature levels
+        :param n_heads      number of attention heads
+        :param n_points     number of sampling points per attention head per feature level
+        """
+        super().__init__()
+        if d_model % n_heads != 0:
+            raise ValueError("d_model must be divisible by n_heads, but got {} and {}".format(d_model, n_heads))
+        _d_per_head = d_model // n_heads
+        # you'd better set _d_per_head to a power of 2 which is more efficient in our CUDA implementation
+        if not _is_power_of_2(_d_per_head):
+            warnings.warn(
+                "You'd better set d_model in MSDeformAttn to make the dimension of each attention head a power of 2 "
+                "which is more efficient in our CUDA implementation."
+            )
+
+        self.im2col_step = 64
+
+        self.d_model = d_model
+        self.n_levels = n_levels
+        self.n_heads = n_heads
+        self.n_points = n_points
+
+        self.sampling_offsets = nn.Linear(d_model, n_heads * n_levels * n_points * 2)
+        self.attention_weights = nn.Linear(d_model, n_heads * n_levels * n_points)
+        self.value_proj = nn.Linear(d_model, d_model)
+        self.output_proj = nn.Linear(d_model, d_model)
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        constant_(self.sampling_offsets.weight.data, 0.0)
+        thetas = torch.arange(self.n_heads, dtype=torch.float32) * (2.0 * math.pi / self.n_heads)
+        grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)
+        grid_init = (
+            (grid_init / grid_init.abs().max(-1, keepdim=True)[0])
+            .view(self.n_heads, 1, 1, 2)
+            .repeat(1, self.n_levels, self.n_points, 1)
+        )
+        for i in range(self.n_points):
+            grid_init[:, :, i, :] *= i + 1
+        with torch.no_grad():
+            self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1))
+        constant_(self.attention_weights.weight.data, 0.0)
+        constant_(self.attention_weights.bias.data, 0.0)
+        xavier_uniform_(self.value_proj.weight.data)
+        constant_(self.value_proj.bias.data, 0.0)
+        xavier_uniform_(self.output_proj.weight.data)
+        constant_(self.output_proj.bias.data, 0.0)
+
+    def forward(
+        self,
+        query,
+        reference_points,
+        input_flatten,
+        input_spatial_shapes,
+        input_level_start_index,
+        input_padding_mask=None,
+    ):
+        """
+        :param query                       (N, Length_{query}, C)
+        :param reference_points            (N, Length_{query}, n_levels, 2), range in [0, 1], top-left (0,0), bottom-right (1, 1), including padding area
+                                        or (N, Length_{query}, n_levels, 4), add additional (w, h) to form reference boxes
+        :param input_flatten               (N, \sum_{l=0}^{L-1} H_l \cdot W_l, C)
+        :param input_spatial_shapes        (n_levels, 2), [(H_0, W_0), (H_1, W_1), ..., (H_{L-1}, W_{L-1})]
+        :param input_level_start_index     (n_levels, ), [0, H_0*W_0, H_0*W_0+H_1*W_1, H_0*W_0+H_1*W_1+H_2*W_2, ..., H_0*W_0+H_1*W_1+...+H_{L-1}*W_{L-1}]
+        :param input_padding_mask          (N, \sum_{l=0}^{L-1} H_l \cdot W_l), True for padding elements, False for non-padding elements
+
+        :return output                     (N, Length_{query}, C)
+        """
+        N, Len_q, _ = query.shape
+        N, Len_in, _ = input_flatten.shape
+        assert (input_spatial_shapes[:, 0] * input_spatial_shapes[:, 1]).sum() == Len_in
+
+        value = self.value_proj(input_flatten)
+        if input_padding_mask is not None:
+            value = value.masked_fill(input_padding_mask[..., None], float(0))
+        value = value.view(N, Len_in, self.n_heads, self.d_model // self.n_heads)
+        sampling_offsets = self.sampling_offsets(query).view(N, Len_q, self.n_heads, self.n_levels, self.n_points, 2)
+        attention_weights = self.attention_weights(query).view(N, Len_q, self.n_heads, self.n_levels * self.n_points)
+        attention_weights = F.softmax(attention_weights, -1).view(N, Len_q, self.n_heads, self.n_levels, self.n_points)
+        # N, Len_q, n_heads, n_levels, n_points, 2
+        if reference_points.shape[-1] == 2:
+            offset_normalizer = torch.stack([input_spatial_shapes[..., 1], input_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.n_points * reference_points[:, :, None, :, None, 2:] * 0.5
+            )
+        else:
+            raise ValueError(
+                "Last dim of reference_points must be 2 or 4, but get {} instead.".format(reference_points.shape[-1])
+            )
+
+        output = FlashDeformAttnFunction.apply(
+            value,
+            input_spatial_shapes,
+            input_level_start_index,
+            sampling_locations,
+            attention_weights,
+            self.im2col_step,
+        )
+        output = self.output_proj(output)
+        return output

DCNv4_op/MANIFEST.in

@@ -0,0 +1,2 @@
+include src/*
+include src/cuda/*

DCNv4_op/__init__.py

@@ -0,0 +1,2 @@
+from .functions import DCNv4Function, FlashDeformAttnFunction
+from .modules import DCNv4, FlashDeformAttn

DCNv4_op/make.sh

@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+# ------------------------------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------------------
+# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+# ------------------------------------------------------------------------------------------------
+
+python setup.py build install

DCNv4_op/scripts/find_best.py

@@ -0,0 +1,61 @@
+import json
+import argparse
+class LineParser:
+    def __init__(self) -> None:
+        self.data = {}
+
+    def parse(self, line):
+        def startswith(line, lst):
+            for ele in lst:
+                if line.startswith(ele):
+                    return True
+            return False
+
+        if not startswith(line, ['1', '2', '3', '4', '5', '6', '7', '8', '9']):
+            return
+
+        eles = line.strip().split()
+        key = eles[0]
+        if key not in self.data:
+            self.data[key] = []
+        
+        self.data[key].append([eles[1], float(eles[2])])
+    
+    def sort(self):
+        for k, v in self.data.items():
+            nv = sorted(v, key=lambda x: x[1])
+            self.data[k] = nv
+    
+    def display_best(self):
+        for k, v in self.data.items():
+            print(f'{k} \t {v[0][0]} \t {v[0][1]:.4f} \t {v[1][0]} \t {v[1][1]:.4f}')
+    
+    def display_best_python(self, output):
+        res = {}
+        def parse(spec):
+            d_stride = int(spec.split('/')[0])
+            thread = int(spec.split('/')[1].split('(')[0])
+            m = int(spec.split('(')[1].split(')')[0])
+            return d_stride, thread, m
+
+        for k, v in self.data.items():
+            res[k] = parse(v[0][0])
+
+        with open(output, "w") as f:
+            json.dump(res, f, indent=4)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--input', type=str)
+    parser.add_argument('--output', type=str)
+    args = parser.parse_args()
+
+    with open(args.input) as f:
+        lines = f.readlines()
+
+    lineparser = LineParser()
+    for line in lines:
+        lineparser.parse(line)
+    lineparser.sort()
+    lineparser.display_best()
+    lineparser.display_best_python(args.output)

DCNv4_op/scripts/search_bwd.sh

@@ -0,0 +1,2 @@
+python search_dcnv4_bwd_engine.py > res_bwd.txt
+python find_best.py --input res_bwd.txt --output table_bwd.py

DCNv4_op/scripts/search_dcnv4.py

@@ -0,0 +1,131 @@
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import time
+import math
+import torch
+import torch.nn as nn
+import math
+from torch.autograd import gradcheck
+import pandas as pd
+from easydict import EasyDict as edict
+import argparse
+
+from torch.cuda import Event
+
+from functions.dcnv3_func import DCNv3Function, dcnv3_core_pytorch
+from functions.dcnv4_func import DCNv4Function
+torch.set_printoptions(threshold=10000)
+
+torch.manual_seed(3)
+
+
+#@torch.no_grad()
+def speed_test(func, args, inputs, name='Unknown'):
+
+    tic = Event(enable_timing=True)
+    toc = Event(enable_timing=True)
+    # warmup
+    for i in range(args.warmup_num):
+        func(*inputs)
+
+    total_time = 0
+    tic.record()
+    for i in range(args.test_num):
+        o = func(*inputs)
+        torch.cuda.synchronize()
+    toc.record()
+
+    avg_time = tic.elapsed_time(toc) / args.test_num
+    # print(
+        # f'>>> {name: <10} finished {args.test_num} running, avg_time: {avg_time:.6f} ms')
+    return avg_time
+
+@torch.no_grad()
+def test(N, H_in, W_in, M, D, spec=None):
+    Kh, Kw = 3, 3
+    remove_center = False
+    P = Kh * Kw - remove_center
+    offset_scale = 2.0
+    pad = 1
+    dilation = 1
+    stride = 1
+    H_out = (H_in + 2 * pad - (dilation * (Kh - 1) + 1)) // stride + 1
+    W_out = (W_in + 2 * pad - (dilation * (Kw - 1) + 1)) // stride + 1
+
+    input = torch.rand(N, H_in, W_in, M*D).cuda()
+    # print(input.shape)
+    offset = (torch.rand(N, H_out, W_out, M*P*2).cuda() * 2 - 1)*2
+    # offset = (torch.rand(N, H_out, W_out, M*P*2).cuda() * 2 - 1)*0
+    mask_origin = torch.rand(N, H_out, W_out, M, P).cuda() + 1e-5
+    mask_origin = mask_origin.half()
+    mask = mask_origin
+    # mask = torch.nn.functional.softmax(mask_origin, dim=-1)
+    offset_mask = torch.cat([offset.unflatten(-1, (M, P * 2)), mask_origin.detach()], dim=-1).flatten(-2)
+
+    im2col_step = 128
+
+    input = input.half()
+    offset = offset.half()
+    mask = mask.half()
+    offset_mask = offset_mask.half()
+
+    dcnv3_args = [
+        input,
+        offset,
+        mask,
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center,
+    ]
+    output_pytorch = DCNv3Function.apply(*dcnv3_args)
+
+    input1 = input.detach()
+
+    def pad(om):
+        padded_zero = int(math.ceil(om.shape[3]/8)*8) - om.shape[3]
+        padded = torch.zeros(om.shape[0], om.shape[1], om.shape[2], padded_zero).to(om)
+        return torch.cat([om, padded], dim=-1)
+
+    dcnv4_args = [
+        input1, pad(offset_mask),
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center, 
+        spec[0], spec[1], 2, None
+        # 8, 512, 2, 256
+    ]
+    output_flash_cuda = DCNv4Function.apply(*dcnv4_args)
+
+    fwdok = torch.allclose(output_flash_cuda, output_pytorch, rtol=1e-2, atol=1e-3)
+    max_abs_err = (output_flash_cuda - output_pytorch).abs().max()
+    max_rel_err = ((output_flash_cuda - output_pytorch).abs() /
+                   (output_pytorch.abs()+ 1e-3)).max()
+    # print('>>> forward half')
+    # print(f'* {fwdok} check_forward_equal_with_pytorch_float: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}')
+    if not fwdok:
+        print(f"Wrong: {N}x{H_in}x{W_in}x{M}x{D} \t {spec[0]}/{spec[1]}({spec[2]})")
+        return
+    # assert(fwdok)
+
+    test_args = edict({'warmup_num': 10000, 'test_num': 10000})
+    
+    exp_time_dcnv4 = speed_test(DCNv4Function.apply, test_args, dcnv4_args, name='exp')
+    torch.cuda.synchronize()
+    print(f"{N}x{H_in}x{W_in}x{M}x{D} \t {spec[0]}/{spec[1]}({spec[2]}): {exp_time_dcnv4}")
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--n", type=int)
+    parser.add_argument("--h", type=int)
+    parser.add_argument("--w", type=int)
+    parser.add_argument("--g", type=int)
+    parser.add_argument("--c", type=int)
+    parser.add_argument("--dstride", type=int)
+    parser.add_argument("--blockthread", type=int)
+    parser.add_argument("--multiplier", type=int)
+    args = parser.parse_args()
+    test(args.n, args.h, args.w, args.g, args.c, (args.dstride, args.blockthread, args.multiplier))
+
+

DCNv4_op/scripts/search_dcnv4_bwd.py

@@ -0,0 +1,200 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import time
+import torch
+import torch.nn as nn
+import math
+from torch.autograd import gradcheck
+import pandas as pd
+from easydict import EasyDict as edict
+import argparse
+
+from torch.cuda import Event
+
+from functions import DCNv4Function, DCNv3Function
+torch.set_printoptions(threshold=10000)
+
+
+
+torch.manual_seed(3)
+
+def speed_test_backward(func, args, inputs, name='Unknown'):
+    # warmup
+    # for i in range(args.warmup_num):
+    #     o = func(*inputs)
+    #     o.sum().backward()
+
+    total_time = 0
+    len_input = len(inputs)
+    for i in range(args.warmup_num + args.test_num):
+        tic = Event(enable_timing=True)
+        toc = Event(enable_timing=True)
+        inputs[0] = inputs[0].detach()
+        inputs[0].requires_grad = True
+        if len_input > 1 and isinstance(inputs[1], torch.Tensor):
+            inputs[1] = inputs[1].detach()
+            inputs[1].requires_grad = True
+        if len_input > 2 and isinstance(inputs[2], torch.Tensor):
+            inputs[2] = inputs[2].detach()
+            inputs[2].requires_grad = True
+
+        o = func(*inputs)
+        torch.cuda.synchronize()
+        tic.record()
+        o.sum().backward()
+        toc.record()
+        torch.cuda.synchronize()
+        _time = tic.elapsed_time(toc)
+        if i >= args.warmup_num:
+            total_time += _time
+        o = o.detach()
+
+    # toc.record()
+    # torch.cuda.synchronize()
+
+    avg_time = total_time / args.test_num
+    #print(
+    #    f'>>> {name: <10} finished {args.test_num} running, avg_time: {avg_time:.6f} ms')
+    return avg_time
+
+# @torch.no_grad()
+def test(N=64, H_in=32, W_in=32, M=4, D=16, spec=None):
+    """
+    64x56x56x128(G=4)
+    2 64: 3.66
+    - offset_mask collection write 3.4022
+    - offset_mask collection 3.1968
+    
+    """
+    Kh, Kw = 3, 3
+    remove_center = False
+    P = Kh * Kw - remove_center
+    offset_scale = 2.0
+    pad = 1
+    dilation = 1
+    stride = 1
+    H_out = (H_in + 2 * pad - (dilation * (Kh - 1) + 1)) // stride + 1
+    W_out = (W_in + 2 * pad - (dilation * (Kw - 1) + 1)) // stride + 1
+
+    additions = [None, None, spec[0], spec[1], False]
+    input = torch.rand(N, H_in, W_in, M*D).cuda() * 10
+    #offset = torch.rand(N, H_out, W_out, M*P*2).cuda() * 0
+    offset = (torch.rand(N, H_out, W_out, M*P*2).cuda() * 2 - 1)*2
+    mask_origin = torch.rand(N, H_out, W_out, M, P).cuda() + 1e-5
+    mask_origin = mask_origin.half()
+    mask_origin.requires_grad = True
+    # offset_mask = torch.cat([offset.unflatten(-1, (M, P, 2)), mask_origin.detach().unsqueeze(-1)], dim=-1).flatten(-3)
+    # mask /= mask.sum(-1, keepdim=True)
+    # mask = torch.nn.functional.softmax(mask_origin, dim=-1, dtype=torch.float32)
+    mask = mask_origin
+    # mask = mask.reshape(N, H_out, W_out, M*P)
+    # offset_mask = torch.cat([offset.unflatten(-1, (M, P, 2)), mask.detach().unsqueeze(-1)], dim=-1).flatten(-3)
+    offset_mask = torch.cat([offset.detach().unflatten(-1, (M, P * 2)), mask_origin.detach()], dim=-1).flatten(-2)
+
+    im2col_step = 128
+
+    input = input.half()
+    offset = offset.half()
+    mask = mask.half()
+    input.requires_grad = True
+    offset.requires_grad = True
+    # mask.requires_grad = True
+    output_pytorch = DCNv3Function.apply(
+        input,
+        offset,
+        mask,
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center)#.detach().cpu()
+    (output_pytorch.sum()/10).backward()
+
+    def pad(om):
+        padded_zero = int(math.ceil(om.shape[3]/8)*8) - om.shape[3]
+        padded = torch.zeros(om.shape[0], om.shape[1], om.shape[2], padded_zero).to(om)
+        return torch.cat([om, padded], dim=-1)
+
+    # value_offset_mask = input.detach()
+    input1 = input.detach()
+    input1.requires_grad = True
+    offset_mask = offset_mask.half()
+    offset_mask.requires_grad = True
+    # offset_mask1.requires_grad = True
+    torch.cuda.profiler.cudart().cudaProfilerStart()
+    output_flash_cuda = DCNv4Function.apply(
+        input1, offset_mask,
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center, *additions)#.detach().cpu()
+    (output_flash_cuda.sum()/10).backward()
+    torch.cuda.profiler.cudart().cudaProfilerStop()
+
+    input_grad = input.grad
+    input2_grad = input1.grad
+    bwdok = torch.allclose(input_grad.float(), input2_grad.float(), rtol=1e-2, atol=1e-3)
+    rel_err = (input_grad.abs() - input2_grad.abs())/(input_grad.abs()+1e-3)
+
+    offset_grad1 = offset.grad
+    offset_grad2 = offset_mask.grad.reshape(N, H_out, W_out, M, P*3)[..., :P*2].reshape(N, H_out, W_out, M*P*2)
+
+    bwdok2 = torch.allclose(offset_grad1.float(), offset_grad2.float(), rtol=1e-2, atol=1e-3)
+    rel_err = (offset_grad1 - offset_grad2).abs() / (offset_grad1.abs()+1e-3)
+
+    mask_grad1 = mask_origin.grad
+    mask_grad2 = offset_mask.grad.reshape(N, H_out, W_out, M, P*3)[..., P*2:].reshape(N, H_out, W_out, M, P)
+
+    bwdok3 = torch.allclose(mask_grad1, mask_grad2, rtol=1e-2, atol=1e-3)
+    rel_err = (mask_grad1 - mask_grad2).abs() / (mask_grad1.abs()+1e-3)
+
+    fwdok = torch.allclose(output_flash_cuda, output_pytorch, rtol=1e-2, atol=1e-3)
+    max_abs_err = (output_flash_cuda - output_pytorch).abs().max()
+    max_rel_err = ((output_flash_cuda - output_pytorch).abs() /
+                   (output_pytorch.abs()+ 1e-3)).max()
+    if not (bwdok and bwdok2 and bwdok3):
+        print(f"Wrong: {N}x{H_in}x{W_in}x{M}x{D} \t {spec[0]}/{spec[1]}({spec[2]})")
+        return
+    # fn_args = [
+    #     input,
+    #     offset,
+    #     mask,
+    #     Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+    #     im2col_step, remove_center
+    # ]
+
+    flash_dcn_fn_args = [
+        input1,
+        offset_mask,
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center, *additions
+    ]
+
+    test_args = edict({'warmup_num': 1000, 'test_num': 1000})
+    try:
+        exp_time = speed_test_backward(DCNv4Function.apply, test_args, flash_dcn_fn_args, name='exp')
+    except:
+        print(f"Wrong: {N}x{H_in}x{W_in}x{M}x{D} \t {spec[0]}/{spec[1]}({spec[2]})")
+        return
+
+    torch.cuda.synchronize()
+    print(f"{N}x{H_in}x{W_in}x{M}x{D} \t {spec[0]}/{spec[1]}({spec[2]}): {exp_time}")
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--n", type=int)
+    parser.add_argument("--h", type=int)
+    parser.add_argument("--w", type=int)
+    parser.add_argument("--g", type=int)
+    parser.add_argument("--c", type=int)
+    parser.add_argument("--dstride", type=int)
+    parser.add_argument("--blockthread", type=int)
+    parser.add_argument("--multiplier", type=int)
+    args = parser.parse_args()
+    test(args.n, args.h, args.w, args.g, args.c, (args.dstride, args.blockthread, args.multiplier))
+
+

DCNv4_op/scripts/search_dcnv4_bwd_engine.py

@@ -0,0 +1,24 @@
+import os
+
+def factors(N):
+    res = []
+    for i in range(1, N+1):
+        if N % i == 0:
+            res.append(i)
+    return res
+
+if __name__ == '__main__':
+    BATCH=64
+    for N, Hin, Win in [(BATCH, 56, 56), (BATCH, 28, 28), (BATCH, 14, 14), (BATCH, 7, 7), 
+                        (1, 200, 320), (1, 100, 160), (1, 50, 80), (1, 25, 40), (1, 64, 64)]:
+        for group_channel in [16, 32, 64]:
+            for group in [4, 5, 6, 7, 8]:
+                for d_stride in [1, 2, 4]:
+                    for m in factors(N*Hin*Win):
+                        if m > 64:
+                            break
+                        block_thread = group * (group_channel//d_stride) * m
+                        if block_thread > 1024:
+                            break
+                        cmd = f"python search_dcnv4_bwd.py --n {N} --h {Hin} --w {Win} --g {group} --c {group_channel} --dstride {d_stride} --blockthread {block_thread} --multiplier {m}"
+                        os.system(cmd)

DCNv4_op/scripts/search_dcnv4_engine.py

@@ -0,0 +1,24 @@
+import os
+
+def factors(N):
+    res = []
+    for i in range(1, N+1):
+        if N % i == 0:
+            res.append(i)
+    return res
+
+if __name__ == '__main__':
+    BATCH=64
+    for group_channel in [16, 32, 64]:
+        for group in [4, 5, 6, 7, 8]:
+            for N, Hin, Win in [(BATCH, 56, 56), (BATCH, 28, 28), (BATCH, 14, 14), (BATCH, 7, 7), 
+                                (1, 200, 320), (1, 100, 160), (1, 50, 80), (1, 25, 40), (1, 64, 64)]:
+                for d_stride in [2, 4, 8, 16]:
+                    for m in factors(N*Hin*Win):
+                        if m > 64:
+                            break
+                        block_thread = group * (group_channel//d_stride) * m
+                        if block_thread > 1024:
+                            break
+                        cmd = f"python search_dcnv4.py --n {N} --h {Hin} --w {Win} --g {group} --c {group_channel} --dstride {d_stride} --blockthread {block_thread} --multiplier {m}"
+                        os.system(cmd)

DCNv4_op/scripts/search_fwd.sh

@@ -0,0 +1,2 @@
+python search_dcnv4_engine.py > res.txt
+python find_best.py --input res.txt --output table.py

DCNv4_op/scripts/test_dcnv4.py

@@ -0,0 +1,144 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import time
+import torch
+import torch.nn as nn
+import math
+from torch.autograd import gradcheck
+import pandas as pd
+from easydict import EasyDict as edict
+
+from torch.cuda import Event
+
+from functions.dcnv3_func import DCNv3Function, dcnv3_core_pytorch
+from functions.dcnv4_func import DCNv4Function
+torch.set_printoptions(threshold=10000)
+
+H_in, W_in = 56, 56
+N, M, D = 64, 4, 32
+
+# H_in, W_in = 28, 28
+# N, M, D = 64, 8, 32
+
+# H_in, W_in = 14, 14
+# N, M, D = 64, 16, 32
+
+# H_in, W_in = 7, 7
+# N, M, D = 64, 32, 32
+
+# H_in, W_in = 8, 8
+# N, M, D = 128, 4, 16
+
+
+Kh, Kw = 3, 3
+remove_center = False
+P = Kh * Kw - remove_center
+offset_scale = 2.0
+pad = 1
+dilation = 1
+stride = 1
+H_out = (H_in + 2 * pad - (dilation * (Kh - 1) + 1)) // stride + 1
+W_out = (W_in + 2 * pad - (dilation * (Kw - 1) + 1)) // stride + 1
+
+torch.manual_seed(3)
+
+#@torch.no_grad()
+def speed_test(func, args, inputs, name='Unknown'):
+
+    tic = Event(enable_timing=True)
+    toc = Event(enable_timing=True)
+    # warmup
+    for i in range(args.warmup_num):
+        func(*inputs)
+
+    total_time = 0
+    tic.record()
+    for i in range(args.test_num):
+        o = func(*inputs)
+        torch.cuda.synchronize()
+    toc.record()
+
+    avg_time = tic.elapsed_time(toc) / args.test_num
+    print(
+        f'>>> {name: <10} finished {args.test_num} running, avg_time: {avg_time:.6f} ms')
+    return avg_time
+
+@torch.no_grad()
+def check_forward_equal_with_pytorch_half():
+    input = torch.rand(N, H_in, W_in, M*D).cuda()
+    print(input.shape)
+    offset = (torch.rand(N, H_out, W_out, M*P*2).cuda() * 2 - 1)*10
+    # offset = (torch.rand(N, H_out, W_out, M*P*2).cuda() * 2 - 1)*0
+    mask_origin = torch.rand(N, H_out, W_out, M, P).cuda() + 1e-5
+    mask_origin = mask_origin.half()
+    mask = mask_origin
+    # mask = torch.nn.functional.softmax(mask_origin, dim=-1)
+    offset_mask = torch.cat([offset.unflatten(-1, (M, P * 2)), mask_origin.detach()], dim=-1).flatten(-2)
+
+    im2col_step = 128
+
+    input = input.half()
+    offset = offset.half()
+    mask = mask.half()
+    offset_mask = offset_mask.half()
+
+    dcnv3_args = [
+        input,
+        offset,
+        mask,
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center,
+    ]
+    output_pytorch = DCNv3Function.apply(*dcnv3_args)
+
+    input1 = input.detach()
+
+    def pad(om):
+        padded_zero = int(math.ceil(om.shape[3]/8)*8) - om.shape[3]
+        padded = torch.zeros(om.shape[0], om.shape[1], om.shape[2], padded_zero).to(om)
+        return torch.cat([om, padded], dim=-1)
+
+    dcnv4_args = [
+        input1, pad(offset_mask),
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center, 8, 512, 2, 256, True, True,
+    ]
+    output_flash_cuda = DCNv4Function.apply(*dcnv4_args)
+
+    fwdok = torch.allclose(output_flash_cuda, output_pytorch, rtol=1e-2, atol=1e-3)
+    max_abs_err = (output_flash_cuda - output_pytorch).abs().max()
+    max_rel_err = ((output_flash_cuda - output_pytorch).abs() /
+                   (output_pytorch.abs()+ 1e-3)).max()
+    print('>>> forward half')
+    print(f'* {fwdok} check_forward_equal_with_pytorch_float: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}')
+    assert(fwdok)
+
+    test_args = edict({'warmup_num': 1000, 'test_num': 1000})
+    
+    exp_time_dcnv4 = speed_test(DCNv4Function.apply, test_args, dcnv4_args, name='exp')
+    exp_time_dcnv3 = speed_test(DCNv3Function.apply, test_args, dcnv3_args, name='exp')
+    torch.cuda.synchronize()
+
+    results = [{}]
+    results[0]['dcnv3_time'] = exp_time_dcnv3
+    results[0]['dcnv4_time'] = exp_time_dcnv4
+    columns = list(results[0].keys())
+
+    outputs = pd.DataFrame(results, columns=columns)
+    with pd.option_context(
+        'display.max_rows', None, 'display.max_columns', None,
+        'display.max_colwidth', None, 'display.width', None,
+        'display.precision', 4, ):
+        print(outputs)
+
+
+if __name__ == '__main__':
+    check_forward_equal_with_pytorch_half()

DCNv4_op/scripts/test_dcnv4_bwd.py

@@ -0,0 +1,221 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import time
+import torch
+import torch.nn as nn
+import math
+from torch.autograd import gradcheck
+import pandas as pd
+from easydict import EasyDict as edict
+
+from torch.cuda import Event
+
+from functions import DCNv4Function, DCNv3Function
+torch.set_printoptions(threshold=10000)
+
+H_in, W_in = 56, 56
+N, M, D = 64, 4, 32
+
+# H_in, W_in = 28, 28
+# N, M, D = 64, 16, 16
+
+# H_in, W_in = 14, 14
+# N, M, D = 64, 32, 16
+
+
+# H_in, W_in = 7, 7
+# N, M, D = 64, 64, 16
+
+# H_in, W_in = 8, 8
+# N, M, D = 128, 4, 16
+
+
+Kh, Kw = 3, 3
+remove_center = False
+P = Kh * Kw - remove_center
+offset_scale = 2.0
+pad = 1
+dilation = 1
+stride = 1
+H_out = (H_in + 2 * pad - (dilation * (Kh - 1) + 1)) // stride + 1
+W_out = (W_in + 2 * pad - (dilation * (Kw - 1) + 1)) // stride + 1
+
+torch.manual_seed(3)
+
+def speed_test_backward(func, args, inputs, name='Unknown'):
+    # warmup
+    # for i in range(args.warmup_num):
+    #     o = func(*inputs)
+    #     o.sum().backward()
+
+    total_time = 0
+    len_input = len(inputs)
+    for i in range(args.warmup_num + args.test_num):
+        tic = Event(enable_timing=True)
+        toc = Event(enable_timing=True)
+        inputs[0] = inputs[0].detach()
+        inputs[0].requires_grad = True
+        if len_input > 1 and isinstance(inputs[1], torch.Tensor):
+            inputs[1] = inputs[1].detach()
+            inputs[1].requires_grad = True
+        if len_input > 2 and isinstance(inputs[2], torch.Tensor):
+            inputs[2] = inputs[2].detach()
+            inputs[2].requires_grad = True
+
+        o = func(*inputs)
+        torch.cuda.synchronize()
+        tic.record()
+        o.sum().backward()
+        toc.record()
+        torch.cuda.synchronize()
+        _time = tic.elapsed_time(toc)
+        if i >= args.warmup_num:
+            total_time += _time
+        o = o.detach()
+
+    # toc.record()
+    # torch.cuda.synchronize()
+
+    avg_time = total_time / args.test_num
+    #print(
+    #    f'>>> {name: <10} finished {args.test_num} running, avg_time: {avg_time:.6f} ms')
+    return avg_time
+
+# @torch.no_grad()
+def check_forward_equal_with_pytorch_half():
+    """
+    64x56x56x128(G=4)
+    2 64: 3.66
+    - offset_mask collection write 3.4022
+    - offset_mask collection 3.1968
+    
+    """
+    additions = [8, 128, 2, 256, False]
+    input = torch.rand(N, H_in, W_in, M*D).cuda() * 10
+    print(input.shape)
+    #offset = torch.rand(N, H_out, W_out, M*P*2).cuda() * 0
+    offset = (torch.rand(N, H_out, W_out, M*P*2).cuda() * 2 - 1)*2
+    mask_origin = torch.rand(N, H_out, W_out, M, P).cuda() + 1e-5
+    mask_origin = mask_origin.half()
+    mask_origin.requires_grad = True
+    # offset_mask = torch.cat([offset.unflatten(-1, (M, P, 2)), mask_origin.detach().unsqueeze(-1)], dim=-1).flatten(-3)
+    # mask /= mask.sum(-1, keepdim=True)
+    # mask = torch.nn.functional.softmax(mask_origin, dim=-1, dtype=torch.float32)
+    mask = mask_origin
+    # mask = mask.reshape(N, H_out, W_out, M*P)
+    # offset_mask = torch.cat([offset.unflatten(-1, (M, P, 2)), mask.detach().unsqueeze(-1)], dim=-1).flatten(-3)
+    offset_mask = torch.cat([offset.detach().unflatten(-1, (M, P * 2)), mask_origin.detach()], dim=-1).flatten(-2)
+
+    im2col_step = 128
+
+    input = input.half()
+    offset = offset.half()
+    mask = mask.half()
+    input.requires_grad = True
+    offset.requires_grad = True
+    # mask.requires_grad = True
+    output_pytorch = DCNv3Function.apply(
+        input,
+        offset,
+        mask,
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center)#.detach().cpu()
+    (output_pytorch.sum()/10).backward()
+
+    def pad(om):
+        padded_zero = int(math.ceil(om.shape[3]/8)*8) - om.shape[3]
+        padded = torch.zeros(om.shape[0], om.shape[1], om.shape[2], padded_zero).to(om)
+        return torch.cat([om, padded], dim=-1)
+
+    # value_offset_mask = input.detach()
+    input1 = input.detach()
+    input1.requires_grad = True
+    offset_mask = offset_mask.half()
+    offset_mask.requires_grad = True
+    # offset_mask1.requires_grad = True
+    torch.cuda.profiler.cudart().cudaProfilerStart()
+    output_flash_cuda = DCNv4Function.apply(
+        input1, offset_mask,
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center, *additions)#.detach().cpu()
+    (output_flash_cuda.sum()/10).backward()
+    torch.cuda.profiler.cudart().cudaProfilerStop()
+
+    input_grad = input.grad
+    input2_grad = input1.grad
+    bwdok = torch.allclose(input_grad.float(), input2_grad.float(), rtol=1e-2, atol=1e-3)
+    print("bwdok")
+    print(bwdok)
+    rel_err = (input_grad.abs() - input2_grad.abs())/(input_grad.abs()+1e-3)
+    print(rel_err.max())
+
+    offset_grad1 = offset.grad
+    offset_grad2 = offset_mask.grad.reshape(N, H_out, W_out, M, P*3)[..., :P*2].reshape(N, H_out, W_out, M*P*2)
+    # print(offset_grad1)
+    # print("====================")
+    # print(offset_grad2)
+    bwdok2 = torch.allclose(offset_grad1.float(), offset_grad2.float(), rtol=1e-2, atol=1e-3)
+    print("bwdok2")
+    print(bwdok2)
+    rel_err = (offset_grad1 - offset_grad2).abs() / (offset_grad1.abs()+1e-3)
+    print(rel_err.max())
+
+    mask_grad1 = mask_origin.grad
+    mask_grad2 = offset_mask.grad.reshape(N, H_out, W_out, M, P*3)[..., P*2:].reshape(N, H_out, W_out, M, P)
+
+    bwdok3 = torch.allclose(mask_grad1, mask_grad2, rtol=1e-2, atol=1e-3)
+    print("bwdok3")
+    print(bwdok3)
+    rel_err = (mask_grad1 - mask_grad2).abs() / (mask_grad1.abs()+1e-3)
+    print(rel_err.max())
+
+    fwdok = torch.allclose(output_flash_cuda, output_pytorch, rtol=1e-2, atol=1e-3)
+    max_abs_err = (output_flash_cuda - output_pytorch).abs().max()
+    max_rel_err = ((output_flash_cuda - output_pytorch).abs() /
+                   (output_pytorch.abs()+ 1e-3)).max()
+    print('>>> forward half')
+    print(f'* {fwdok} check_forward_equal_with_pytorch_float: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}')
+    fn_args = [
+        input,
+        offset,
+        mask,
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center
+    ]
+
+    flash_dcn_fn_args = [
+        input1,
+        offset_mask,
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center, *additions
+    ]
+
+    test_args = edict({'warmup_num': 1000, 'test_num': 1000})
+    exp_time = speed_test_backward(
+            DCNv4Function.apply, test_args, flash_dcn_fn_args, name='exp')
+    exp_time_base = speed_test_backward(
+            DCNv3Function.apply, test_args, fn_args, name='exp')
+
+    results = [{}]
+    results[0]['time'] = exp_time
+    results[0]['time_base'] = exp_time_base
+    columns = list(results[0].keys())
+
+    outputs = pd.DataFrame(results, columns=columns)
+    with pd.option_context(
+        'display.max_rows', None, 'display.max_columns', None,
+        'display.max_colwidth', None, 'display.width', None,
+        'display.precision', 4, ):
+        print(outputs)
+
+
+if __name__ == '__main__':
+    check_forward_equal_with_pytorch_half()

DCNv4_op/scripts/test_flash_deform_attn.py

@@ -0,0 +1,174 @@
+# ------------------------------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------------------
+# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+# ------------------------------------------------------------------------------------------------
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+from easydict import EasyDict as edict
+from torch.cuda import Event
+import pandas as pd
+
+import time
+import torch
+import torch.nn as nn
+from torch.autograd import gradcheck
+
+from functions import MSDeformAttnFunction, FlashDeformAttnFunction, ms_deform_attn_core_pytorch
+
+
+# N, M, D = 1, 4, 8
+# # Lq, L, P = 2, 2, 2
+# # shapes = torch.as_tensor([(6, 4), (3, 2)], dtype=torch.long).cuda()
+# Lq, L, P = 1, 2, 8
+# shapes = torch.as_tensor([(8, 16), (4, 8)], dtype=torch.long).cuda()
+
+# N, M, D = 1, 8, 32
+# # Lq, L, P = 2, 2, 2
+# # shapes = torch.as_tensor([(6, 4), (3, 2)], dtype=torch.long).cuda()
+# Lq, L, P = 300, 4, 4
+# # shapes = torch.as_tensor([(134, 151), (67, 76), (34, 38), (17, 19)], dtype=torch.long).cuda()
+# # shapes = torch.as_tensor([(134, 151), (67, 76), (34, 38), (16, 16)], dtype=torch.long).cuda()
+# # shapes = torch.as_tensor([(134, 151), (67, 76), (34, 38), (17, 19)], dtype=torch.long).cuda()
+# # shapes = torch.as_tensor([(17, 19), (4, 4)], dtype=torch.long).cuda()
+# shapes = torch.as_tensor([(100, 151), (50, 76), (25, 38), (13, 19)], dtype=torch.long).cuda()
+# # shapes = torch.as_tensor([(110, 151)], dtype=torch.long).cuda()
+
+# B:6
+# H:232
+# W:400
+# G:5
+# D: 16
+# channels: 80
+# kernel: 3 points = 3 * 3
+# num_split = 45 = kernel *kernel * G
+
+H = 256
+W = 256
+N, M, D = 1, 8, 32
+Lq, L, P = 100*152, 4, 8
+
+shapes = torch.Tensor([[100, 152], [ 50,  76], [ 25,  38], [ 13,  19]]).long().cuda()
+
+# x = x.reshape([B, H*W, G, D + self.num_split * 3])
+# shapes = torch.as_tensor([(H, W)], dtype=torch.long).cuda()
+# shapes = torch.as_tensor([(H, W), (H // 2, W // 2)], dtype=torch.long).cuda()
+# shapes = torch.as_tensor([(H, W), (H // 2, W // 2), (H // 4, W // 4), (H // 8, W // 8)], dtype=torch.long).cuda()
+
+level_start_index = torch.cat((shapes.new_zeros((1,)), shapes.prod(1).cumsum(0)[:-1]))
+S = sum([(H * W).item() for H, W in shapes])
+print(S)
+
+def get_reference_points(spatial_shapes, device):
+    reference_points_list = []
+    for lvl, (H_, W_) in enumerate(spatial_shapes):
+
+        ref_y, ref_x = torch.meshgrid(torch.linspace(0.5, H_ - 0.5, H_, dtype=torch.float32, device=device),
+                                        torch.linspace(0.5, W_ - 0.5, W_, dtype=torch.float32, device=device))
+        ref_y = ref_y.reshape(-1)[None] / (H_)
+        ref_x = ref_x.reshape(-1)[None] / (W_)
+        ref = torch.stack((ref_x, ref_y), -1)
+        reference_points_list.append(ref)
+    reference_points = torch.cat(reference_points_list, 1)
+    # reference_points = reference_points[:, :, None] * valid_ratios[:, None]
+    return reference_points
+
+
+torch.manual_seed(3)
+
+@torch.no_grad()
+def speed_test(func, args, inputs, name='Unknown'):
+
+    tic = Event(enable_timing=True)
+    toc = Event(enable_timing=True)
+    # warmup
+    for i in range(args.warmup_num):
+        func(*inputs)
+
+    tic.record()
+    for i in range(args.test_num):
+        func(*inputs)
+    toc.record()
+    torch.cuda.synchronize()
+
+    avg_time = tic.elapsed_time(toc) / args.test_num
+    print(
+        f'>>> {name: <10} finished {args.test_num} running, avg_time: {avg_time:.6f} ms')
+    return avg_time
+
+
+@torch.no_grad()
+def check_forward_equal_with_pytorch_half():
+    value = torch.rand(N, S, M, D).cuda() * 0.01
+    # offset = (torch.rand(N, Lq, M, L, P, 2).cuda() * 2 - 1) / 10
+    sampling_locations = torch.rand(N, Lq, M, L, P, 2).cuda()
+    attention_weights = torch.rand(N, Lq, M, L, P).cuda() + 1e-5
+    sampling_loc_attn = torch.cat([sampling_locations.reshape(N, Lq, M, L*P*2), attention_weights.reshape(N, Lq, M, L*P)], dim=-1)
+    attention_weights = torch.nn.functional.softmax(attention_weights.flatten(-2, -1), dim=-1).unflatten(-1, (L, P))
+
+
+    im2col_step = 128
+
+    flash_fn_args = (
+        value.half(),
+        shapes,
+        level_start_index,
+        sampling_loc_attn.half(),
+        im2col_step,
+        P, 16
+    )
+    output_cuda = (
+        FlashDeformAttnFunction.apply(*flash_fn_args)
+        .detach()
+        .cpu()
+    ).double()
+    
+    fn_args = (
+        value,
+        shapes,
+        level_start_index,
+        sampling_locations,
+        attention_weights,
+        im2col_step,
+    )
+
+    output_pytorch = (
+        MSDeformAttnFunction.apply(*fn_args)
+        .detach().double()
+        .cpu()
+    )
+
+    max_abs_err = (output_cuda - output_pytorch).abs().max()
+    max_rel_err = ((output_cuda - output_pytorch).abs() / output_pytorch.abs()).max()
+    fwdok = torch.allclose(output_cuda, output_pytorch, rtol=1e-2, atol=1e-3)
+
+    print(
+        f"* {fwdok} check_forward_equal_with_pytorch_float: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}"
+    )
+
+    test_args = edict({'warmup_num': 1000, 'test_num': 1000})
+    exp_time_base = speed_test(
+            MSDeformAttnFunction.apply, test_args, fn_args, name='exp')
+    exp_time = speed_test(
+            FlashDeformAttnFunction.apply, test_args, flash_fn_args, name='exp')
+
+    results = [{}]
+    results[0]['time'] = exp_time
+    results[0]['time_base'] = exp_time_base
+    columns = list(results[0].keys())
+
+    outputs = pd.DataFrame(results, columns=columns)
+    with pd.option_context(
+        'display.max_rows', None, 'display.max_columns', None,
+        'display.max_colwidth', None, 'display.width', None,
+        'display.precision', 4, ):
+        print(outputs)
+
+
+if __name__ == "__main__":
+    check_forward_equal_with_pytorch_half()
+

DCNv4_op/scripts/test_flash_deform_attn_backward.py

@@ -0,0 +1,194 @@
+# ------------------------------------------------------------------------------------------------
+# Deformable DETR
+# Copyright (c) 2020 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------------------
+# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+# ------------------------------------------------------------------------------------------------
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+from easydict import EasyDict as edict
+from torch.cuda import Event
+import pandas as pd
+
+import time
+import torch
+import torch.nn as nn
+from torch.autograd import gradcheck
+
+from functions import MSDeformAttnFunction, ms_deform_attn_core_pytorch, FlashDeformAttnFunction
+
+
+H = 256
+W = 256
+N, M, D = 1, 8, 16
+Lq, L, P = H * W, 1, 8
+
+# x = x.reshape([B, H*W, G, D + self.num_split * 3])
+shapes = torch.as_tensor([(H, W)], dtype=torch.long).cuda()
+# shapes = torch.as_tensor([(H, W), (H // 2, W // 2)], dtype=torch.long).cuda()
+# shapes = torch.as_tensor([(H, W), (H // 2, W // 2), (H // 4, W // 4), (H // 8, W // 8)], dtype=torch.long).cuda()
+
+H = 256
+W = 256
+N, M, D = 1, 8, 32
+Lq, L, P = 100*152, 4, 8
+
+shapes = torch.Tensor([[100, 152], [ 50,  76], [ 25,  38], [ 13,  19]]).long().cuda()
+
+level_start_index = torch.cat((shapes.new_zeros((1,)), shapes.prod(1).cumsum(0)[:-1]))
+S = sum([(H * W).item() for H, W in shapes])
+
+def get_reference_points(spatial_shapes, device):
+    reference_points_list = []
+    for lvl, (H_, W_) in enumerate(spatial_shapes):
+
+        ref_y, ref_x = torch.meshgrid(torch.linspace(0.5, H_ - 0.5, H_, dtype=torch.float32, device=device),
+                                        torch.linspace(0.5, W_ - 0.5, W_, dtype=torch.float32, device=device))
+        ref_y = ref_y.reshape(-1)[None] / (H_)
+        ref_x = ref_x.reshape(-1)[None] / (W_)
+        ref = torch.stack((ref_x, ref_y), -1)
+        reference_points_list.append(ref)
+    reference_points = torch.cat(reference_points_list, 1)
+    # reference_points = reference_points[:, :, None] * valid_ratios[:, None]
+    return reference_points
+
+
+torch.manual_seed(3)
+
+@torch.no_grad()
+def speed_test(func, args, inputs, name='Unknown'):
+
+    tic = Event(enable_timing=True)
+    toc = Event(enable_timing=True)
+    # warmup
+    for i in range(args.warmup_num):
+        func(*inputs)
+
+    tic.record()
+    for i in range(args.test_num):
+        func(*inputs)
+    toc.record()
+    torch.cuda.synchronize()
+
+    avg_time = tic.elapsed_time(toc) / args.test_num
+    print(
+        f'>>> {name: <10} finished {args.test_num} running, avg_time: {avg_time:.6f} ms')
+    return avg_time
+
+
+def check_forward_equal_with_pytorch_half():
+    value = torch.rand(N, S, M, D).cuda() * 0.01
+    offset = (torch.rand(N, Lq, M, L, P, 2).cuda() * 2 - 1) / 10
+    sampling_locations = torch.rand(N, Lq, M, L, P, 2).cuda()
+    attention_weights_origin = torch.rand(N, Lq, M, L, P).cuda() + 1e-5
+    attention_weights_origin.requires_grad = True
+    sampling_loc_attn = torch.cat([sampling_locations.detach().reshape(N, Lq, M, L*P*2), attention_weights_origin.detach().reshape(N, Lq, M, L*P)], dim=-1)
+
+    attention_weights = torch.nn.functional.softmax(attention_weights_origin.flatten(-2, -1), dim=-1).unflatten(-1, (L, P))
+
+
+    im2col_step = 128
+
+    value.requires_grad = True
+    sampling_loc_attn.requires_grad = True
+    output_cuda = (
+        FlashDeformAttnFunction.apply(
+            value.float(),
+            shapes,
+            level_start_index,
+            sampling_loc_attn.float(),
+            im2col_step,
+        )
+    )
+    (output_cuda.float().sum()/10).backward()
+
+
+    value1 = value.detach()
+    value1.requires_grad = True
+    sampling_locations.requires_grad = True
+    #attention_weights.requires_grad = True
+    output_pytorch = (
+        ms_deform_attn_core_pytorch(value1, shapes, sampling_locations, attention_weights)
+    )
+    (output_pytorch.sum()/10).backward()
+
+    max_abs_err = (output_cuda.float() - output_pytorch).abs().max()
+    max_rel_err = ((output_cuda.float() - output_pytorch).abs() / output_pytorch.abs()).max()
+    fwdok = torch.allclose(output_cuda.float(), output_pytorch, rtol=1e-2, atol=1e-3)
+    print(fwdok)
+    print(max_abs_err, max_rel_err)
+    #exit()
+
+    bwdok1 = torch.allclose(value.grad, value1.grad, rtol=1e-2, atol=1e-3)
+    print(bwdok1)
+    # rel_err = (sampling_locations.grad - sampling_loc_attn.grad[..., :L*P*2].reshape(*sampling_locations.shape)).abs()/(sampling_locations.grad.abs()+1e-3)
+    # print(rel_err.max())
+
+    locgrad1 = sampling_locations.grad
+    locgrad2 = sampling_loc_attn.grad[..., :L*P*2].reshape(*sampling_locations.shape)
+    bwdok2 = torch.allclose(locgrad1, locgrad2, rtol=1e-2, atol=1e-3)
+    print(bwdok2)
+    rel_err = (locgrad1 - locgrad2).abs()/(locgrad1.abs()+1e-3)
+    print(rel_err.max())
+
+    attngrad1 = attention_weights_origin.grad
+    attngrad2 = sampling_loc_attn.grad[..., L*P*2:].reshape(*attention_weights_origin.shape)
+    bwdok3 = torch.allclose(locgrad1, locgrad2, rtol=1e-2, atol=1e-3)
+    print(bwdok3)
+    rel_err = (attngrad1 - attngrad2).abs()/(attngrad1.abs()+1e-3)
+    print(rel_err.max())
+    exit()
+    #exit()
+
+    # pdb.set_trace()
+    max_abs_err = (output_cuda - output_pytorch).abs().max()
+    max_rel_err = ((output_cuda - output_pytorch).abs() / output_pytorch.abs()).max()
+    fwdok = torch.allclose(output_cuda, output_pytorch, rtol=1e-2, atol=1e-3)
+
+    print(
+        f"* {fwdok} check_forward_equal_with_pytorch_float: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}"
+    )
+
+
+    fn_args = (
+        value,
+        shapes,
+        level_start_index,
+        sampling_locations,
+        attention_weights,
+        im2col_step,
+    )
+
+    flash_dcn_fn_args = (
+        value.half(),
+        shapes,
+        level_start_index,
+        sampling_loc_attn.half(),
+        im2col_step,
+    )
+
+
+    test_args = edict({'warmup_num': 50, 'test_num': 100})
+    exp_time = speed_test(
+            FlashMSDeformAttnFunction.apply, test_args, flash_dcn_fn_args, name='exp')
+    exp_time_base = speed_test(
+            MSDeformAttnFunction.apply, test_args, fn_args, name='exp')
+
+    results = [{}]
+    results[0]['time'] = exp_time
+    results[0]['time_base'] = exp_time_base
+    columns = list(results[0].keys())
+
+    outputs = pd.DataFrame(results, columns=columns)
+    with pd.option_context(
+        'display.max_rows', None, 'display.max_columns', None,
+        'display.max_colwidth', None, 'display.width', None,
+        'display.precision', 4, ):
+        print(outputs)
+
+
+if __name__ == "__main__":
+    check_forward_equal_with_pytorch_half()

DCNv4_op/setup.py

@@ -0,0 +1,72 @@
+# ------------------------------------------------------------------------------------------------
+# Deformable Convolution v4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------------------------------
+# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+# ------------------------------------------------------------------------------------------------
+
+import os
+import glob
+
+import torch
+
+from torch.utils.cpp_extension import CUDA_HOME
+from torch.utils.cpp_extension import CppExtension
+from torch.utils.cpp_extension import CUDAExtension
+
+from setuptools import find_packages
+from setuptools import setup
+
+requirements = ["torch", "torchvision"]
+
+def get_extensions():
+    this_dir = os.path.dirname(os.path.abspath(__file__))
+    extensions_dir = os.path.join(this_dir, "src")
+
+    main_file = glob.glob(os.path.join(extensions_dir, "*.cpp"))
+    source_cpu = glob.glob(os.path.join(extensions_dir, "cpu", "*.cpp"))
+    source_cuda = glob.glob(os.path.join(extensions_dir, "cuda", "*.cu"))
+
+    sources = main_file + source_cpu
+    extension = CppExtension
+    extra_compile_args = {"cxx": []}
+    define_macros = []
+
+    if torch.cuda.is_available() and CUDA_HOME is not None:
+        extension = CUDAExtension
+        sources += source_cuda
+        define_macros += [("WITH_CUDA", None)]
+        extra_compile_args["nvcc"] = [
+            "-DCUDA_HAS_FP16=1",
+            "-D__CUDA_NO_HALF_OPERATORS__",
+            "-D__CUDA_NO_HALF_CONVERSIONS__",
+            "-D__CUDA_NO_HALF2_OPERATORS__",
+            "-O3",
+        ]
+    else:
+        raise NotImplementedError('Cuda is not available')
+
+    sources = [os.path.join(extensions_dir, s) for s in sources]
+    include_dirs = [extensions_dir]
+    ext_modules = [
+        extension(
+            "DCNv4.ext",
+            sources,
+            include_dirs=include_dirs,
+            define_macros=define_macros,
+            extra_compile_args=extra_compile_args,
+        )
+    ]
+    return ext_modules
+
+setup(
+    name="DCNv4",
+    version="1.0.0",
+    author="Yuwen Xiong, Feng Wang",
+    url="",
+    description="PyTorch Wrapper for CUDA Functions of DCNv4",
+    packages=['DCNv4', 'DCNv4/functions', 'DCNv4/modules'],
+    ext_modules=get_extensions(),
+    cmdclass={"build_ext": torch.utils.cpp_extension.BuildExtension},
+)

DCNv4_op/src/cuda/common.h

@@ -0,0 +1,212 @@
+#ifndef FMSDACOMMON
+#define FMSDACOMMON
+#include <algorithm>
+#include <cstdio>
+#include <cstring>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <THC/THCAtomics.cuh>
+
+#include <ATen/ATen.h>
+#include <ATen/OpMathType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cooperative_groups.h>
+#include <cooperative_groups/memcpy_async.h>
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+constexpr int kWarpSize = 32;
+#define opmath_t at::opmath_type<scalar_t>
+
+inline int GET_BLOCKS(const int N, const int num_threads) {
+  return (N + num_threads - 1) / num_threads;
+}
+
+#define CUDA_KERNEL_LOOP(i, n)                                                 \
+  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < (n);                 \
+       i += blockDim.x * gridDim.x)
+
+inline bool check_backward_warpp(int d_stride, int D){
+  int n_group_threads = D / d_stride;
+  return (n_group_threads <= kWarpSize) && (kWarpSize % n_group_threads == 0);
+}
+
+template <typename scalar_t, typename transfer_t, int c_per_thread>
+__device__ void ms_deform_attn_im2col_bilinear(
+    opmath_t out_reg_array[], const scalar_t *&p_value, const int &height,
+    const int &width, const opmath_t &h_px, const opmath_t &w_px,
+    const opmath_t &attn, const int &w_stride, const int &base_ptr) {
+
+  const int h_low = floor(h_px);
+  const int w_low = floor(w_px);
+  const int h_high = h_low + 1;
+  const int w_high = w_low + 1;
+  const opmath_t lh = h_px - h_low;
+  const opmath_t lw = w_px - w_low;
+  const opmath_t hh = 1 - lh;
+  const opmath_t hw = 1 - lw;
+
+  const opmath_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
+
+  const int h_stride = width * w_stride;
+  const int h_low_ptr_offset = h_low * h_stride;
+  const int h_high_ptr_offset = h_low_ptr_offset + h_stride;
+  const int w_low_ptr_offset = w_low * w_stride;
+  const int w_high_ptr_offset = w_low_ptr_offset + w_stride;
+
+  int idx1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr;
+  int idx2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr;
+  int idx3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr;
+  int idx4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr;
+
+  scalar_t v1_array[c_per_thread] = {0.};
+  scalar_t v2_array[c_per_thread] = {0.};
+  scalar_t v3_array[c_per_thread] = {0.};
+  scalar_t v4_array[c_per_thread] = {0.};
+
+  if (h_low >= 0 && w_low >= 0) {
+    auto p1 = p_value + idx1;
+    *(transfer_t *)(v1_array) = *(transfer_t *)(p1);
+  }
+
+  if (h_low >= 0 && w_high < width) {
+    auto p2 = p_value + idx2;
+    *(transfer_t *)(v2_array) = *(transfer_t *)(p2);
+  }
+  if (h_high < height && w_low >= 0) {
+    auto p3 = p_value + idx3;
+    *(transfer_t *)(v3_array) = *(transfer_t *)(p3);
+  }
+  if (h_high < height && w_high < width) {
+    auto p4 = p_value + idx4;
+    *(transfer_t *)(v4_array) = *(transfer_t *)(p4);
+  }
+#pragma unroll
+  for (int i = 0; i < c_per_thread; i++) {
+    out_reg_array[i] +=
+        (opmath_t)attn *
+        (w1 * (opmath_t)v1_array[i] + w2 * (opmath_t)v2_array[i] +
+         w3 * (opmath_t)v3_array[i] + w4 * (opmath_t)v4_array[i]);
+  }
+}
+
+template <typename scalar_t, typename transfer_t, int c_per_thread>
+__device__ void ms_deform_attn_col2im_bilinear(
+    const scalar_t *&p_value, const int &height, const int &width,
+    const opmath_t &h_px, const opmath_t &w_px, const opmath_t &attn,
+    const int &w_stride, const int &base_ptr, const opmath_t offset_scale_h,
+    const opmath_t offset_scale_w, const scalar_t *&top_grad,
+    opmath_t *&grad_im, opmath_t *grad_offset) {
+
+  const int h_low = floor(h_px);
+  const int w_low = floor(w_px);
+  const int h_high = h_low + 1;
+  const int w_high = w_low + 1;
+  const opmath_t lh = h_px - h_low;
+  const opmath_t lw = w_px - w_low;
+  const opmath_t hh = 1 - lh;
+  const opmath_t hw = 1 - lw;
+
+  const opmath_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
+
+  scalar_t _top_grad_array[c_per_thread] = {0.};
+  *(transfer_t *)(_top_grad_array) = *(transfer_t *)(top_grad);
+
+  opmath_t top_grad_array[c_per_thread] = {0.};
+  for (int i = 0; i < c_per_thread; ++i) {
+    top_grad_array[i] = (opmath_t)(_top_grad_array[i]);
+  }
+
+  const int h_stride = width * w_stride;
+  const int h_low_ptr_offset = h_low * h_stride;
+  const int h_high_ptr_offset = h_low_ptr_offset + h_stride;
+  const int w_low_ptr_offset = w_low * w_stride;
+  const int w_high_ptr_offset = w_low_ptr_offset + w_stride;
+
+  int idx1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr;
+  int idx2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr;
+  int idx3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr;
+  int idx4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr;
+
+  scalar_t v1_array[c_per_thread] = {0.};
+  scalar_t v2_array[c_per_thread] = {0.};
+  scalar_t v3_array[c_per_thread] = {0.};
+  scalar_t v4_array[c_per_thread] = {0.};
+
+  opmath_t grad_h_weight[c_per_thread] = {0.};
+  opmath_t grad_w_weight[c_per_thread] = {0.};
+
+  if (h_low >= 0 && w_low >= 0) {
+    auto p1 = p_value + idx1;
+    *(transfer_t *)(v1_array) = *(transfer_t *)(p1);
+#pragma unroll
+    for (int i = 0; i < c_per_thread; ++i) {
+      grad_h_weight[i] -= hw * v1_array[i];
+      grad_w_weight[i] -= hh * v1_array[i];
+      atomicAdd(grad_im + idx1 + i, top_grad_array[i] * attn * w1);
+    }
+  }
+
+  if (h_low >= 0 && w_high < width) {
+    auto p2 = p_value + idx2;
+    *(transfer_t *)(v2_array) = *(transfer_t *)(p2);
+#pragma unroll
+    for (int i = 0; i < c_per_thread; ++i) {
+      grad_h_weight[i] -= lw * v2_array[i];
+      grad_w_weight[i] += hh * v2_array[i];
+      atomicAdd(grad_im + idx2 + i, top_grad_array[i] * attn * w2);
+    }
+  }
+  if (h_high < height && w_low >= 0) {
+    auto p3 = p_value + idx3;
+    *(transfer_t *)(v3_array) = *(transfer_t *)(p3);
+#pragma unroll
+    for (int i = 0; i < c_per_thread; ++i) {
+      grad_h_weight[i] += hw * v3_array[i];
+      grad_w_weight[i] -= lh * v3_array[i];
+      atomicAdd(grad_im + idx3 + i, top_grad_array[i] * attn * w3);
+    }
+  }
+  if (h_high < height && w_high < width) {
+    auto p4 = p_value + idx4;
+    *(transfer_t *)(v4_array) = *(transfer_t *)(p4);
+#pragma unroll
+    for (int i = 0; i < c_per_thread; ++i) {
+      grad_h_weight[i] += lw * v4_array[i];
+      grad_w_weight[i] += lh * v4_array[i];
+      atomicAdd(grad_im + idx4 + i, top_grad_array[i] * attn * w4);
+    }
+  }
+
+  opmath_t _grad_offset_x = 0;
+  opmath_t _grad_offset_y = 0;
+#pragma unroll
+  for (int i = 0; i < c_per_thread; ++i) {
+    _grad_offset_x +=
+        grad_w_weight[i] * top_grad_array[i]; // channel aware term
+    _grad_offset_y +=
+        grad_h_weight[i] * top_grad_array[i]; // channel aware term
+  }
+  _grad_offset_x *= (offset_scale_w * attn); // channel shared term
+  _grad_offset_y *= (offset_scale_h * attn); // channel shared term
+
+  *grad_offset = _grad_offset_x;
+  *(grad_offset + 1) = _grad_offset_y;
+
+  opmath_t current_val;
+  opmath_t _grad_offset_z = 0;
+#pragma unroll
+  for (int i = 0; i < c_per_thread; i++) {
+    current_val = (opmath_t)(w1 * v1_array[i] + w2 * v2_array[i] +
+                             w3 * v3_array[i] + w4 * v4_array[i]);
+    _grad_offset_z += current_val * top_grad_array[i];
+  }
+  *(grad_offset + 2) = _grad_offset_z;
+}
+
+
+
+#endif

DCNv4_op/src/cuda/dcnv4_col2im_cuda.cuh

@@ -0,0 +1,561 @@
+#include <algorithm>
+#include <cstdio>
+#include <cstring>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <THC/THCAtomics.cuh>
+
+#include <ATen/ATen.h>
+#include <ATen/OpMathType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cooperative_groups.h>
+#include <cooperative_groups/memcpy_async.h>
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include "common.h"
+
+template <typename scalar_t, int d_stride, typename transfer_t, int L, int K,
+          bool softmax>
+__global__ void backward_kernel_dcn(
+    const scalar_t *p_value, const scalar_t *p_offset,
+    const scalar_t *grad_output, const int G, const int D, const int Q,
+    const int kernel_h, const int kernel_w, const int stride_h,
+    const int stride_w, const int pad_h, const int pad_w, const int dilation_h,
+    const int dilation_w, const int height_in, const int width_in,
+    const int height_out, const int width_out, const opmath_t offset_scale,
+    const int remove_center, const int block_multiplier, opmath_t *grad_im,
+    opmath_t *grad_offset, const int padded_offset_dim) {
+
+  extern __shared__ char _s[];
+
+  const int &qi = (blockIdx.x * block_multiplier % Q) + threadIdx.z;
+  const int &bi = blockIdx.x * block_multiplier / Q;
+
+  const int &di_s = threadIdx.x * d_stride;
+  const int &gi = threadIdx.y;
+  constexpr int li = 0;
+
+  opmath_t *const cache_g_mask_before_softmax = (opmath_t *)(_s); // mG x K
+  opmath_t *const cache_grad_offset =
+      (opmath_t *)(cache_g_mask_before_softmax +
+                   block_multiplier * G * K); // mG x blockDim.x x 3
+  opmath_t *const p_mask_shm =
+      (opmath_t *)(cache_grad_offset + block_multiplier * G * blockDim.x * 3) +
+      (threadIdx.z * G + gi) * K;
+
+  const scalar_t *p_offset_ptr = p_offset + (bi*Q + qi)*padded_offset_dim + gi*K*3;
+
+  const int mask_length = K;
+  const int num_thread = (D / d_stride);
+  const int num_iter = mask_length / num_thread;
+  const int remainder = mask_length - num_iter * num_thread;
+
+  const scalar_t *top_grad = grad_output + ((bi * Q + qi) * G + gi) * D + di_s;
+
+  __syncthreads();
+  for (int i = 0; i < num_iter; i++) {
+    *(p_mask_shm + num_thread * i + threadIdx.x) =
+        *(scalar_t *)(p_offset_ptr + K * 2 + num_thread * i + threadIdx.x);
+  }
+  if (remainder > 0 && threadIdx.x < remainder) {
+    *(p_mask_shm + num_thread * num_iter + threadIdx.x) = *(
+        scalar_t *)(p_offset_ptr + K * 2 + num_thread * num_iter + threadIdx.x);
+  }
+
+  if (softmax) {
+    __syncthreads();
+    // transfer offset from global memory to shared memory >
+
+    // Calculate softmax over L and K
+    if (threadIdx.x == 0) { // gi != 0, di = 0, li = 0
+      opmath_t softmax_max = -1e100;
+      opmath_t softmax_sum = 0.0;
+
+      // get max
+      for (int j = 0; j < K; j++) {
+        softmax_max = max(softmax_max, p_mask_shm[j]);
+      }
+
+      // get sumexp
+      for (int j = 0; j < K; j++) {
+        opmath_t exp_results = exp(p_mask_shm[j] - softmax_max);
+        p_mask_shm[j] = exp_results;
+        softmax_sum += exp_results;
+      }
+
+      // normalize
+      for (int j = 0; j < K; j++) {
+        p_mask_shm[j] /= softmax_sum;
+      }
+    }
+
+    __syncthreads();
+  }
+
+  int offset_idx = 0;
+  int mask_idx = 0;
+
+  const int w_stride = G * D;
+  const int base_ptr = gi * D + di_s;
+  const scalar_t *p_value_ptr =
+      p_value + (bi * (height_in * width_in)) * (G * D);
+  opmath_t *grad_im_ptr = grad_im + (bi * (height_in * width_in)) * (G * D);
+
+  const int p0_w = ((dilation_w * (kernel_w - 1)) >> 1) - pad_w +
+                   (qi % width_out) * stride_w;
+  const int p0_h = ((dilation_h * (kernel_h - 1)) >> 1) - pad_h +
+                   (qi / width_out) * stride_h;
+  const opmath_t p0_w_ =
+      p0_w - ((dilation_w * (kernel_w - 1)) >> 1) * offset_scale;
+  const opmath_t p0_h_ =
+      p0_h - ((dilation_h * (kernel_h - 1)) >> 1) * offset_scale;
+  const int center_h = kernel_h / 2;
+  const int center_w = kernel_w / 2;
+
+  grad_offset += (bi*Q + qi)*padded_offset_dim + gi*K*3;
+  opmath_t *grad_offset_softmax = grad_offset + K * 2;
+
+  int cache_grad_off_idx =
+      ((threadIdx.z * G + threadIdx.y) * blockDim.x + threadIdx.x) * 3;
+  for (int i = 0; i < kernel_w; ++i) {
+    for (int j = 0; j < kernel_h; ++j) {
+      if (i != center_w || j != center_h || !remove_center) {
+        const opmath_t w_im =
+            p0_w_ + (i * dilation_w + (opmath_t)p_offset_ptr[offset_idx]) *
+                        offset_scale;
+        const opmath_t h_im =
+            p0_h_ + (j * dilation_h + (opmath_t)p_offset_ptr[offset_idx + 1]) *
+                        offset_scale;
+        const opmath_t attn = p_mask_shm[mask_idx];
+        cache_grad_offset[cache_grad_off_idx] = 0;
+        cache_grad_offset[cache_grad_off_idx + 1] = 0;
+        cache_grad_offset[cache_grad_off_idx + 2] = 0;
+
+        if (h_im > -1 && w_im > -1 && h_im < height_in && w_im < width_in) {
+          ms_deform_attn_col2im_bilinear<scalar_t, transfer_t, d_stride>(
+              p_value_ptr, height_in, width_in, h_im, w_im, attn, w_stride,
+              base_ptr, offset_scale, offset_scale, top_grad, grad_im_ptr,
+              cache_grad_offset + cache_grad_off_idx);
+        }
+
+        // aggregated across different channel for offset
+        
+        __syncthreads();
+        if (threadIdx.x == 0) { //
+          int _didx = (threadIdx.z * G + threadIdx.y) * blockDim.x * 3;
+          opmath_t _grad_w = cache_grad_offset[_didx],
+                   _grad_h = cache_grad_offset[_didx + 1],
+                   _grad_a = cache_grad_offset[_didx + 2];
+
+          for (int c_id = 1; c_id < blockDim.x; ++c_id) {
+            _grad_w += cache_grad_offset[_didx + 3 * c_id];
+            _grad_h += cache_grad_offset[_didx + 3 * c_id + 1];
+            _grad_a += cache_grad_offset[_didx + 3 * c_id + 2];
+          }
+
+          *(grad_offset) = _grad_w;     // B x H x W x G x L x K x 3
+          *(grad_offset + 1) = _grad_h; // B x H x W x G x L x K x 3
+          if (softmax) {
+            cache_g_mask_before_softmax[(threadIdx.z * G + threadIdx.y) * K +
+                                        mask_idx] = _grad_a * attn;
+          }
+          else{
+            grad_offset_softmax[mask_idx] = _grad_a;
+          }
+        }
+        __syncthreads();
+
+        offset_idx += 2;
+        mask_idx += 1;
+        grad_offset += 2;
+      }
+    }
+  }
+  // backward for softmax
+  if(softmax){
+    if (threadIdx.x == 0) {
+      const opmath_t* group_g_mask = cache_g_mask_before_softmax + (threadIdx.z*G + threadIdx.y)*K;
+      #pragma unroll
+      for (int i = 0; i < K; ++i) {
+        opmath_t sum = 0.;
+        for (int j = 0; j < K; ++j) {
+          sum += group_g_mask[j]; // dL/di * di/dj
+        }
+        *(grad_offset_softmax) = group_g_mask[i] - p_mask_shm[i] * sum;
+
+        grad_offset_softmax += 1;
+      }
+    }
+    __syncthreads();
+  }
+}
+
+template <typename scalar_t, int d_stride, typename transfer_t, int L, int K,
+          bool softmax>
+__global__ void backward_kernel_dcn_warp_primitive(
+    const scalar_t *p_value, const scalar_t *p_offset,
+    const scalar_t *grad_output, const int G, const int D, const int Q,
+    const int kernel_h, const int kernel_w, const int stride_h,
+    const int stride_w, const int pad_h, const int pad_w, const int dilation_h,
+    const int dilation_w, const int height_in, const int width_in,
+    const int height_out, const int width_out, const opmath_t offset_scale,
+    const int remove_center, const int block_multiplier, opmath_t *grad_im,
+    opmath_t *grad_offset, const int padded_offset_dim) {
+
+  extern __shared__ char _s[];
+
+  const int &qi = (blockIdx.x * block_multiplier % Q) + threadIdx.z;
+  const int &bi = blockIdx.x * block_multiplier / Q;
+
+
+  const int &di_s = threadIdx.x * d_stride;
+  const int &gi = threadIdx.y;
+
+  constexpr int li = 0;
+  const int tid = (threadIdx.z * blockDim.y + threadIdx.y)*blockDim.x + threadIdx.x;
+
+  const int lane_id = tid % kWarpSize;
+
+  // find the position of current group in the current warp
+  const int group_per_warp = kWarpSize / blockDim.x;
+  const int group_in_warp_id = (threadIdx.z * G + threadIdx.y) % group_per_warp;
+  const unsigned lane_mask = ((1 << blockDim.x) - 1) << (group_in_warp_id * blockDim.x);
+
+  opmath_t *const p_mask_shm = (opmath_t *)(_s) + (threadIdx.z * G + gi) * K;
+  opmath_t *cache_g_mask_before_softmax = (opmath_t *)((opmath_t *)(_s) + block_multiplier * G * K) +
+                                          (threadIdx.z*G+gi)*K; // only used by threadIdx.x = 0
+
+  const scalar_t *p_offset_ptr = p_offset + (bi*Q + qi)*padded_offset_dim + gi*K*3;
+
+  const int mask_length = K;
+  const int num_thread = (D / d_stride);
+  const int num_iter = mask_length / num_thread;
+  const int remainder = mask_length - num_iter * num_thread;
+
+  const scalar_t *top_grad = grad_output + ((bi * Q + qi) * G + gi) * D + di_s;
+
+  __syncthreads();
+  for (int i = 0; i < num_iter; i++) {
+    *(p_mask_shm + num_thread * i + threadIdx.x) =
+        *(scalar_t *)(p_offset_ptr + K * 2 + num_thread * i + threadIdx.x);
+  }
+  if (remainder > 0 && threadIdx.x < remainder) {
+    *(p_mask_shm + num_thread * num_iter + threadIdx.x) = *(
+        scalar_t *)(p_offset_ptr + K * 2 + num_thread * num_iter + threadIdx.x);
+  }
+
+  if (softmax) {
+    __syncthreads();
+    // transfer offset from global memory to shared memory >
+
+    // Calculate softmax over L and K
+    if (threadIdx.x == 0) { // gi != 0, di = 0, li = 0
+      opmath_t softmax_max = -1e100;
+      opmath_t softmax_sum = 0.0;
+
+      // get max
+      for (int j = 0; j < K; j++) {
+        softmax_max = max(softmax_max, p_mask_shm[j]);
+      }
+
+      // get sumexp
+      for (int j = 0; j < K; j++) {
+        opmath_t exp_results = exp(p_mask_shm[j] - softmax_max);
+        p_mask_shm[j] = exp_results;
+        softmax_sum += exp_results;
+      }
+
+      // normalize
+      for (int j = 0; j < K; j++) {
+        p_mask_shm[j] /= softmax_sum;
+      }
+    }
+
+    __syncthreads();
+  }
+
+  int offset_idx = 0;
+  int mask_idx = 0;
+
+  const int w_stride = G * D;
+  const int base_ptr = gi * D + di_s;
+  const scalar_t *p_value_ptr =
+      p_value + (bi * (height_in * width_in)) * (G * D);
+  opmath_t *grad_im_ptr = grad_im + (bi * (height_in * width_in)) * (G * D);
+
+  const int p0_w = ((dilation_w * (kernel_w - 1)) >> 1) - pad_w +
+                   (qi % width_out) * stride_w;
+  const int p0_h = ((dilation_h * (kernel_h - 1)) >> 1) - pad_h +
+                   (qi / width_out) * stride_h;
+  const opmath_t p0_w_ =
+      p0_w - ((dilation_w * (kernel_w - 1)) >> 1) * offset_scale;
+  const opmath_t p0_h_ =
+      p0_h - ((dilation_h * (kernel_h - 1)) >> 1) * offset_scale;
+  const int center_h = kernel_h / 2;
+  const int center_w = kernel_w / 2;
+
+  grad_offset += (bi * Q + qi)*padded_offset_dim + gi*K*3;
+  opmath_t *grad_offset_softmax = grad_offset + K * 2;
+
+  int cache_grad_off_idx =
+      ((threadIdx.z * G + threadIdx.y) * blockDim.x + threadIdx.x) * 3;
+
+  opmath_t reg_grad_offset[3] = {0.};
+  for (int i = 0; i < kernel_w; ++i) {
+    for (int j = 0; j < kernel_h; ++j) {
+      if (i != center_w || j != center_h || !remove_center) {
+        const opmath_t w_im =
+            p0_w_ + (i * dilation_w + (opmath_t)p_offset_ptr[offset_idx]) *
+                        offset_scale;
+        const opmath_t h_im =
+            p0_h_ + (j * dilation_h + (opmath_t)p_offset_ptr[offset_idx + 1]) *
+                        offset_scale;
+        const opmath_t attn = p_mask_shm[mask_idx];
+        reg_grad_offset[0] = 0;
+        reg_grad_offset[1] = 0;
+        reg_grad_offset[2] = 0;
+
+        if (h_im > -1 && w_im > -1 && h_im < height_in && w_im < width_in) {
+          ms_deform_attn_col2im_bilinear<scalar_t, transfer_t, d_stride>(
+              p_value_ptr, height_in, width_in, h_im, w_im, attn, w_stride,
+              base_ptr, offset_scale, offset_scale, top_grad, grad_im_ptr,
+              reg_grad_offset);
+        }
+
+        // aggregated across different channel for offset
+        for (uint32_t offset = blockDim.x>>1; offset > 0; offset >>= 1){
+          reg_grad_offset[0] += __shfl_down_sync(lane_mask, reg_grad_offset[0], offset);
+          reg_grad_offset[1] += __shfl_down_sync(lane_mask, reg_grad_offset[1], offset);
+          reg_grad_offset[2] += __shfl_down_sync(lane_mask, reg_grad_offset[2], offset);
+        }
+
+        if (threadIdx.x == 0) { //
+          *(grad_offset) = reg_grad_offset[0];     // B x H x W x G x L x K x 3
+          *(grad_offset + 1) = reg_grad_offset[1]; // B x H x W x G x L x K x 3
+          if (softmax) {
+            cache_g_mask_before_softmax[mask_idx] = reg_grad_offset[2] * attn;
+          }
+          else{
+            grad_offset_softmax[mask_idx] = reg_grad_offset[2];
+          }
+        }
+        offset_idx += 2;
+        mask_idx += 1;
+        grad_offset += 2;
+      }
+    }
+  }
+  // backward for softmax
+  if(softmax){
+    if (threadIdx.x == 0) {
+      opmath_t sum = 0.;
+      #pragma unroll
+      for (int i=0; i < K; ++i){
+        sum += cache_g_mask_before_softmax[i];
+      }
+      #pragma unroll
+      for (int i = 0; i < K; ++i) {
+        *(grad_offset_softmax) = cache_g_mask_before_softmax[i] - p_mask_shm[i] * sum;
+        grad_offset_softmax += 1;
+      }
+    }
+  }
+}
+
+template <typename scalar_t, typename stride_type, int d_stride>
+void _dcnv4_col2im_cuda(
+    cudaStream_t stream,
+    const scalar_t *value,       // B, H * W, (G * D)
+    const scalar_t *p_offset,    // B, H * W, (G*K*3)
+    const scalar_t *grad_output, // B, H_out*W_out, G * D
+    const int kernel_h, const int kernel_w, const int stride_h,
+    const int stride_w, const int pad_h, const int pad_w, const int dilation_h,
+    const int dilation_w, const int G, const int D, const int B,
+    const int height_in, const int width_in, const int height_out,
+    const int width_out, const opmath_t offset_scale, const int remove_center,
+    opmath_t *grad_im, opmath_t *grad_offset, const int block_thread, 
+    const bool softmax, const int padded_offset_dim) {
+
+  constexpr int L = 1;
+
+  auto kernel =
+      backward_kernel_dcn_warp_primitive<scalar_t, d_stride, stride_type, 1, 9, false>;
+
+  int N = height_in * width_in;
+  int Q = height_out * width_out;
+  int K = kernel_h * kernel_w;
+
+  if (remove_center) {
+    K -= 1;
+  }
+
+  if (softmax) {
+    switch (K) {
+    case 9:
+      if(check_backward_warpp(d_stride, D)){
+        kernel = backward_kernel_dcn_warp_primitive<scalar_t, d_stride, stride_type, 1, 9, true>;
+      }
+      else{
+        kernel = backward_kernel_dcn<scalar_t, d_stride, stride_type, 1, 9, true>;
+      }
+      break;
+    case 8:
+      if(check_backward_warpp(d_stride, D)){
+        kernel = backward_kernel_dcn_warp_primitive<scalar_t, d_stride, stride_type, 1, 8, true>;
+      }
+      else {
+        kernel = backward_kernel_dcn<scalar_t, d_stride, stride_type, 1, 8, true>;
+      }
+      break;
+    default:
+      printf("K=%ld\n", K);
+      throw std::invalid_argument("invalid kernel shape");
+    }
+  } else {
+    switch (K) {
+    case 9:
+      if(check_backward_warpp(d_stride, D)){
+        kernel = backward_kernel_dcn_warp_primitive<scalar_t, d_stride, stride_type, 1, 9, false>;
+      }
+      else{
+        kernel = backward_kernel_dcn<scalar_t, d_stride, stride_type, 1, 9, false>;
+      }
+      break;
+    case 8:
+      if(check_backward_warpp(d_stride, D)){
+        kernel = backward_kernel_dcn_warp_primitive<scalar_t, d_stride, stride_type, 1, 8, false>;
+      }
+      else {
+        kernel = backward_kernel_dcn<scalar_t, d_stride, stride_type, 1, 8, false>;
+      }
+      break;
+    default:
+      printf("K=%ld\n", K);
+      throw std::invalid_argument("invalid kernel shape");
+    }
+  }
+
+  const int block_multiplier = block_thread / (D / d_stride) / G;
+  assert((B*Q) % block_multiplier == 0);
+
+  dim3 num_blocks(B*Q / block_multiplier);
+  dim3 num_threads(D / d_stride, G, block_multiplier);
+
+  const int blockdimX = D / d_stride;
+
+  int shm_size = sizeof(opmath_t) * (G * block_multiplier * K) * 2;
+  if(!check_backward_warpp(d_stride, D)){
+    shm_size = sizeof(opmath_t) * ((G * block_multiplier * K) * 2 + G * block_multiplier * blockdimX * 3);
+  }
+
+  cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
+                       shm_size);
+
+  kernel<<<num_blocks, num_threads, shm_size, stream>>>(
+      value, p_offset, grad_output, G, D, Q, kernel_h, kernel_w, stride_h,
+      stride_w, pad_h, pad_w, dilation_h, dilation_w, height_in, width_in,
+      height_out, width_out, offset_scale, remove_center, block_multiplier,
+      grad_im, grad_offset, padded_offset_dim);
+
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess) {
+    printf("error in dcnv4_im2col_cuda: %s\n", cudaGetErrorString(err));
+    printf("launch arguments: gridDim=(%d, %d, %d), blockDim=(%d, %d, %d), "
+           "shm_size=%d\n\n",
+           num_blocks.x, num_blocks.y, num_blocks.z, num_threads.x,
+           num_threads.y, num_threads.z, shm_size);
+    AT_ASSERTM(false, "kernel launch error");
+  }
+}
+
+template <typename scalar_t>
+void dcnv4_col2im_cuda(
+    cudaStream_t stream,
+    const scalar_t *value,       // B, H * W, (G * D)
+    const scalar_t *p_offset,    // B, H * W, (G*K*3)
+    const scalar_t *grad_output, // B, H_out*W_out, G * D
+    const int kernel_h, const int kernel_w, const int stride_h,
+    const int stride_w, const int pad_h, const int pad_w, const int dilation_h,
+    const int dilation_w, const int G, const int D, const int B,
+    const int height_in, const int width_in, const int height_out,
+    const int width_out, const opmath_t offset_scale, const int remove_center,
+    opmath_t *grad_im, opmath_t *grad_offset, const int d_stride,
+    const int block_thread, const bool softmax, const int padded_offset_dim) {
+
+  assert(D % d_stride == 0);
+  const int size_scalar = sizeof(scalar_t);
+  if (size_scalar == 2) {
+    switch (d_stride) {
+    case 1:
+      _dcnv4_col2im_cuda<scalar_t, scalar_t, 1>(
+          stream, value, p_offset, grad_output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center, grad_im,
+          grad_offset, block_thread, softmax, padded_offset_dim);
+      break;
+    case 2:
+      _dcnv4_col2im_cuda<scalar_t, uint, 2>(
+          stream, value, p_offset, grad_output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center, grad_im,
+          grad_offset, block_thread, softmax, padded_offset_dim);
+      break;
+    case 4:
+      _dcnv4_col2im_cuda<scalar_t, uint2, 4>(
+          stream, value, p_offset, grad_output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center, grad_im,
+          grad_offset, block_thread, softmax, padded_offset_dim);
+      break;
+    case 8:
+      _dcnv4_col2im_cuda<scalar_t, uint4, 8>(
+          stream, value, p_offset, grad_output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center, grad_im,
+          grad_offset, block_thread, softmax, padded_offset_dim);
+      break;
+    case 16:
+      _dcnv4_col2im_cuda<scalar_t, ulonglong4, 16>(
+          stream, value, p_offset, grad_output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center, grad_im,
+          grad_offset, block_thread, softmax, padded_offset_dim);
+      break;
+    }
+  } else {
+    assert(size_scalar == 4);
+    switch (d_stride) {
+    case 1:
+      _dcnv4_col2im_cuda<scalar_t, uint, 1>(
+          stream, value, p_offset, grad_output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center, grad_im,
+          grad_offset, block_thread, softmax, padded_offset_dim);
+      break;
+    case 2:
+      _dcnv4_col2im_cuda<scalar_t, uint2, 2>(
+          stream, value, p_offset, grad_output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center, grad_im,
+          grad_offset, block_thread, softmax, padded_offset_dim);
+      break;
+    case 4:
+      _dcnv4_col2im_cuda<scalar_t, uint4, 4>(
+          stream, value, p_offset, grad_output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center, grad_im,
+          grad_offset, block_thread, softmax, padded_offset_dim);
+      break;
+    case 8:
+      _dcnv4_col2im_cuda<scalar_t, ulonglong4, 8>(
+          stream, value, p_offset, grad_output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center, grad_im,
+          grad_offset, block_thread, softmax, padded_offset_dim);
+      break;
+    }
+  }
+}

DCNv4_op/src/cuda/dcnv4_cuda.cu

@@ -0,0 +1,175 @@
+/*!
+**************************************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************************************
+* Modified from
+*https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+#include "cuda/dcnv4_im2col_cuda.cuh"
+#include "cuda/dcnv4_col2im_cuda.cuh"
+#include <vector>
+
+#include <ATen/ATen.h>
+#include <ATen/OpMathType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cooperative_groups.h>
+#include <cooperative_groups/memcpy_async.h>
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <torch/torch.h>
+
+at::Tensor dcnv4_cuda_forward(
+    const at::Tensor &value,
+    const at::Tensor &p_offset,
+    const int kernel_h, const int kernel_w, const int stride_h,
+    const int stride_w, const int pad_h, const int pad_w, const int dilation_h,
+    const int dilation_w, const int group, const int group_channels,
+    const float offset_scale, const int im2col_step, const int remove_center,
+    const int d_stride, const int block_thread, const bool softmax) {
+  AT_ASSERTM(value.is_contiguous(), "input tensor has to be contiguous");
+  AT_ASSERTM(value.type().is_cuda(), "input must be a CUDA tensor");
+  AT_ASSERTM(p_offset.is_contiguous(), "input tensor has to be contiguous");
+  AT_ASSERTM(p_offset.type().is_cuda(), "input must be a CUDA tensor");
+
+  const int batch = value.size(0);
+  const int height_in = value.size(1);
+  const int width_in = value.size(2);
+  const int channels = value.size(3);
+  const int padded_offset_dim = p_offset.size(3);
+
+  // tensor core requirement
+  assert(padded_offset_dim % 8 == 0);
+
+  const int height_out =
+      (height_in + 2 * pad_h - (dilation_h * (kernel_h - 1) + 1)) / stride_h +
+      1;
+  const int width_out =
+      (width_in + 2 * pad_w - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;
+
+  const int im2col_step_ = std::min(batch, im2col_step);
+  AT_ASSERTM(batch % im2col_step_ == 0, "batch(", batch,
+             ") must divide im2col_step(", im2col_step_, ")");
+  AT_ASSERTM(
+      channels == (group * group_channels),
+      "Input channels and group times group channels wont match: (%d vs %d).",
+      channels, group * group_channels);
+
+  auto output = at::zeros(
+      {batch, height_out, width_out, group * group_channels}, value.options());
+
+  const int batch_n = im2col_step_;
+  auto output_n = output.view({batch / batch_n, batch_n, height_out, width_out,
+                               group * group_channels});
+  auto per_value_size = height_in * width_in * channels;
+  auto per_offset_size = height_out * width_out * padded_offset_dim;
+
+  for (int n = 0; n < batch / im2col_step_; ++n) {
+    auto columns = output_n.select(0, n);
+    AT_DISPATCH_FLOATING_TYPES_AND2(
+        at::ScalarType::Half, at::ScalarType::BFloat16, value.scalar_type(),
+        "dcnv4_forward_cuda", ([&] {
+          dcnv4_im2col_cuda(
+              at::cuda::getCurrentCUDAStream(),
+              value.data_ptr<scalar_t>() + n * im2col_step_ * per_value_size,
+              p_offset.data_ptr<scalar_t>() +
+                  n * im2col_step_ * per_offset_size,
+              columns.data_ptr<scalar_t>(), kernel_h, kernel_w, stride_h,
+              stride_w, pad_h, pad_w, dilation_h, dilation_w, group,
+              group_channels, batch_n, height_in, width_in, height_out,
+              width_out, offset_scale, remove_center, d_stride, block_thread,
+              softmax, padded_offset_dim);
+        }));
+  }
+
+  return output;
+}
+
+std::vector<at::Tensor>
+dcnv4_cuda_backward(
+    const at::Tensor &value, 
+    const at::Tensor &p_offset, 
+    const int kernel_h, const int kernel_w, const int stride_h,
+    const int stride_w, const int pad_h, const int pad_w, const int dilation_h,
+    const int dilation_w, const int group, const int group_channels,
+    const float offset_scale, const int im2col_step, const at::Tensor &grad_output, 
+    const int remove_center, const int d_stride, const int block_thread,
+    const bool softmax) {
+  AT_ASSERTM(value.is_contiguous(), "input tensor has to be contiguous");
+  AT_ASSERTM(p_offset.is_contiguous(), "offset tensor has to be contiguous");
+  AT_ASSERTM(grad_output.is_contiguous(), 
+             "grad_output tensor has to be contiguous");
+
+  AT_ASSERTM(value.type().is_cuda(), "input must be a CUDA tensor");
+  AT_ASSERTM(p_offset.type().is_cuda(), "offset must be a CUDA tensor");
+  AT_ASSERTM(grad_output.type().is_cuda(),
+             "grad_output must be a CUDA tensor");
+
+  const int batch = value.size(0);
+  const int height_in = value.size(1);
+  const int width_in = value.size(2);
+  const int channels = value.size(3);
+  const int padded_offset_dim = p_offset.size(3);
+  assert(padded_offset_dim % 8 == 0);
+
+  const int height_out =
+      (height_in + 2 * pad_h - (dilation_h * (kernel_h - 1) + 1)) / stride_h +
+      1;
+  const int width_out =
+      (width_in + 2 * pad_w - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;
+
+  const int im2col_step_ = std::min(batch, im2col_step);
+  AT_ASSERTM(batch % im2col_step_ == 0, "batch(", batch,
+             ") must divide im2col_step(", im2col_step_, ")");
+  AT_ASSERTM(
+      channels == (group * group_channels),
+      "Input channels and group times group channels wont match: (%d vs %d).",
+      channels, group * group_channels);
+  
+  auto dtype = value.dtype();
+  if (dtype == at::kHalf){
+    dtype = at::kFloat;
+  }
+
+  auto grad_input = at::zeros_like(value, dtype);
+  auto grad_offset = at::zeros_like(p_offset, dtype);
+
+  const int batch_n = im2col_step_;
+  auto grad_output_n = grad_output.view({batch / batch_n, batch_n, height_out, width_out,
+                               group, group_channels});
+  auto per_value_size = height_in * width_in * channels;
+  auto per_offset_size = height_out * width_out * padded_offset_dim;
+
+  for (int n = 0; n < batch / im2col_step_; ++n) {
+    auto columns = grad_output_n.select(0, n);
+    AT_DISPATCH_FLOATING_TYPES_AND_HALF(value.type(),
+        "dcnv4_backward_cuda", ([&] {
+          dcnv4_col2im_cuda(
+              at::cuda::getCurrentCUDAStream(),
+              value.data_ptr<scalar_t>() + n * im2col_step_ * per_value_size,
+              p_offset.data_ptr<scalar_t>() +
+                  n * im2col_step_ * per_offset_size,
+              columns.data_ptr<scalar_t>(), kernel_h, kernel_w, stride_h,
+              stride_w, pad_h, pad_w, dilation_h, dilation_w, group,
+              group_channels, batch_n, height_in, width_in, height_out,
+              width_out, offset_scale, remove_center,
+              grad_input.data<opmath_t>() + n * im2col_step_ * per_value_size,
+              grad_offset.data<opmath_t>() +
+                  n * im2col_step_ * per_offset_size,
+              d_stride, block_thread, softmax, padded_offset_dim
+              );
+        }));
+  }
+
+  if(value.dtype() == torch::kHalf){
+    return {grad_input.to(torch::kHalf), grad_offset.to(torch::kHalf)};
+  }
+  else{
+    return {grad_input, grad_offset};
+  }
+}
+

DCNv4_op/src/cuda/dcnv4_cuda.h

@@ -0,0 +1,33 @@
+/*!
+**************************************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************************************
+* Modified from
+*https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+#pragma once
+#include <torch/extension.h>
+
+at::Tensor dcnv4_cuda_forward(
+    const at::Tensor &value,
+    const at::Tensor &p_offset,
+    const int kernel_h, const int kernel_w, const int stride_h,
+    const int stride_w, const int pad_h, const int pad_w, const int dilation_h,
+    const int dilation_w, const int group, const int group_channels,
+    const float offset_scale, const int im2col_step, const int remove_center,
+    const int d_stride, const int block_thread, const bool softmax);
+
+std::vector<at::Tensor>
+dcnv4_cuda_backward(
+    const at::Tensor &value,
+    const at::Tensor &p_offset,
+    const int kernel_h, const int kernel_w, const int stride_h,
+    const int stride_w, const int pad_h, const int pad_w, const int dilation_h,
+    const int dilation_w, const int group, const int group_channels,
+    const float offset_scale, const int im2col_step, const at::Tensor &grad_output, 
+    const int remove_center, const int d_stride, const int block_thread, 
+    const bool softmax);

DCNv4_op/src/cuda/dcnv4_im2col_cuda.cuh

@@ -0,0 +1,416 @@
+#include <algorithm>
+#include <cstdio>
+#include <cstring>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <THC/THCAtomics.cuh>
+
+#include <ATen/ATen.h>
+#include <ATen/OpMathType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cooperative_groups.h>
+#include <cooperative_groups/memcpy_async.h>
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include "common.h"
+
+template <typename scalar_t, int d_stride, typename transfer_t, int L, int K,
+          bool softmax>
+__global__ void forward_kernel_dcn(
+    const scalar_t *p_value, const scalar_t *p_offset, scalar_t *p_output,
+    const int G, const int D, const int Q, const int kernel_h,
+    const int kernel_w, const int stride_h, const int stride_w, const int pad_h,
+    const int pad_w, const int dilation_h, const int dilation_w,
+    const int height_in, const int width_in, const int height_out,
+    const int width_out, const opmath_t offset_scale, const int remove_center,
+    const int block_multiplier, const int padded_offset_dim) {
+
+  const int &qi = (blockIdx.x * block_multiplier % Q) + threadIdx.z;
+  const int &bi = blockIdx.x * block_multiplier / Q;
+
+  const int &di_s = threadIdx.x * d_stride;
+  const int &gi = threadIdx.y;
+  constexpr int li = 0;
+
+  extern __shared__ char _s[];
+  opmath_t *const p_mask_shm =
+      (opmath_t *)(_s) + ((threadIdx.z * G + gi) * L + li) * K;
+
+  opmath_t p_out_shm[d_stride] = {0.};
+
+  const scalar_t *p_offset_ptr = p_offset + (bi*Q + qi)*padded_offset_dim + gi*K*3;
+
+  const int mask_length = K;
+  const int num_thread = (D / d_stride);
+  const int num_iter = mask_length / num_thread;
+  const int remainder = mask_length - num_iter * num_thread;
+
+  for (int i = 0; i < num_iter; i++) {
+    *(p_mask_shm + num_thread * i + threadIdx.x) =
+        *(scalar_t *)(p_offset_ptr + K * 2 + num_thread * i + threadIdx.x);
+  }
+  if (remainder > 0 && threadIdx.x < remainder) {
+    *(p_mask_shm + num_thread * num_iter + threadIdx.x) = *(
+        scalar_t *)(p_offset_ptr + K * 2 + num_thread * num_iter + threadIdx.x);
+  }
+
+  int mask_idx;
+  if (softmax) {
+    __syncthreads();
+
+    // Calculate softmax over L and K
+    if (threadIdx.x == 0) { // gi != 0, di = 0, li = 0
+      opmath_t softmax_max = -1e100;
+      opmath_t softmax_sum = 0.0;
+
+      // get max
+      // #pragma unroll
+      for (int j = 0; j < K; j++) {
+        softmax_max = max(softmax_max, p_mask_shm[j]);
+      }
+
+      // get sumexp
+      // #pragma unroll
+      for (int j = 0; j < K; j++) {
+        opmath_t exp_results = exp(p_mask_shm[j] - softmax_max);
+        p_mask_shm[j] = exp_results;
+        softmax_sum += exp_results;
+      }
+
+      // normalize
+      // #pragma unroll
+      for (int j = 0; j < K; j++) {
+        p_mask_shm[j] /= softmax_sum;
+      }
+    }
+
+    __syncthreads();
+  }
+  int offset_idx = 0;
+  mask_idx = 0;
+
+  const int w_stride = G * D;
+  const int base_ptr = gi * D + di_s;
+  const scalar_t *p_value_ptr =
+      p_value + (bi * (height_in * width_in)) * (G * D);
+
+  const int p0_w = ((dilation_w * (kernel_w - 1)) >> 1) - pad_w +
+                   (qi % width_out) * stride_w;
+  const int p0_h = ((dilation_h * (kernel_h - 1)) >> 1) - pad_h +
+                   (qi / width_out) * stride_h;
+  const opmath_t p0_w_ =
+      p0_w - ((dilation_w * (kernel_w - 1)) >> 1) * offset_scale;
+  const opmath_t p0_h_ =
+      p0_h - ((dilation_h * (kernel_h - 1)) >> 1) * offset_scale;
+  const int center_h = kernel_h / 2;
+  const int center_w = kernel_w / 2;
+
+  int out_idx = ((bi * Q + qi) * G + gi) * D + di_s;
+
+  for (int i = 0; i < kernel_w; ++i) {
+    for (int j = 0; j < kernel_h; ++j) {
+      if (i != center_w || j != center_h || !remove_center) {
+        const opmath_t w_im =
+            p0_w_ + (i * dilation_w + (opmath_t)p_offset_ptr[offset_idx]) *
+                        offset_scale;
+        const opmath_t h_im =
+            p0_h_ + (j * dilation_h + (opmath_t)p_offset_ptr[offset_idx + 1]) *
+                        offset_scale;
+        const opmath_t attn = p_mask_shm[mask_idx];
+
+        if (h_im > -1 && w_im > -1 && h_im < height_in && w_im < width_in) {
+          ms_deform_attn_im2col_bilinear<scalar_t, transfer_t, d_stride>(
+              p_out_shm, p_value_ptr, height_in, width_in, h_im, w_im, attn,
+              w_stride, base_ptr);
+        }
+        offset_idx += 2;
+        mask_idx += 1;
+      }
+    }
+  }
+  scalar_t *fp16_regs = (scalar_t *)(p_out_shm);
+#pragma unroll
+  for (int ds = 0; ds < d_stride; ds++) {
+    fp16_regs[ds] = p_out_shm[ds];
+  }
+  *(transfer_t *)(p_output + out_idx) = *(transfer_t *)(p_out_shm);
+}
+
+template <typename scalar_t, int d_stride, typename transfer_t, int L, int K,
+          bool softmax>
+__global__ void forward_kernel_dcn_reg(
+    const scalar_t *p_value, const scalar_t *p_offset, scalar_t *p_output,
+    const int G, const int D, const int Q, const int kernel_h,
+    const int kernel_w, const int stride_h, const int stride_w, const int pad_h,
+    const int pad_w, const int dilation_h, const int dilation_w,
+    const int height_in, const int width_in, const int height_out,
+    const int width_out, const opmath_t offset_scale, const int remove_center,
+    const int block_multiplier, const int padded_offset_dim) {
+
+  const int &qi = (blockIdx.x * block_multiplier % Q) + threadIdx.z;
+  const int &bi = blockIdx.x * block_multiplier / Q;
+
+  const int &di_s = threadIdx.x * d_stride;
+  const int &gi = threadIdx.y;
+  constexpr int li = 0;
+
+  opmath_t p_mask_shm[K] = {0.};
+  opmath_t p_out_shm[d_stride] = {0.};
+
+  const scalar_t *p_offset_ptr = p_offset + (bi*Q + qi)*padded_offset_dim + gi*K*3;
+  const int mask_length = K;
+  const int num_thread = (D / d_stride);
+  const int num_iter = mask_length / num_thread;
+  const int remainder = mask_length - num_iter * num_thread;
+
+  for (int i=0; i < K; i++){
+    p_mask_shm[i] = *(p_offset_ptr + K*2 + i);
+  }
+
+  if (softmax) {
+    // Calculate softmax over L and K
+      opmath_t softmax_max = -1e100;
+      opmath_t softmax_sum = 0.0;
+      // get max
+      for (int j = 0; j < K; j++) {
+        softmax_max = max(softmax_max, p_mask_shm[j]);
+      }
+
+      // get sumexp
+      for (int j = 0; j < K; j++) {
+        opmath_t exp_results = exp(p_mask_shm[j] - softmax_max);
+        p_mask_shm[j] = exp_results;
+        softmax_sum += exp_results;
+      }
+
+      // normalize
+      for (int j = 0; j < K; j++) {
+        p_mask_shm[j] /= softmax_sum;
+      }
+  }
+
+  int offset_idx = 0;
+  int mask_idx = 0;
+
+  const int w_stride = G * D;
+  const int base_ptr = gi * D + di_s;
+  const scalar_t *p_value_ptr =
+      p_value + (bi * (height_in * width_in)) * (G * D);
+
+  const int p0_w = ((dilation_w * (kernel_w - 1)) >> 1) - pad_w +
+                   (qi % width_out) * stride_w;
+  const int p0_h = ((dilation_h * (kernel_h - 1)) >> 1) - pad_h +
+                   (qi / width_out) * stride_h;
+  const opmath_t p0_w_ =
+      p0_w - ((dilation_w * (kernel_w - 1)) >> 1) * offset_scale;
+  const opmath_t p0_h_ =
+      p0_h - ((dilation_h * (kernel_h - 1)) >> 1) * offset_scale;
+  const int center_h = kernel_h / 2;
+  const int center_w = kernel_w / 2;
+
+  int out_idx = ((bi * Q + qi) * G + gi) * D + di_s;
+
+  for (int i = 0; i < kernel_w; ++i) {
+    for (int j = 0; j < kernel_h; ++j) {
+      if (i != center_w || j != center_h || !remove_center) {
+        const opmath_t w_im =
+            p0_w_ + (i * dilation_w + (opmath_t)p_offset_ptr[offset_idx]) *
+                        offset_scale;
+        const opmath_t h_im =
+            p0_h_ + (j * dilation_h + (opmath_t)p_offset_ptr[offset_idx + 1]) *
+                        offset_scale;
+        const opmath_t attn = p_mask_shm[mask_idx];
+
+        if (h_im > -1 && w_im > -1 && h_im < height_in && w_im < width_in) {
+          ms_deform_attn_im2col_bilinear<scalar_t, transfer_t, d_stride>(
+              p_out_shm, p_value_ptr, height_in, width_in, h_im, w_im, attn,
+              w_stride, base_ptr);
+        }
+        offset_idx += 2;
+        mask_idx += 1;
+      }
+    }
+  }
+  scalar_t *fp16_regs = (scalar_t *)(p_out_shm);
+#pragma unroll
+  for (int ds = 0; ds < d_stride; ds++) {
+    fp16_regs[ds] = p_out_shm[ds];
+  }
+
+  *(transfer_t *)(p_output + out_idx) = *(transfer_t *)(p_out_shm);
+}
+
+template <typename scalar_t, typename stride_type, int d_stride>
+void _dcnv4_im2col_cuda(cudaStream_t stream,
+                              const scalar_t *value,    // B, H * W, (G * D)
+                              const scalar_t *p_offset, // B, H * W, G * K * 3)
+                              scalar_t *output,         // B, H_out*W_out, G * D
+                              const int kernel_h, const int kernel_w,
+                              const int stride_h, const int stride_w,
+                              const int pad_h, const int pad_w,
+                              const int dilation_h, const int dilation_w,
+                              const int G, const int D, const int B,
+                              const int height_in, const int width_in,
+                              const int height_out, const int width_out,
+                              const opmath_t offset_scale,
+                              const int remove_center, const int block_thread,
+                              const int softmax,
+                              const int padded_offset_dim) {
+
+  constexpr int L = 1;
+
+  auto kernel = forward_kernel_dcn_reg<scalar_t, d_stride, stride_type, 1, 9, true>;
+
+  int N = height_in * width_in;
+  int Q = height_out * width_out;
+  int K = kernel_h * kernel_w;
+
+  if (remove_center) {
+    K -= 1;
+  }
+  if (softmax) {
+    switch (K) {
+    case 9:
+      kernel = forward_kernel_dcn_reg<scalar_t, d_stride, stride_type, 1, 9, true>;
+      break;
+    case 8:
+      kernel = forward_kernel_dcn_reg<scalar_t, d_stride, stride_type, 1, 8, true>;
+      break;
+    default:
+      printf("K=%ld\n", K);
+      throw std::invalid_argument("invalid kernel shape");
+    }
+  } else {
+    switch (K) {
+    case 9:
+      kernel = forward_kernel_dcn_reg<scalar_t, d_stride, stride_type, 1, 9, false>;
+      break;
+    case 8:
+      kernel = forward_kernel_dcn_reg<scalar_t, d_stride, stride_type, 1, 8, false>;
+    break;
+    default:
+      printf("K=%ld\n", K);
+      throw std::invalid_argument("invalid kernel shape");
+    }
+  }
+
+  const int block_multiplier = block_thread / (D / d_stride) / G;
+  assert((B*Q) % block_multiplier == 0);
+
+  dim3 num_blocks(B*Q / block_multiplier);
+  dim3 num_threads(D / d_stride, G, block_multiplier);
+
+  int shm_size = 0;
+
+  cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
+                       shm_size);
+
+  kernel<<<num_blocks, num_threads, shm_size, stream>>>(
+      value, p_offset, output, G, D, Q, kernel_h, kernel_w, stride_h, stride_w,
+      pad_h, pad_w, dilation_h, dilation_w, height_in, width_in, height_out,
+      width_out, offset_scale, remove_center, block_multiplier, padded_offset_dim);
+
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess) {
+    printf("error in dcnv4_im2col_cuda: %s\n", cudaGetErrorString(err));
+    printf("launch arguments: gridDim=(%d, %d, %d), blockDim=(%d, %d, %d), "
+           "shm_size=%d\n\n",
+           num_blocks.x, num_blocks.y, num_blocks.z, num_threads.x,
+           num_threads.y, num_threads.z, shm_size);
+    AT_ASSERTM(false, "kernel launch error");
+  }
+}
+
+template <typename scalar_t>
+void dcnv4_im2col_cuda(
+    cudaStream_t stream,
+    const scalar_t *value,    // B, H * W, (G * D)
+    const scalar_t *p_offset, // B, H * W, G * K * 3)
+    scalar_t *output,         // B, H_out*W_out, G * D
+    const int kernel_h, const int kernel_w, const int stride_h,
+    const int stride_w, const int pad_h, const int pad_w, const int dilation_h,
+    const int dilation_w, const int G, const int D, const int B,
+    const int height_in, const int width_in, const int height_out,
+    const int width_out, const opmath_t offset_scale, const int remove_center,
+    const int d_stride, const int block_thread, const bool softmax,
+    const int padded_offset_dim) {
+
+  assert(D % d_stride == 0);
+  if (sizeof(scalar_t) == 2) {
+    switch (d_stride) {
+    case 1:
+      _dcnv4_im2col_cuda<scalar_t, scalar_t, 1>(
+          stream, value, p_offset, output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center,
+          block_thread, softmax, padded_offset_dim);
+      break;
+    case 2:
+      _dcnv4_im2col_cuda<scalar_t, uint, 2>(
+          stream, value, p_offset, output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center,
+          block_thread, softmax, padded_offset_dim);
+      break;
+    case 4:
+      _dcnv4_im2col_cuda<scalar_t, uint2, 4>(
+          stream, value, p_offset, output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center,
+          block_thread, softmax, padded_offset_dim);
+      break;
+    case 8:
+      _dcnv4_im2col_cuda<scalar_t, uint4, 8>(
+          stream, value, p_offset, output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center,
+          block_thread, softmax, padded_offset_dim);
+      break;
+    case 16:
+      _dcnv4_im2col_cuda<scalar_t, ulonglong4, 16>(
+          stream, value, p_offset, output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center,
+          block_thread, softmax, padded_offset_dim);
+      break;
+    }
+  } else {
+    assert(sizeof(scalar_t) == 4);
+    switch (d_stride) {
+    case 1:
+      _dcnv4_im2col_cuda<scalar_t, uint, 1>(
+          stream, value, p_offset, output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center,
+          block_thread, softmax, padded_offset_dim);
+      break;
+    case 2:
+      _dcnv4_im2col_cuda<scalar_t, uint2, 2>(
+          stream, value, p_offset, output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center,
+          block_thread, softmax, padded_offset_dim);
+      break;
+    case 4:
+      _dcnv4_im2col_cuda<scalar_t, uint4, 4>(
+          stream, value, p_offset, output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center,
+          block_thread, softmax, padded_offset_dim);
+      break;
+    case 8:
+      _dcnv4_im2col_cuda<scalar_t, ulonglong4, 8>(
+          stream, value, p_offset, output, kernel_h, kernel_w, stride_h,
+          stride_w, pad_h, pad_w, dilation_h, dilation_w, G, D, B, height_in,
+          width_in, height_out, width_out, offset_scale, remove_center,
+          block_thread, softmax, padded_offset_dim);
+      break;
+    default:
+      printf("not supported for d_stride > 8 for fp32");
+      throw std::invalid_argument("invalid d_stride");
+    }
+  }
+}

DCNv4_op/src/cuda/flash_deform_attn_cuda.cu

@@ -0,0 +1,162 @@
+/*!
+**************************************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************************************
+* Modified from
+*https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+#include "cuda/flash_deform_im2col_cuda.cuh"
+#include "cuda/flash_deform_col2im_cuda.cuh"
+#include <vector>
+
+#include <ATen/ATen.h>
+#include <ATen/OpMathType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cooperative_groups.h>
+#include <cooperative_groups/memcpy_async.h>
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <torch/torch.h>
+
+at::Tensor flash_deform_attn_cuda_forward(
+    const at::Tensor &value, const at::Tensor &spatial_shapes,
+    const at::Tensor &level_start_index, const at::Tensor &sampling_loc_attn,
+    const int im2col_step = 64, const int K=8, const int d_stride=8, 
+    const int block_thread=0) {
+  AT_ASSERTM(value.is_contiguous(), "value tensor has to be contiguous");
+  AT_ASSERTM(spatial_shapes.is_contiguous(),
+             "spatial_shapes tensor has to be contiguous");
+  AT_ASSERTM(level_start_index.is_contiguous(),
+             "level_start_index tensor has to be contiguous");
+  AT_ASSERTM(sampling_loc_attn.is_contiguous(),
+             "sampling_loc_attn tensor has to be contiguous");
+
+  AT_ASSERTM(value.type().is_cuda(), "value must be a CUDA tensor");
+  AT_ASSERTM(spatial_shapes.type().is_cuda(),
+             "spatial_shapes must be a CUDA tensor");
+  AT_ASSERTM(level_start_index.type().is_cuda(),
+             "level_start_index must be a CUDA tensor");
+  AT_ASSERTM(sampling_loc_attn.type().is_cuda(),
+             "sampling_loc_attn must be a CUDA tensor");
+
+  const int batch = value.size(0);
+  const int spatial_size = value.size(1);
+  const int num_heads = value.size(2);
+  const int num_channels = value.size(3);
+
+  const int num_levels = spatial_shapes.size(0);
+  const int num_query = sampling_loc_attn.size(1);
+  const int num_point = K;
+
+  const int im2col_step_ = std::min(batch, im2col_step);
+  AT_ASSERTM(batch % im2col_step_ == 0, "batch(", batch,
+             ") must divide im2col_step(", im2col_step_, ")");
+
+  auto output =
+      at::zeros({batch, num_query, num_heads, num_channels}, value.options());
+
+  auto per_value_size = spatial_size * num_heads * num_channels;
+  auto per_offset_size = num_query * num_heads * num_levels * num_point * 3;
+  auto per_out_size = num_query * num_heads * num_channels;
+
+  for (int n = 0; n < batch / im2col_step_; ++n) {
+    AT_DISPATCH_FLOATING_TYPES_AND2(
+        at::ScalarType::Half, at::ScalarType::BFloat16, value.scalar_type(),
+        "flash_deform_attn_forward_cuda", ([&] {
+          flash_deformable_im2col_cuda(
+              at::cuda::getCurrentCUDAStream(),
+              value.data_ptr<scalar_t>() + n * im2col_step_ * per_value_size,
+              spatial_shapes.data<int64_t>(), level_start_index.data<int64_t>(),
+              sampling_loc_attn.data_ptr<scalar_t>() +
+                  n * im2col_step_ * per_offset_size,
+              output.data_ptr<scalar_t>() + n * im2col_step_ * per_out_size,
+              im2col_step_, spatial_size, num_heads, num_channels, num_levels,
+              num_query, num_point, d_stride, block_thread, true);
+        }));
+  }
+  output = output.view({batch, num_query, num_heads * num_channels});
+  return output;
+}
+
+std::vector<at::Tensor>
+flash_deform_attn_cuda_backward(
+    const at::Tensor &value, const at::Tensor &spatial_shapes,
+    const at::Tensor &level_start_index, const at::Tensor &sampling_loc_attn,
+    const at::Tensor &grad_output, const int im2col_step = 64, const int K=8,
+    const int d_stride=2, const int block_thread=0) {
+  AT_ASSERTM(value.is_contiguous(), "value tensor has to be contiguous");
+  AT_ASSERTM(spatial_shapes.is_contiguous(),
+             "spatial_shapes tensor has to be contiguous");
+  AT_ASSERTM(level_start_index.is_contiguous(),
+             "level_start_index tensor has to be contiguous");
+  AT_ASSERTM(sampling_loc_attn.is_contiguous(),
+             "sampling_loc_attn tensor has to be contiguous");
+  AT_ASSERTM(grad_output.is_contiguous(),
+             "grad_output tensor has to be contiguous");
+
+  AT_ASSERTM(value.type().is_cuda(), "value must be a CUDA tensor");
+  AT_ASSERTM(spatial_shapes.type().is_cuda(),
+             "spatial_shapes must be a CUDA tensor");
+  AT_ASSERTM(level_start_index.type().is_cuda(),
+             "level_start_index must be a CUDA tensor");
+  AT_ASSERTM(sampling_loc_attn.type().is_cuda(),
+             "sampling_loc_attn must be a CUDA tensor");
+  AT_ASSERTM(grad_output.type().is_cuda(),
+             "grad_output must be a CUDA tensor");
+
+  const int batch = value.size(0);
+  const int spatial_size = value.size(1);
+  const int num_heads = value.size(2);
+  const int num_channels = value.size(3);
+
+  const int num_levels = spatial_shapes.size(0);
+  const int num_query = sampling_loc_attn.size(1);
+  const int num_point = K;
+
+  const int im2col_step_ = std::min(batch, im2col_step);
+  AT_ASSERTM(batch % im2col_step_ == 0, "batch(", batch,
+             ") must divide im2col_step(", im2col_step_, ")");
+
+  auto dtype = value.dtype();
+  if (dtype == at::kHalf){
+    dtype = at::kFloat;
+  }
+
+  auto grad_input = at::zeros_like(value, dtype);
+  auto grad_offset = at::zeros_like(sampling_loc_attn, dtype);
+
+  auto per_value_size = spatial_size * num_heads * num_channels;
+  auto per_offset_size = num_query * num_heads * num_levels * num_point * 3;
+  auto per_out_size = num_query * num_heads * num_channels;
+
+  for (int n = 0; n < batch / im2col_step_; ++n) {
+    AT_DISPATCH_FLOATING_TYPES_AND_HALF(value.type(),
+        "flash_deform_attn_backward_cuda", ([&] {
+          flash_deformable_col2im_cuda(
+              at::cuda::getCurrentCUDAStream(),
+              value.data_ptr<scalar_t>() + n * im2col_step_ * per_value_size,
+              spatial_shapes.data<int64_t>(), level_start_index.data<int64_t>(),
+              sampling_loc_attn.data_ptr<scalar_t>() +
+                  n * im2col_step_ * per_offset_size,
+              grad_output.data_ptr<scalar_t>() + n * im2col_step_ * per_out_size,
+              im2col_step_, spatial_size, num_heads, num_channels, num_levels,
+              num_query, num_point, 
+              grad_input.data<opmath_t>() + n * im2col_step_ * per_value_size,
+              grad_offset.data<opmath_t>() + n * im2col_step_ * per_offset_size,
+              d_stride, block_thread
+              );
+        }));
+  }
+
+  if(value.dtype() == torch::kHalf){
+    return {grad_input.to(torch::kHalf), grad_offset.to(torch::kHalf)};
+  }
+  else{
+    return {grad_input, grad_offset};
+  }
+}

DCNv4_op/src/cuda/flash_deform_attn_cuda.h

@@ -0,0 +1,25 @@
+/*!
+**************************************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************************************
+* Modified from
+*https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+#pragma once
+#include <torch/extension.h>
+
+at::Tensor flash_deform_attn_cuda_forward(
+    const at::Tensor &value, const at::Tensor &spatial_shapes,
+    const at::Tensor &level_start_index, const at::Tensor &sampling_loc_attn,
+    const int im2col_step, const int K, const int d_stride, const int block_thread);
+
+std::vector<at::Tensor>
+flash_deform_attn_cuda_backward(
+    const at::Tensor &value, const at::Tensor &spatial_shapes,
+    const at::Tensor &level_start_index, const at::Tensor &sampling_loc_attn,
+    const at::Tensor &grad_output, const int im2col_step, const int K,
+    const int d_stride, const int block_thread);

DCNv4_op/src/cuda/flash_deform_col2im_cuda.cuh

@@ -0,0 +1,580 @@
+#include <algorithm>
+#include <cstdio>
+#include <cstring>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <THC/THCAtomics.cuh>
+
+#include <ATen/ATen.h>
+#include <ATen/OpMathType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cooperative_groups.h>
+#include <cooperative_groups/memcpy_async.h>
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include "common.h"
+
+template <typename scalar_t, int d_stride, typename transfer_t, int L, int K>
+__global__ void
+backward_kernel(const scalar_t *p_value, const int64_t *data_spatial_shapes,
+                const int64_t *data_level_start_index, const scalar_t *p_offset,
+                const scalar_t *grad_output, const int N, const int G,
+                const int D, const int Q, 
+                const int block_multiplier, opmath_t *grad_im,
+                opmath_t *grad_offset) {
+
+  extern __shared__ char _s[];
+
+  const int &qi = (blockIdx.x * block_multiplier % Q) + threadIdx.z;
+  const int &bi = blockIdx.x * block_multiplier / Q;
+
+  const int &di_s = threadIdx.x * d_stride;
+  const int &gi = threadIdx.y;
+
+  opmath_t *cache_g_mask_before_softmax =
+      (opmath_t *)(_s); // (block_multiplier*G) * (L * K)
+  opmath_t *cache_grad_offset =
+      (opmath_t *)(cache_g_mask_before_softmax +
+                   block_multiplier * G * L *
+                       K); // (block_multiplier*G*D/d_stride*3)
+  opmath_t *const p_mask_shm =
+      ((opmath_t *)(cache_grad_offset +
+                    block_multiplier * G * D / d_stride * 3)) +
+      (threadIdx.z * G + gi) * L * K; // G*block_multiplier * L * K
+
+  const scalar_t *p_offset_ptr =
+      p_offset + (((bi * Q + qi) * G + gi) * L) * K * 3;
+  const int mask_length = L * K;
+  const int num_thread = (D / d_stride);
+  const int num_iter = mask_length / num_thread;
+  const int remainder = mask_length - num_iter * num_thread;
+  const scalar_t *top_grad = grad_output + ((bi * Q + qi) * G + gi) * D + di_s;
+
+  for (int i = 0; i < num_iter; i++) {
+    *(p_mask_shm + num_thread * i + threadIdx.x) =
+        *(scalar_t *)(p_offset_ptr + L * K * 2 + num_thread * i + threadIdx.x);
+  }
+  if (remainder > 0 && threadIdx.x < remainder) {
+    *(p_mask_shm + num_thread * num_iter + threadIdx.x) =
+        *(scalar_t *)(p_offset_ptr + L * K * 2 + num_thread * num_iter +
+                      threadIdx.x);
+  }
+
+  __syncthreads();
+  // Calculate softmax over L and K
+  if (threadIdx.x == 0) { // gi != 0, di = 0, li = 0
+    opmath_t softmax_max = -1e100;
+    opmath_t softmax_sum = 0.0;
+
+    // get max
+    for (int j = 0; j < L * K; j++) {
+      softmax_max = max(softmax_max, p_mask_shm[j]);
+    }
+
+    // get sumexp
+    for (int j = 0; j < L * K; j++) {
+      opmath_t exp_results = exp(p_mask_shm[j] - softmax_max);
+      p_mask_shm[j] = exp_results;
+      softmax_sum += exp_results;
+    }
+
+    // normalize
+    for (int j = 0; j < L * K; j++) {
+      p_mask_shm[j] /= softmax_sum;
+    }
+  }
+
+  __syncthreads();
+
+  int offset_idx = 0;
+  int mask_idx = 0;
+  const int w_stride = G * D;
+  const int base_ptr = gi * D + di_s;
+
+  for (int li = 0; li < L; li++) {
+
+    const int spatial_h = data_spatial_shapes[li * 2];
+    const int spatial_w = data_spatial_shapes[li * 2 + 1];
+    const int level_start_id = data_level_start_index[li];
+    const scalar_t *p_value_ptr = p_value + (bi * N + level_start_id) * G * D;
+    opmath_t *grad_im_ptr = grad_im + (bi * N + level_start_id) * G * D;
+
+    int cache_grad_off_idx =
+        ((threadIdx.z * G + threadIdx.y) * blockDim.x + threadIdx.x) * 3;
+    for (int ki = 0; ki < K; ki++) {
+      const opmath_t loc_w = p_offset_ptr[offset_idx];
+      const opmath_t loc_h = p_offset_ptr[offset_idx + 1];
+      const opmath_t attn = p_mask_shm[mask_idx];
+      const opmath_t h_im = loc_h * spatial_h - 0.5;
+      const opmath_t w_im = loc_w * spatial_w - 0.5;
+      // for cache_grad_offset (mG) x D/d x 3
+      cache_grad_offset[cache_grad_off_idx] = 0;
+      cache_grad_offset[cache_grad_off_idx + 1] = 0;
+      cache_grad_offset[cache_grad_off_idx + 2] = 0;
+
+      if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
+        ms_deform_attn_col2im_bilinear<scalar_t, transfer_t, d_stride>(
+            p_value_ptr, spatial_h, spatial_w, h_im, w_im, attn, w_stride,
+            base_ptr, spatial_h, spatial_w, top_grad, grad_im_ptr,
+            cache_grad_offset + cache_grad_off_idx);
+
+        // aggregate across different channel for offset
+        __syncthreads();
+        if (threadIdx.x == 0) {
+          int _didx = (threadIdx.z * G + threadIdx.y) * blockDim.x * 3;
+          opmath_t _grad_w = cache_grad_offset[_didx];
+          opmath_t _grad_h = cache_grad_offset[_didx + 1];
+          opmath_t _grad_a = cache_grad_offset[_didx + 2];
+          for (int c_id = 1; c_id < blockDim.x; ++c_id) {
+            _grad_w += cache_grad_offset[_didx + 3 * c_id];
+            _grad_h += cache_grad_offset[_didx + 3 * c_id + 1];
+            _grad_a += cache_grad_offset[_didx + 3 * c_id + 2];
+          }
+
+          grad_offset[((bi * Q + qi) * G + gi) * L * K * 3 + li * K * 2 +
+                      ki * 2] = _grad_w;
+          grad_offset[((bi * Q + qi) * G + gi) * L * K * 3 + li * K * 2 +
+                      ki * 2 + 1] = _grad_h;
+          cache_g_mask_before_softmax
+              [((threadIdx.y + threadIdx.z * G) * L + li) * K + ki] = _grad_a;
+        }
+      }
+      __syncthreads();
+
+      offset_idx += 2;
+      mask_idx += 1;
+    }
+  }
+  // backward for softmax
+  if (threadIdx.x == 0) {
+    for (int i = 0; i < L * K; ++i) {
+      opmath_t grad_i = 0.;
+      const opmath_t *group_g_mask = cache_g_mask_before_softmax +
+                                      (threadIdx.y + threadIdx.z * G) * L * K;
+      for (int j = 0; j < L * K; ++j) {
+        if (i != j) {
+          grad_i -= group_g_mask[j] * p_mask_shm[i] * p_mask_shm[j];
+        } else {
+          grad_i += group_g_mask[i] * p_mask_shm[i] * (1 - p_mask_shm[i]);
+        }
+      }
+      grad_offset[((bi * Q + qi) * G + gi) * L * K * 3 + L * K * 2 + i] =
+          grad_i;
+    }
+  }
+  __syncthreads();
+}
+
+template <typename scalar_t, int d_stride, typename transfer_t, int L, int K>
+__global__ void
+backward_kernel_warp_primitive(const scalar_t *p_value, const int64_t *data_spatial_shapes,
+                const int64_t *data_level_start_index, const scalar_t *p_offset,
+                const scalar_t *grad_output, const int N, const int G,
+                const int D, const int Q, 
+                const int block_multiplier, opmath_t *grad_im,
+                opmath_t *grad_offset) {
+
+  extern __shared__ char _s[];
+
+  const int &qi = (blockIdx.x * block_multiplier % Q) + threadIdx.z;
+  const int &bi = blockIdx.x * block_multiplier / Q;
+
+  const int &di_s = threadIdx.x * d_stride;
+  const int &gi = threadIdx.y;
+
+  const int tid = (threadIdx.z * blockDim.y + threadIdx.y)*blockDim.x + threadIdx.x;
+  const int lane_id = tid % kWarpSize;
+  const int group_per_warp = kWarpSize / blockDim.x;
+  const int group_in_warp_id = (threadIdx.z * G + threadIdx.y) % group_per_warp;
+  const unsigned lane_mask = ((1 << blockDim.x) - 1) << (group_in_warp_id * blockDim.x);
+
+  opmath_t *cache_g_mask_before_softmax =
+      (opmath_t *)(_s); // (block_multiplier*G) * (L * K)
+
+  opmath_t *const p_mask_shm =
+      ((opmath_t *)(cache_g_mask_before_softmax + block_multiplier * G * L * K)) +
+        (threadIdx.z * G + gi) * L * K; // G*block_multiplier * L * K
+
+  const scalar_t *p_offset_ptr =
+      p_offset + (((bi * Q + qi) * G + gi) * L) * K * 3;
+  const int mask_length = L * K;
+  const int num_thread = (D / d_stride);
+  const int num_iter = mask_length / num_thread;
+  const int remainder = mask_length - num_iter * num_thread;
+  const scalar_t *top_grad = grad_output + ((bi * Q + qi) * G + gi) * D + di_s;
+
+  for (int i = 0; i < num_iter; i++) {
+    *(p_mask_shm + num_thread * i + threadIdx.x) =
+        *(scalar_t *)(p_offset_ptr + L * K * 2 + num_thread * i + threadIdx.x);
+  }
+  if (remainder > 0 && threadIdx.x < remainder) {
+    *(p_mask_shm + num_thread * num_iter + threadIdx.x) =
+        *(scalar_t *)(p_offset_ptr + L * K * 2 + num_thread * num_iter +
+                      threadIdx.x);
+  }
+
+  __syncthreads();
+  // Calculate softmax over L and K
+  if (threadIdx.x == 0) { // gi != 0, di = 0, li = 0
+    opmath_t softmax_max = -1e100;
+    opmath_t softmax_sum = 0.0;
+
+    // get max
+    for (int j = 0; j < L * K; j++) {
+      softmax_max = max(softmax_max, p_mask_shm[j]);
+    }
+
+    // get sumexp
+    for (int j = 0; j < L * K; j++) {
+      opmath_t exp_results = exp(p_mask_shm[j] - softmax_max);
+      p_mask_shm[j] = exp_results;
+      softmax_sum += exp_results;
+    }
+
+    // normalize
+    for (int j = 0; j < L * K; j++) {
+      p_mask_shm[j] /= softmax_sum;
+    }
+  }
+
+  __syncthreads();
+
+  int offset_idx = 0;
+  int mask_idx = 0;
+  const int w_stride = G * D;
+  const int base_ptr = gi * D + di_s;
+
+  for (int li = 0; li < L; li++) {
+
+    const int spatial_h = data_spatial_shapes[li * 2];
+    const int spatial_w = data_spatial_shapes[li * 2 + 1];
+    const int level_start_id = data_level_start_index[li];
+    const scalar_t *p_value_ptr = p_value + (bi * N + level_start_id) * G * D;
+    opmath_t *grad_im_ptr = grad_im + (bi * N + level_start_id) * G * D;
+
+    int cache_grad_off_idx =
+        ((threadIdx.z * G + threadIdx.y) * blockDim.x + threadIdx.x) * 3;
+
+    opmath_t reg_grad_offset[3] = {0.};
+    for (int ki = 0; ki < K; ki++) {
+      const opmath_t loc_w = p_offset_ptr[offset_idx];
+      const opmath_t loc_h = p_offset_ptr[offset_idx + 1];
+      const opmath_t attn = p_mask_shm[mask_idx];
+      const opmath_t h_im = loc_h * spatial_h - 0.5;
+      const opmath_t w_im = loc_w * spatial_w - 0.5;
+      reg_grad_offset[0] = 0;
+      reg_grad_offset[1] = 0;
+      reg_grad_offset[2] = 0;
+
+      if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
+        ms_deform_attn_col2im_bilinear<scalar_t, transfer_t, d_stride>(
+            p_value_ptr, spatial_h, spatial_w, h_im, w_im, attn, w_stride,
+            base_ptr, spatial_h, spatial_w, top_grad, grad_im_ptr,
+            reg_grad_offset);
+
+        // aggregate across different channel for offset
+        for (uint32_t offset = blockDim.x>>1; offset > 0; offset >>= 1){
+          reg_grad_offset[0] += __shfl_down_sync(lane_mask, reg_grad_offset[0], offset);
+          reg_grad_offset[1] += __shfl_down_sync(lane_mask, reg_grad_offset[1], offset);
+          reg_grad_offset[2] += __shfl_down_sync(lane_mask, reg_grad_offset[2], offset);
+        }
+
+        if (threadIdx.x == 0) {
+          grad_offset[((bi * Q + qi) * G + gi) * L * K * 3 + li * K * 2 +
+                      ki * 2] = reg_grad_offset[0];
+          grad_offset[((bi * Q + qi) * G + gi) * L * K * 3 + li * K * 2 +
+                      ki * 2 + 1] = reg_grad_offset[1];
+          cache_g_mask_before_softmax
+              [((threadIdx.y + threadIdx.z * G) * L + li) * K + ki] = reg_grad_offset[2];
+        }
+      }
+      __syncthreads();
+
+      offset_idx += 2;
+      mask_idx += 1;
+    }
+  }
+  // backward for softmax
+  if (threadIdx.x == 0) {
+    for (int i = 0; i < L * K; ++i) {
+      opmath_t grad_i = 0.;
+      const opmath_t *group_g_mask = cache_g_mask_before_softmax +
+                                      (threadIdx.y + threadIdx.z * G) * L * K;
+      for (int j = 0; j < L * K; ++j) {
+        if (i != j) {
+          grad_i -= group_g_mask[j] * p_mask_shm[i] * p_mask_shm[j];
+        } else {
+          grad_i += group_g_mask[i] * p_mask_shm[i] * (1 - p_mask_shm[i]);
+        }
+      }
+      grad_offset[((bi * Q + qi) * G + gi) * L * K * 3 + L * K * 2 + i] =
+          grad_i;
+    }
+  }
+  __syncthreads();
+}
+
+template <typename scalar_t, typename stride_type, int K, int d_stride>
+void _flash_deformable_col2im_cuda(
+    cudaStream_t stream,
+    const scalar_t *value,                 // B, N, G, D
+    const int64_t *data_spatial_shapes,    // L * 2
+    const int64_t *data_level_start_index, // L
+    const scalar_t *offset,                // B, N, G, L, K, 3
+    const scalar_t *grad_output,           // B, N, G, D
+    const int B, const int N, const int G, const int D, const int L,
+    const int Q, opmath_t *grad_im, opmath_t *grad_offset,
+    const int block_thread) {
+
+  assert(D % d_stride == 0);
+
+  const int block_multiplier = block_thread / (D / d_stride) / G;
+  assert((B*Q) % block_multiplier == 0);
+  dim3 num_blocks(B*Q / block_multiplier);
+  dim3 num_threads(D / d_stride, G, block_multiplier);
+
+  int shm_size;
+  if(check_backward_warpp(d_stride, D)){
+    shm_size =
+      sizeof(opmath_t) * (block_multiplier * G * L * K) +
+      sizeof(opmath_t) * (G * block_multiplier * L * K);
+  }
+  else{
+    shm_size =
+      sizeof(opmath_t) * (block_multiplier * G * L * K) +
+      sizeof(opmath_t) * (G * block_multiplier * L * K) + 
+      sizeof(opmath_t) * (G * block_multiplier * D / d_stride * 3);
+  }
+
+  auto kernel = backward_kernel_warp_primitive<scalar_t, d_stride, stride_type, 1, K>;
+
+  switch (L) {
+  case 1:
+    if(check_backward_warpp(d_stride, D)){
+      kernel = backward_kernel_warp_primitive<scalar_t, d_stride, stride_type, 1, K>;
+    } else {
+      kernel = backward_kernel<scalar_t, d_stride, stride_type, 1, K>;
+    }
+    break;
+  case 2:
+    if(check_backward_warpp(d_stride, D)){
+      kernel = backward_kernel_warp_primitive<scalar_t, d_stride, stride_type, 2, K>;
+    } else {
+      kernel = backward_kernel<scalar_t, d_stride, stride_type, 2, K>;
+    }
+    break;
+  case 3:
+    if(check_backward_warpp(d_stride, D)){
+      kernel = backward_kernel_warp_primitive<scalar_t, d_stride, stride_type, 3, K>;
+    } else {
+      kernel = backward_kernel<scalar_t, d_stride, stride_type, 3, K>;
+    }
+    break;
+  case 4:
+    if(check_backward_warpp(d_stride, D)){
+      kernel = backward_kernel_warp_primitive<scalar_t, d_stride, stride_type, 4, K>;
+    } else {
+      kernel = backward_kernel<scalar_t, d_stride, stride_type, 4, K>;
+    }
+    break;
+  case 5:
+    if(check_backward_warpp(d_stride, D)){
+      kernel = backward_kernel_warp_primitive<scalar_t, d_stride, stride_type, 5, K>;
+    } else {
+      kernel = backward_kernel<scalar_t, d_stride, stride_type, 5, K>;
+    }
+    break;
+  default:
+    printf("L=%ld\n", L);
+    throw std::invalid_argument("invalid number of scales");
+  }
+  cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
+                       shm_size);
+
+  kernel<<<num_blocks, num_threads, shm_size, stream>>>(
+      value, data_spatial_shapes, data_level_start_index, offset, grad_output,
+      N, G, D, Q, block_multiplier, grad_im, grad_offset);
+
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess) {
+    printf("error in flash_deformable_im2col_cuda: %s\n",
+           cudaGetErrorString(err));
+    printf("launch arguments: gridDim=(%d, %d, %d), blockDim=(%d, %d, %d), "
+           "shm_size=%d, Q=%d\n\n",
+           num_blocks.x, num_blocks.y, num_blocks.z, num_threads.x,
+           num_threads.y, num_threads.z, shm_size, Q);
+    AT_ASSERTM(false, "kernel launch error");
+  }
+}
+
+template <typename scalar_t, int K>
+void flash_deformable_col2im_cuda_inner(
+    cudaStream_t stream,
+    const scalar_t *value,                 // B, N, G, D
+    const int64_t *data_spatial_shapes,    // L * 2
+    const int64_t *data_level_start_index, // L
+    const scalar_t *offset,                // B, N, G, L, K, 3
+    const scalar_t *grad_output,           // B, N, G, D
+    const int B, const int N, const int G, const int D, const int L,
+    const int Q, opmath_t *grad_im, opmath_t *grad_offset, 
+    const int d_stride, const int block_thread) {
+
+  assert(D % d_stride == 0);
+  if(sizeof(scalar_t) == 2) {
+    switch(d_stride) {
+    case 1:
+      _flash_deformable_col2im_cuda<scalar_t, scalar_t, K, 1>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          grad_output,            // B, N, G, D
+          B, N, G, D, L, Q, grad_im, grad_offset,
+          block_thread);
+      break;
+    case 2:
+      _flash_deformable_col2im_cuda<scalar_t, uint, K, 2>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          grad_output,            // B, N, G, D
+          B, N, G, D, L, Q, grad_im, grad_offset,
+          block_thread);
+      break;
+    case 4:
+      _flash_deformable_col2im_cuda<scalar_t, uint2, K, 4>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          grad_output,            // B, N, G, D
+          B, N, G, D, L, Q, grad_im, grad_offset,
+          block_thread);
+      break;
+    case 8:
+      _flash_deformable_col2im_cuda<scalar_t, uint4, K, 8>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          grad_output,            // B, N, G, D
+          B, N, G, D, L, Q, grad_im, grad_offset,
+          block_thread);
+      break;
+    case 16:
+      _flash_deformable_col2im_cuda<scalar_t, ulonglong4, K, 16>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          grad_output,            // B, N, G, D
+          B, N, G, D, L, Q, grad_im, grad_offset,
+          block_thread);
+      break;
+    default:
+      printf("not supported for d_stride > 16 for fp16");
+      throw std::invalid_argument("invalid d_stride");
+    }
+  } else {
+    assert(sizeof(scalar_t) == 4);
+    switch(d_stride) {
+    case 1:  
+      _flash_deformable_col2im_cuda<scalar_t, scalar_t, K, 1>(
+        stream,
+        value,                  // B, N, G, D
+        data_spatial_shapes,    // L * 2
+        data_level_start_index, // L
+        offset,                 // B, N, G, L, K, 3
+        grad_output,            // B, N, G, D
+        B, N, G, D, L, Q, grad_im, grad_offset,
+        block_thread);
+      break;
+    case 2:  
+      _flash_deformable_col2im_cuda<scalar_t, uint2, K, 2>(
+        stream,
+        value,                  // B, N, G, D
+        data_spatial_shapes,    // L * 2
+        data_level_start_index, // L
+        offset,                 // B, N, G, L, K, 3
+        grad_output,            // B, N, G, D
+        B, N, G, D, L, Q, grad_im, grad_offset,
+        block_thread);
+      break;
+    case 4:  
+      _flash_deformable_col2im_cuda<scalar_t, uint4, K, 4>(
+        stream,
+        value,                  // B, N, G, D
+        data_spatial_shapes,    // L * 2
+        data_level_start_index, // L
+        offset,                 // B, N, G, L, K, 3
+        grad_output,            // B, N, G, D
+        B, N, G, D, L, Q, grad_im, grad_offset,
+        block_thread);
+      break;
+    case 8:  
+      _flash_deformable_col2im_cuda<scalar_t, ulonglong4, K, 8>(
+        stream,
+        value,                  // B, N, G, D
+        data_spatial_shapes,    // L * 2
+        data_level_start_index, // L
+        offset,                 // B, N, G, L, K, 3
+        grad_output,            // B, N, G, D
+        B, N, G, D, L, Q, grad_im, grad_offset,
+        block_thread);
+      break;
+    default:
+      printf("not supported for d_stride > 8 for fp32");
+      throw std::invalid_argument("invalid d_stride");
+    }
+  }
+}
+
+template <typename scalar_t>
+void flash_deformable_col2im_cuda(
+    cudaStream_t stream,
+    const scalar_t *value,                 // B, N, G, D
+    const int64_t *data_spatial_shapes,    // L * 2
+    const int64_t *data_level_start_index, // L
+    const scalar_t *offset,                // B, N, G, L, K, 3
+    const scalar_t *grad_output,           // B, N, G, D
+    const int B, const int N, const int G, const int D, const int L,
+    const int Q, const int K, opmath_t *grad_im, opmath_t *grad_offset,
+    const int d_stride, const int block_thread) {
+
+  switch (K) {
+  case 4:
+    flash_deformable_col2im_cuda_inner<scalar_t, 4>(
+        stream,
+        value,                  // B, N, G, D
+        data_spatial_shapes,    // L * 2
+        data_level_start_index, // L
+        offset,                 // B, N, G, L, K, 3
+        grad_output,            // B, N, G, D
+        B, N, G, D, L, Q, grad_im, grad_offset,
+        d_stride, block_thread);
+    break;
+  case 8:
+    flash_deformable_col2im_cuda_inner<scalar_t, 8>(
+        stream,
+        value,                  // B, N, G, D
+        data_spatial_shapes,    // L * 2
+        data_level_start_index, // L
+        offset,                 // B, N, G, L, K, 3
+        grad_output,            // B, N, G, D
+        B, N, G, D, L, Q, grad_im, grad_offset,
+        d_stride, block_thread);
+    break;
+  default:
+    printf("not supported for K not in [4, 8]");
+    throw std::invalid_argument("invalid K");
+  }
+}

DCNv4_op/src/cuda/flash_deform_im2col_cuda.cuh

@@ -0,0 +1,451 @@
+/*!
+**************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************
+* Modified from DCN (https://github.com/msracver/Deformable-ConvNets)
+* Copyright (c) 2018 Microsoft
+**************************************************************************
+*/
+#include <algorithm>
+#include <cstdio>
+#include <cstring>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <THC/THCAtomics.cuh>
+
+#include <ATen/ATen.h>
+#include <ATen/OpMathType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cooperative_groups.h>
+#include <cooperative_groups/memcpy_async.h>
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include "common.h"
+
+template <typename scalar_t, int d_stride, typename transfer_t, int L, int K>
+__global__ void
+forward_kernel(const scalar_t *p_value, const int64_t *data_spatial_shapes,
+               const int64_t *data_level_start_index, const scalar_t *p_offset,
+               scalar_t *p_output, const int N, const int G, const int D,
+               const int Q, const int block_multiplier) {
+
+  extern __shared__ char _s[];
+
+  const int &qi = (blockIdx.x * block_multiplier % Q) + threadIdx.z;
+  const int &bi = blockIdx.x * block_multiplier / Q;
+
+  const int &di_s = threadIdx.x * d_stride;
+  const int &gi = threadIdx.y;
+
+  opmath_t p_out_shm[d_stride] = {0.};
+  opmath_t *const p_mask_shm =
+      (opmath_t *)(_s) + (threadIdx.z * G + gi) * L * K;
+
+  const scalar_t *p_offset_ptr =
+      p_offset + (((bi * Q + qi) * G + gi) * L) * K * 3;
+  const int mask_length = L * K;
+  const int num_thread = (D / d_stride);
+  const int num_iter = mask_length / num_thread;
+  const int remainder = mask_length - num_iter * num_thread;
+  for (int i = 0; i < num_iter; i++) {
+    *(p_mask_shm + num_thread * i + threadIdx.x) =
+        *(scalar_t *)(p_offset_ptr + L * K * 2 + num_thread * i + threadIdx.x);
+  }
+  if (remainder > 0 && threadIdx.x < remainder) {
+    *(p_mask_shm + num_thread * num_iter + threadIdx.x) =
+        *(scalar_t *)(p_offset_ptr + L * K * 2 + num_thread * num_iter +
+                      threadIdx.x);
+  }
+
+  __syncthreads();
+  // Calculate softmax over L and K
+  if (threadIdx.x == 0) { // di = 0
+    opmath_t softmax_max = -1e100;
+    opmath_t softmax_sum = 0.0;
+
+    // get max
+    for (int j = 0; j < L * K; j++) {
+      softmax_max = max(softmax_max, p_mask_shm[j]);
+    }
+
+    // get sumexp
+    for (int j = 0; j < L * K; j++) {
+      opmath_t exp_results = exp(p_mask_shm[j] - softmax_max);
+      p_mask_shm[j] = exp_results;
+      softmax_sum += exp_results;
+    }
+
+    // normalize
+    for (int j = 0; j < L * K; j++) {
+      p_mask_shm[j] /= softmax_sum;
+    }
+  }
+
+  __syncthreads();
+
+  int offset_idx = 0;
+  int mask_idx = 0;
+  const int w_stride = G * D;
+  const int base_ptr = gi * D + di_s;
+
+  for (int li = 0; li < L; li++) {
+    const int spatial_h = data_spatial_shapes[li * 2];
+    const int spatial_w = data_spatial_shapes[li * 2 + 1];
+    const int level_start_id = data_level_start_index[li];
+    const scalar_t *p_value_ptr = p_value + (bi * N + level_start_id) * G * D;
+
+    for (int ki = 0; ki < K; ki++) {
+      const opmath_t loc_w = p_offset_ptr[offset_idx];
+      const opmath_t loc_h = p_offset_ptr[offset_idx + 1];
+      const opmath_t attn = p_mask_shm[mask_idx];
+      const opmath_t h_im = loc_h * spatial_h - 0.5;
+      const opmath_t w_im = loc_w * spatial_w - 0.5;
+      if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
+        ms_deform_attn_im2col_bilinear<scalar_t, transfer_t, d_stride>(
+            p_out_shm, p_value_ptr, spatial_h, spatial_w, h_im, w_im, attn,
+            w_stride, base_ptr);
+      }
+      offset_idx += 2;
+      mask_idx += 1;
+    }
+  }
+
+  int out_idx = ((bi * Q + qi) * G + gi) * D + di_s;
+
+  scalar_t *fp16_regs = (scalar_t *)(p_out_shm);
+#pragma unroll
+  for (int ds = 0; ds < d_stride; ds++) {
+    fp16_regs[ds] = p_out_shm[ds];
+  }
+
+  *(transfer_t *)(p_output + out_idx) = *(transfer_t *)(p_out_shm);
+}
+
+template <typename scalar_t, int d_stride, typename transfer_t, int L, int K>
+__global__ void
+forward_kernel_reg(const scalar_t *p_value, const int64_t *data_spatial_shapes,
+               const int64_t *data_level_start_index, const scalar_t *p_offset,
+               scalar_t *p_output, const int N, const int G, const int D,
+               const int Q, const int block_multiplier) {
+
+  const int &qi = (blockIdx.x * block_multiplier % Q) + threadIdx.z;
+  const int &bi = blockIdx.x * block_multiplier / Q;
+
+  const int &di_s = threadIdx.x * d_stride;
+  const int &gi = threadIdx.y;
+
+  opmath_t p_out_shm[d_stride] = {0.};
+  opmath_t p_mask_shm[L*K] = {0.};
+
+  const scalar_t *p_offset_ptr =
+      p_offset + (((bi * Q + qi) * G + gi) * L) * K * 3;
+
+  for (int i=0; i < L*K; i++){
+    p_mask_shm[i] = *(p_offset_ptr + L * K * 2 + i);
+  }
+
+  // Calculate softmax over L and K
+  opmath_t softmax_max = -1e100;
+  opmath_t softmax_sum = 0.0;
+
+  // get max
+  for (int j = 0; j < L * K; j++) {
+    softmax_max = max(softmax_max, p_mask_shm[j]);
+  }
+
+  // get sumexp
+  for (int j = 0; j < L * K; j++) {
+    opmath_t exp_results = exp(p_mask_shm[j] - softmax_max);
+    p_mask_shm[j] = exp_results;
+    softmax_sum += exp_results;
+  }
+
+  // normalize
+  for (int j = 0; j < L * K; j++) {
+    p_mask_shm[j] /= softmax_sum;
+  }
+
+  int offset_idx = 0;
+  int mask_idx = 0;
+  const int w_stride = G * D;
+  const int base_ptr = gi * D + di_s;
+
+  for (int li = 0; li < L; li++) {
+    const int spatial_h = data_spatial_shapes[li * 2];
+    const int spatial_w = data_spatial_shapes[li * 2 + 1];
+    const int level_start_id = data_level_start_index[li];
+    const scalar_t *p_value_ptr = p_value + (bi * N + level_start_id) * G * D;
+
+    for (int ki = 0; ki < K; ki++) {
+      const opmath_t loc_w = p_offset_ptr[offset_idx];
+      const opmath_t loc_h = p_offset_ptr[offset_idx + 1];
+      const opmath_t attn = p_mask_shm[mask_idx];
+      const opmath_t h_im = loc_h * spatial_h - 0.5;
+      const opmath_t w_im = loc_w * spatial_w - 0.5;
+      if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
+        ms_deform_attn_im2col_bilinear<scalar_t, transfer_t, d_stride>(
+            p_out_shm, p_value_ptr, spatial_h, spatial_w, h_im, w_im, attn,
+            w_stride, base_ptr);
+      }
+      offset_idx += 2;
+      mask_idx += 1;
+    }
+  }
+
+  int out_idx = ((bi * Q + qi) * G + gi) * D + di_s;
+
+  scalar_t *fp16_regs = (scalar_t *)(p_out_shm);
+#pragma unroll
+  for (int ds = 0; ds < d_stride; ds++) {
+    fp16_regs[ds] = p_out_shm[ds];
+  }
+
+  *(transfer_t *)(p_output + out_idx) = *(transfer_t *)(p_out_shm);
+}
+
+template <typename scalar_t, typename stride_type, int K, int d_stride>
+void _flash_deformable_im2col_cuda(
+    cudaStream_t stream,
+    const scalar_t *value,                 // B, N, G, D
+    const int64_t *data_spatial_shapes,    // L * 2
+    const int64_t *data_level_start_index, // L
+    const scalar_t *offset,                // B, N, G, L, K, 3
+    scalar_t *output,                      // B, N, G, D
+    const int B, const int N, const int G, const int D, const int L,
+    const int Q, const int block_thread, 
+    const bool _use_reg) {
+
+  assert(D % d_stride == 0);
+
+  const int block_multiplier = block_thread / (D / d_stride) / G;;
+  assert((B*Q) % block_multiplier == 0);
+  dim3 num_blocks(B*Q / block_multiplier);
+  dim3 num_threads(D / d_stride, G, block_multiplier);
+
+  const int shm_size = 0;
+
+  auto kernel = forward_kernel_reg<scalar_t, d_stride, stride_type, 1, K>;
+
+  switch (L) {
+  case 1:
+    kernel = forward_kernel_reg<scalar_t, d_stride, stride_type, 1, K>;
+    break;
+  case 2:
+    kernel = forward_kernel_reg<scalar_t, d_stride, stride_type, 2, K>;
+    break;
+  case 3:
+    kernel = forward_kernel_reg<scalar_t, d_stride, stride_type, 3, K>;
+    break;
+  case 4:
+    kernel = forward_kernel_reg<scalar_t, d_stride, stride_type, 4, K>;
+    break;
+  case 5:
+    kernel = forward_kernel_reg<scalar_t, d_stride, stride_type, 5, K>;
+    break;
+  default:
+    printf("L=%ld\n", L);
+    throw std::invalid_argument("invalid number of scales");
+  }
+
+  cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
+                       shm_size);
+
+  kernel<<<num_blocks, num_threads, shm_size, stream>>>(
+      value, data_spatial_shapes, data_level_start_index, offset, output, N, G,
+      D, Q, block_multiplier);
+
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess) {
+    printf("error in flash_deformable_im2col_cuda: %s\n",
+           cudaGetErrorString(err));
+    printf("launch arguments: gridDim=(%d, %d, %d), blockDim=(%d, %d, %d), "
+           "shm_size=%d, Q=%d\n\n",
+           num_blocks.x, num_blocks.y, num_blocks.z, num_threads.x,
+           num_threads.y, num_threads.z, shm_size, Q);
+    AT_ASSERTM(false, "kernel launch error");
+  }
+}
+
+template <typename scalar_t, int K>
+void flash_deformable_im2col_cuda_inner(
+    cudaStream_t stream,
+    const scalar_t *value,                 // B, N, G, D
+    const int64_t *data_spatial_shapes,    // L * 2
+    const int64_t *data_level_start_index, // L
+    const scalar_t *offset,                // B, N, G, L, K, 3
+    scalar_t *output,                      // B, N, G, D
+    const int B, const int N, const int G, const int D, const int L,
+    const int Q, const int d_stride, 
+    const int block_thread,
+    const bool _use_reg) {
+  
+  assert(D % d_stride == 0);
+  if(sizeof(scalar_t) == 2) {
+    switch(d_stride) {
+    case 1:
+      _flash_deformable_im2col_cuda<scalar_t, scalar_t, K, 1>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          output,                 // B, N, G, D
+          B, N, G, D, L, Q,
+          block_thread,
+          _use_reg);
+      break;
+    case 2:
+      _flash_deformable_im2col_cuda<scalar_t, uint, K, 2>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          output,                 // B, N, G, D
+          B, N, G, D, L, Q,
+          block_thread,
+          _use_reg);
+      break;
+    case 4:
+      _flash_deformable_im2col_cuda<scalar_t, uint2, K, 4>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          output,                 // B, N, G, D
+          B, N, G, D, L, Q,
+          block_thread,
+          _use_reg);
+      break;
+    case 8:
+      _flash_deformable_im2col_cuda<scalar_t, uint4, K, 8>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          output,                 // B, N, G, D
+          B, N, G, D, L, Q,
+          block_thread,
+          _use_reg);
+      break;
+    case 16:
+      _flash_deformable_im2col_cuda<scalar_t, ulonglong4, K, 16>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          output,                 // B, N, G, D
+          B, N, G, D, L, Q,
+          block_thread,
+          _use_reg);
+      break;
+    default:
+      printf("not supported for d_stride > 16 for fp16");
+      throw std::invalid_argument("invalid d_stride");
+    }
+  } else {
+    assert(sizeof(scalar_t) == 4);
+    switch(d_stride) {
+    case 1:
+      _flash_deformable_im2col_cuda<scalar_t, scalar_t, K, 1>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          output,                 // B, N, G, D
+          B, N, G, D, L, Q,
+          block_thread,
+          _use_reg);
+      break;
+    case 2:
+      _flash_deformable_im2col_cuda<scalar_t, uint2, K, 2>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          output,                 // B, N, G, D
+          B, N, G, D, L, Q,
+          block_thread,
+          _use_reg);
+      break;
+    case 4:
+      _flash_deformable_im2col_cuda<scalar_t, uint4, K, 4>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          output,                 // B, N, G, D
+          B, N, G, D, L, Q,
+          block_thread,
+          _use_reg);
+      break;
+    case 8:
+      _flash_deformable_im2col_cuda<scalar_t, ulonglong4, K, 8>(
+          stream,
+          value,                  // B, N, G, D
+          data_spatial_shapes,    // L * 2
+          data_level_start_index, // L
+          offset,                 // B, N, G, L, K, 3
+          output,                 // B, N, G, D
+          B, N, G, D, L, Q,
+          block_thread,
+          _use_reg);
+      break;
+    default:
+      printf("not supported for d_stride > 8 for fp32");
+      throw std::invalid_argument("invalid d_stride");
+    }
+  }
+}
+
+template <typename scalar_t>
+void flash_deformable_im2col_cuda(
+    cudaStream_t stream,
+    const scalar_t *value,                 // B, N, G, D
+    const int64_t *data_spatial_shapes,    // L * 2
+    const int64_t *data_level_start_index, // L
+    const scalar_t *offset,                // B, N, G, L, K, 3
+    scalar_t *output,                      // B, N, G, D
+    const int B, const int N, const int G, const int D, const int L,
+    const int Q, const int K, const int d_stride, 
+    const int block_thread,
+    const bool _use_reg) {
+  switch (K) {
+  case 4:
+    flash_deformable_im2col_cuda_inner<scalar_t, 4>(
+        stream,
+        value,                  // B, N, G, D
+        data_spatial_shapes,    // L * 2
+        data_level_start_index, // L
+        offset,                 // B, N, G, L, K, 3
+        output,                 // B, N, G, D
+        B, N, G, D, L, Q, d_stride, 
+        block_thread, _use_reg);
+    break;
+  case 8:
+    flash_deformable_im2col_cuda_inner<scalar_t, 8>(
+        stream,
+        value,                  // B, N, G, D
+        data_spatial_shapes,    // L * 2
+        data_level_start_index, // L
+        offset,                 // B, N, G, L, K, 3
+        output,                 // B, N, G, D
+        B, N, G, D, L, Q, d_stride,
+        block_thread, _use_reg);
+    break;
+  default:
+    printf("not supported for K not in [4, 8]");
+    throw std::invalid_argument("invalid K");
+  }
+}

DCNv4_op/src/dcnv4.h

@@ -0,0 +1,107 @@
+/*!
+**************************************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************************************
+* Modified from
+*https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+#pragma once
+
+#ifdef WITH_CUDA
+#include "cuda/dcnv4_cuda.h"
+#include "cuda/flash_deform_attn_cuda.h"
+#endif
+
+at::Tensor flash_deform_attn_forward(const at::Tensor &value,
+                                        const at::Tensor &spatial_shapes,
+                                        const at::Tensor &level_start_index,
+                                        const at::Tensor &sampling_loc_attn,
+                                        const int im2col_step, const int K,
+                                        const int d_stride, const int block_thread) {
+  if (value.device().is_cuda()) {
+#ifdef WITH_CUDA
+    return flash_deform_attn_cuda_forward(value, spatial_shapes,
+                                             level_start_index,
+                                             sampling_loc_attn, im2col_step, 
+                                             K, d_stride, block_thread);
+#else
+    AT_ERROR("Not compiled with GPU support");
+#endif
+  }
+  AT_ERROR("Not implemented on the CPU");
+}
+
+std::vector<at::Tensor>
+flash_deform_attn_backward(const at::Tensor &value, 
+                              const at::Tensor &spatial_shapes,
+                              const at::Tensor &level_start_index, 
+                              const at::Tensor &sampling_loc_attn,
+                              const at::Tensor &grad_output, 
+                              const int im2col_step, 
+                              const int K, 
+                              const int d_stride, const int block_thread){
+  if (value.device().is_cuda()) {
+#ifdef WITH_CUDA
+    return flash_deform_attn_cuda_backward(value, 
+                                              spatial_shapes,
+                                              level_start_index,
+                                              sampling_loc_attn, 
+                                              grad_output,
+                                              im2col_step, 
+                                              K, d_stride, 
+                                              block_thread);
+#else
+    AT_ERROR("Not compiled with GPU support");
+#endif
+  }
+  AT_ERROR("Not implemented on the CPU");
+}
+
+at::Tensor dcnv4_forward(
+    const at::Tensor &value,
+    const at::Tensor &p_offset,
+    const int kernel_h, const int kernel_w, const int stride_h,
+    const int stride_w, const int pad_h, const int pad_w, const int dilation_h,
+    const int dilation_w, const int group, const int group_channels,
+    const float offset_scale, const int im2col_step, const int remove_center,
+    const int d_stride, const int block_thread, const bool softmax) {
+  if (value.device().is_cuda()) {
+#ifdef WITH_CUDA
+    return dcnv4_cuda_forward(
+        value, p_offset, kernel_h, kernel_w, stride_h, stride_w, pad_h,
+        pad_w, dilation_h, dilation_w, group, group_channels, offset_scale,
+        im2col_step, remove_center, d_stride, block_thread, softmax);
+#else
+    AT_ERROR("Not compiled with GPU support");
+#endif
+  }
+  AT_ERROR("Not implemented on the CPU");
+}
+
+std::vector<at::Tensor>
+dcnv4_backward(
+    const at::Tensor &value, 
+    const at::Tensor &p_offset, 
+    const int kernel_h, const int kernel_w, const int stride_h,
+    const int stride_w, const int pad_h, const int pad_w, const int dilation_h,
+    const int dilation_w, const int group, const int group_channels,
+    const float offset_scale, const int im2col_step, const at::Tensor &grad_output, 
+    const int remove_center, const int d_stride, const int block_thread,
+    const bool softmax){
+  if (value.device().is_cuda()) {
+#ifdef WITH_CUDA
+    return dcnv4_cuda_backward(
+        value, p_offset, kernel_h, kernel_w, stride_h, stride_w, pad_h,
+        pad_w, dilation_h, dilation_w, group, group_channels, offset_scale,
+        im2col_step, grad_output, remove_center, d_stride, block_thread,
+        softmax);
+#else
+    AT_ERROR("Not compiled with GPU support");
+#endif
+  }
+  AT_ERROR("Not implemented on the CPU");
+}

DCNv4_op/src/vision.cpp

@@ -0,0 +1,21 @@
+/*!
+**************************************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************************************
+* Modified from
+*https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+#include "dcnv4.h"
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("flash_deform_attn_forward", &flash_deform_attn_forward,
+        "flash_deform_attn_forward");
+  m.def("flash_deform_attn_backward", &flash_deform_attn_backward,
+        "flash_deform_attn_backward");
+  m.def("dcnv4_forward", &dcnv4_forward, "dcnv4_forward");
+  m.def("dcnv4_backward", &dcnv4_backward, "dcnv4_backward");
+}

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 OpenGVLab
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,158 @@
+
+
+# [DCNv4](https://arxiv.org/pdf/2401.06197.pdf)
+
+
+## News
+- `Jan 15, 2024`: 🚀 Compared with InternImage, the new FlashInternImage powered with DCNv4 has faster inference speed, faster convergence, and better performance!!!
+- `Jan 15, 2024`: 🚀 "DCNv4" is released！
+
+
+## Introduction
+We introduce Deformable Convolution v4 (DCNv4), a highly efficient and effective operator designed for a broad spectrum of vision applications. DCNv4 addresses the limitations of its predecessor, DCNv3, with two key enhancements: 1. removing softmax normalization in spatial aggregation to enhance its dynamic property and expressive power and 2. optimizing memory access to minimize redundant operations for speedup. These improvements result in a significantly faster convergence compared to DCNv3 and a substantial increase in processing speed, with DCNv4 achieving more than three times the forward speed.
+DCNv4 demonstrates exceptional performance across various tasks, including image classification, instance and semantic segmentation, and notably, image generation. 
+When integrated into generative models like U-Net in the latent diffusion model, DCNv4 outperforms its baseline, underscoring its possibility to enhance generative models.
+In practical applications, replacing DCNv3 with DCNv4 in the InternImage model to create FlashInternImage results in up to 80\% speed increase and further performance improvement without further modifications.
+The advancements in speed and efficiency of DCNv4, combined with its robust performance across diverse vision tasks, show its potential as a foundational building block for future vision models.
+
+## Released Models
+
+
+
+<details>
+<summary> ImageNet Image Classification </summary>
+<br>
+<div>
+
+|      name      |   pretrain   | resolution | acc@1 | #param | download                                                                              |
+| :------------: | :----------: | :--------: | :---: | :----: | :---------------------------------------------------------------------------------------------------------------------------------------------------------------: |
+| FlashInternImage-T  | ImageNet-1K  |  224x224   | 83.6  |  30M   |        [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/flash_intern_image_t_1k_224.pth) \| [cfg](classification/configs/flash_intern_image_t_1k_224.yaml)       |
+| FlashInternImage-S  | ImageNet-1K  |  224x224   | 84.4  |  50M   |        [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/flash_intern_image_s_1k_224.pth) \| [cfg](classification/configs/flash_intern_image_s_1k_224.yaml)       |
+| FlashInternImage-B  | ImageNet-1K  |  224x224   | 84.9  |  97M   |      [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/flash_intern_image_b_1k_224.pth) \| [cfg](classification/configs/flash_intern_image_b_1k_224.yaml)       |
+| FlashInternImage-L  | ImageNet-22K |  384x384   | 88.1  |  223M  | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/flash_internimage_l_22kto1k_384.pth) \| [cfg](classification/configs/flash_intern_image_l_22kto1k_384.yaml)  |
+</div>
+
+</details>
+
+<details>
+<summary> COCO Object Detection and Instance Segmentation </summary>
+<br>
+<div>
+
+|    backbone   |method |  schd | box mAP | mask mAP |Config | Download | 
+| :-----------------:| :----------:  |  :---------: | :-----:  |:------: | :-----: | :---: |
+| FlashInternImage-T  |Mask-RCNN|          1x      |  48.0   |   43.1    |  [config](./detection/configs/coco/mask_rcnn_flash_intern_image_t_fpn_1x_coco.py) | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask_rcnn_flash_internimage_t_fpn_1x_coco.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask_rcnn_flash_internimage_t_fpn_1x_coco.log) |
+| FlashInternImage-T  |Mask-RCNN |          3x      |  49.5   |   44.0     | [config](./detection/configs/coco/mask_rcnn_flash_intern_image_t_fpn_3x_coco.py) | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask_rcnn_flash_internimage_t_fpn_3x_coco.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask_rcnn_flash_internimage_t_fpn_3x_coco.log) |
+| FlashInternImage-S   |Mask-RCNN|          1x      |  49.2   |   44.0    |  [config](./detection/configs/coco/mask_rcnn_flash_intern_image_s_fpn_1x_coco.py) | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask_rcnn_flash_internimage_s_fpn_1x_coco.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask_rcnn_flash_internimage_s_fpn_1x_coco.log) |
+| FlashInternImage-S  |Mask-RCNN |          3x      |  50.5   |   44.9   | [config](./detection/configs/coco/mask_rcnn_flash_intern_image_s_fpn_3x_coco.py) | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask_rcnn_flash_internimage_s_fpn_3x_coco.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask_rcnn_flash_internimage_s_fpn_3x_coco.log) |
+| FlashInternImage-B  |Mask-RCNN |          1x      |  50.1   |   44.5  | [config](./detection/configs/coco/mask_rcnn_flash_intern_image_b_fpn_1x_coco.py) | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask_rcnn_flash_internimage_b_fpn_1x_coco.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask_rcnn_flash_internimage_b_fpn_1x_coco.log) |
+| FlashInternImage-B   |Mask-RCNN|          3x      |  50.6   |   45.4  |  [config](./detection/configs/coco/mask_rcnn_flash_intern_image_b_fpn_3x_coco.py)| [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask_rcnn_flash_internimage_b_fpn_3x_coco.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask_rcnn_flash_internimage_b_fpn_3x_coco.log) |
+
+|    backbone    |   method|      schd | box mAP | mask mAP | Config | Download |
+| :------------:|  :---------: |  :---------: | :-----: | :------: | :---: | :---: |
+| FlashInternImage-L |Cascade Mask R-CNN |        1x      |  55.6   |   48.2     | [config](./detection/configs/coco/cascade_flash_intern_image_l_fpn_1x_coco.py) | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/cascade_flash_internimage_l_fpn_1x_coco.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/cascade_flash_internimage_l_fpn_1x_coco.log)
+| FlashInternImage-L |Cascade Mask R-CNN |        3x      |  56.7   |   48.9    | [config](./detection/configs/coco/cascade_flash_intern_image_l_fpn_3x_coco.py) | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/cascade_flash_internimage_l_fpn_3x_coco.pth)  |
+
+|    backbone   |method |  lr type     | pretrain    |       schd | box mAP  | Config | Download |
+| :------------: |  :---------: |  :---------: |:---------: | :---------: | :-----: |  :---: | :-----: |
+| FlashInternImage-T  |DINO| layer-wise lr    | ImageNet-1K  |     1x      |  54.7   |     [config](./detection/configs/coco/dino_4scale_flash_internimage_t_1x_coco.py)     | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/dino_4scale_flash_internimage_t_1x_coco.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/dino_4scale_flash_internimage_t_1x_coco.json) |
+| FlashInternImage-S  |DINO | layer-wise lr    | ImageNet-1K  |     1x      |  55.3   |    [config](./detection/configs/coco/dino_4scale_flash_internimage_s_1x_coco.py)     | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/dino_4scale_flash_internimage_s_1x_coco.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/dino_4scale_flash_internimage_s_1x_coco.log) |
+| FlashInternImage-B   |DINO| layer-wise lr    | ImageNet-1K  |     1x      |  56.0     | [config](./detection/configs/coco/dino_4scale_flash_internimage_b_1x_coco.py)     | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/dino_4scale_flash_internimage_b_1x_coco.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/dino_4scale_flash_internimage_b_1x_coco.log) |
+| FlashInternImage-L  |DINO | 0.1x backbone lr | ImageNet-22K |     1x      |  58.8     |  [config](./detection/configs/coco/dino_4scale_flash_internimage_l_1x_coco.py) | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/dino_4scale_flash_internimage_l_1x_coco.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/dino_4scale_flash_internimage_l_1x_coco.log) |
+
+
+</div>
+
+</details>
+
+
+<details>
+<summary> ADE20K Semantic Segmentation </summary>
+<br>
+<div>
+
+| backbone      |method | resolution | mIoU (ss/ms) | Config | Download            |
+|:--------------:|:----------:|:----------:|:-----------:|:-----------:|:----------:
+| FlashInternImage-T|UperNet  | 512x512    | 49.3 / 50.3   | [config](./segmentation/configs/ade20k/upernet_flash_internimage_t_512_160k_ade20k.py) | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/upernet_flash_internimage_t_512_160k_ade20k.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/upernet_flash_internimage_t_512_160k_ade20k.log)   | 
+| FlashInternImage-S |UperNet   | 512x512    | 50.6 / 51.6     | [config](./segmentation/configs/ade20k/upernet_flash_internimage_s_512_160k_ade20k.py)  | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/upernet_flash_internimage_s_512_160k_ade20k.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/upernet_flash_internimage_s_512_160k_ade20k.log)  | 
+| FlashInternImage-B |UperNet   | 512x512    | 52.0 / 52.6       | [config](./segmentation/configs/ade20k/upernet_flash_internimage_b_512_160k_ade20k.py) | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/upernet_flash_internimage_b_512_160k_ade20k.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/upernet_flash_internimage_s_512_160k_ade20k.log)  | 
+| FlashInternImage-L  |UperNet  | 640x640    | 55.6 / 56.0    | [config](./segmentation/configs/ade20k/upernet_flash_internimage_l_640_160k_ade20k.py)| [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/upernet_flash_internimage_l_640_160k_ade20k.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/upernet_flash_internimage_l_640_160k_ade20k.log)  | 
+
+
+| backbone      |method | resolution | mIoU (ss) | Config | Download            |
+|:--------------:|:----------:|:----------:|:-----------:|:-----------:|:----------:
+| FlashInternImage-T  |Mask2Former| 512x512    | 51.2   | [config](./segmentation/configs/ade20k/mask2former_flash_internimage_t_512_160k_ade20k_ss.py) | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask2former_flash_internimage_t_512_160k_ade20k_ss.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask2former_flash_internimage_t_512_160k_ade20k_ss.log)   | 
+| FlashInternImage-S   |Mask2Former| 640x640    | 52.6     | [config](./segmentation/configs/ade20k/mask2former_flash_internimage_s_640_160k_ade20k_ss.py)  | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask2former_flash_internimage_s_640_160k_ade20k_ss.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask2former_flash_internimage_s_640_160k_ade20k_ss.log)  | 
+| FlashInternImage-B   |Mask2Former| 640x640    |  53.4       | [config](./segmentation/configs/ade20k/mask2former_flash_internimage_b_640_160k_ade20k_ss.py) | [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask2former_flash_internimage_b_640_160k_ade20k_ss.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask2former_flash_internimage_b_640_160k_ade20k_ss.log)  | 
+| FlashInternImage-L   |Mask2Former| 640x640    | 56.7     | [config](./segmentation/configs/ade20k/mask2former_flash_internimage_l_640_160k_ade20k_ss.py)| [ckpt](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask2former_flash_internimage_l_640_160k_ade20k_ss.pth) \| [log](https://huggingface.co/OpenGVLab/DCNv4/resolve/main/mask2former_flash_internimage_l_640_160k_ade20k_ss.log)  | 
+
+
+
+</div>
+
+</details>
+
+## Citations
+
+If this work is helpful for your research, please consider citing the following BibTeX entry.
+
+```bibtex
+
+@article{xiong2024efficient,
+      title={Efficient Deformable ConvNets: Rethinking Dynamic and Sparse Operator for Vision Applications}, 
+      author={Yuwen Xiong and Zhiqi Li and Yuntao Chen and Feng Wang and Xizhou Zhu and Jiapeng Luo and Wenhai Wang and Tong Lu and Hongsheng Li and Yu Qiao and Lewei Lu and Jie Zhou and Jifeng Dai},
+      journal={arXiv preprint arXiv:2401.06197},
+      year={2024}
+}
+
+@article{wang2022internimage,
+  title={InternImage: Exploring Large-Scale Vision Foundation Models with Deformable Convolutions},
+  author={Wang, Wenhai and Dai, Jifeng and Chen, Zhe and Huang, Zhenhang and Li, Zhiqi and Zhu, Xizhou and Hu, Xiaowei and Lu, Tong and Lu, Lewei and Li, Hongsheng and others},
+  journal={arXiv preprint arXiv:2211.05778},
+  year={2022}
+}
+
+@inproceedings{zhu2022uni,
+  title={Uni-perceiver: Pre-training unified architecture for generic perception for zero-shot and few-shot tasks},
+  author={Zhu, Xizhou and Zhu, Jinguo and Li, Hao and Wu, Xiaoshi and Li, Hongsheng and Wang, Xiaohua and Dai, Jifeng},
+  booktitle={CVPR},
+  pages={16804--16815},
+  year={2022}
+}
+
+@article{zhu2022uni,
+  title={Uni-perceiver-moe: Learning sparse generalist models with conditional moes},
+  author={Zhu, Jinguo and Zhu, Xizhou and Wang, Wenhai and Wang, Xiaohua and Li, Hongsheng and Wang, Xiaogang and Dai, Jifeng},
+  journal={arXiv preprint arXiv:2206.04674},
+  year={2022}
+}
+
+@article{li2022uni,
+  title={Uni-Perceiver v2: A Generalist Model for Large-Scale Vision and Vision-Language Tasks},
+  author={Li, Hao and Zhu, Jinguo and Jiang, Xiaohu and Zhu, Xizhou and Li, Hongsheng and Yuan, Chun and Wang, Xiaohua and Qiao, Yu and Wang, Xiaogang and Wang, Wenhai and others},
+  journal={arXiv preprint arXiv:2211.09808},
+  year={2022}
+}
+
+@article{yang2022bevformer,
+  title={BEVFormer v2: Adapting Modern Image Backbones to Bird's-Eye-View Recognition via Perspective Supervision},
+  author={Yang, Chenyu and Chen, Yuntao and Tian, Hao and Tao, Chenxin and Zhu, Xizhou and Zhang, Zhaoxiang and Huang, Gao and Li, Hongyang and Qiao, Yu and Lu, Lewei and others},
+  journal={arXiv preprint arXiv:2211.10439},
+  year={2022}
+}
+
+@article{su2022towards,
+  title={Towards All-in-one Pre-training via Maximizing Multi-modal Mutual Information},
+  author={Su, Weijie and Zhu, Xizhou and Tao, Chenxin and Lu, Lewei and Li, Bin and Huang, Gao and Qiao, Yu and Wang, Xiaogang and Zhou, Jie and Dai, Jifeng},
+  journal={arXiv preprint arXiv:2211.09807},
+  year={2022}
+}
+
+@inproceedings{li2022bevformer,
+  title={Bevformer: Learning bird’s-eye-view representation from multi-camera images via spatiotemporal transformers},
+  author={Li, Zhiqi and Wang, Wenhai and Li, Hongyang and Xie, Enze and Sima, Chonghao and Lu, Tong and Qiao, Yu and Dai, Jifeng},
+  booktitle={ECCV},
+  pages={1--18},
+  year={2022},
+}
+```

classification/README.md

@@ -0,0 +1,190 @@
+# FlashInternImage for Image Classification
+
+This folder contains the implementation of the FlashInternImage for image classification.
+
+## Usage
+
+### Install
+
+- Clone this repo:
+
+```bash
+git clone https://github.com/OpenGVLab/DCNv4.git
+cd DCNv4
+```
+
+- Create a conda virtual environment and activate it:
+
+```bash
+conda create -n internimage python=3.7 -y
+conda activate internimage
+```
+
+- Install `CUDA>=10.2` with `cudnn>=7` following
+  the [official installation instructions](https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html)
+- Install `PyTorch>=1.10.0` and `torchvision>=0.9.0` with `CUDA>=10.2`:
+
+For examples, to install torch==1.11 with CUDA==11.3:
+```bash
+pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113  -f https://download.pytorch.org/whl/torch_stable.html
+```
+
+- Install `timm==0.6.11` and `mmcv-full==1.5.0`:
+
+```bash
+pip install -U openmim
+mim install mmcv-full==1.5.0
+pip install timm==0.6.11 mmdet==2.28.1
+```
+
+- Install other requirements:
+
+```bash
+pip install opencv-python termcolor yacs pyyaml scipy
+```
+
+- Install DCNv4
+```bash
+pip install DCNv4==latest
+```
+
+### Data Preparation
+
+We use standard ImageNet dataset, you can download it from http://image-net.org/. We provide the following two ways to
+load data:
+
+- For standard folder dataset, move validation images to labeled sub-folders. The file structure should look like:
+  ```bash
+  $ tree data
+  imagenet
+  ├── train
+  │   ├── class1
+  │   │   ├── img1.jpeg
+  │   │   ├── img2.jpeg
+  │   │   └── ...
+  │   ├── class2
+  │   │   ├── img3.jpeg
+  │   │   └── ...
+  │   └── ...
+  └── val
+      ├── class1
+      │   ├── img4.jpeg
+      │   ├── img5.jpeg
+      │   └── ...
+      ├── class2
+      │   ├── img6.jpeg
+      │   └── ...
+      └── ...
+ 
+  ```
+- To boost the slow speed when reading images from massive small files, we also support zipped ImageNet, which includes
+  four files:
+    - `train.zip`, `val.zip`: which store the zipped folder for train and validate splits.
+    - `train.txt`, `val.txt`: which store the relative path in the corresponding zip file and ground truth
+      label. Make sure the data folder looks like this:
+
+  ```bash
+  $ tree data
+  data
+  └── ImageNet-Zip
+      ├── train_map.txt
+      ├── train.zip
+      ├── val_map.txt
+      └── val.zip
+  
+  $ head -n 5 meta_data/val.txt
+  ILSVRC2012_val_00000001.JPEG	65
+  ILSVRC2012_val_00000002.JPEG	970
+  ILSVRC2012_val_00000003.JPEG	230
+  ILSVRC2012_val_00000004.JPEG	809
+  ILSVRC2012_val_00000005.JPEG	516
+  
+  $ head -n 5 meta_data/train.txt
+  n01440764/n01440764_10026.JPEG	0
+  n01440764/n01440764_10027.JPEG	0
+  n01440764/n01440764_10029.JPEG	0
+  n01440764/n01440764_10040.JPEG	0
+  n01440764/n01440764_10042.JPEG	0
+  ```
+- For ImageNet-22K dataset, make a folder named `fall11_whole` and move all images to labeled sub-folders in this
+  folder. Then download the train-val split
+  file ([ILSVRC2011fall_whole_map_train.txt](https://github.com/SwinTransformer/storage/releases/download/v2.0.1/ILSVRC2011fall_whole_map_train.txt)
+  & [ILSVRC2011fall_whole_map_val.txt](https://github.com/SwinTransformer/storage/releases/download/v2.0.1/ILSVRC2011fall_whole_map_val.txt))
+  , and put them in the parent directory of `fall11_whole`. The file structure should look like:
+
+  ```bash
+    $ tree imagenet22k/
+    imagenet22k/
+    └── fall11_whole
+        ├── n00004475
+        ├── n00005787
+        ├── n00006024
+        ├── n00006484
+        └── ...
+  ```
+
+### Evaluation
+
+To evaluate a pretrained `FlashInternImage` on ImageNet val, run:
+
+```bash
+python -m torch.distributed.launch --nproc_per_node <num-of-gpus-to-use> --master_port 12345 main.py --eval \
+--cfg <config-file> --resume <checkpoint> --data-path <imagenet-path> 
+```
+
+For example, to evaluate the `FlashInternImage-B` with a single GPU:
+
+```bash
+python -m torch.distributed.launch --nproc_per_node 1 --master_port 12345 main.py --eval \
+--cfg configs/flash_intern_image_b_1k_224.yaml --resume flash_intern_image_b_1k_224.pth --data-path <imagenet-path>
+```
+
+### Training from Scratch on ImageNet-1K
+
+> The paper results were obtained from models trained with configs in `configs/without_lr_decay`.
+
+To train an `InternImage` on ImageNet from scratch, run:
+
+```bash
+python -m torch.distributed.launch --nproc_per_node <num-of-gpus-to-use> --master_port 12345  main.py \ 
+--cfg <config-file> --data-path <imagenet-path> [--batch-size <batch-size-per-gpu> --output <output-directory> --tag <job-tag>]
+```
+
+### Manage Jobs with Slurm.
+
+For example, to train `FlashInternImage` with 8 GPU on a single node for 300 epochs, run:
+
+`FlashInternImage-T`:
+
+```bash
+GPUS=8 sh train_in1k.sh <partition> <job-name> configs/flash_intern_image_t_1k_224.yaml --resume flash_intern_image_t_1k_224.pth --eval
+```
+
+`FlashInternImage-S`:
+
+```bash
+GPUS=8 sh train_in1k.sh <partition> <job-name> configs/flash_intern_image_s_1k_224.yaml --resume flash_intern_image_s_1k_224.pth --eval
+```
+
+
+<!-- 
+### Test pretrained model on ImageNet-22K
+
+For example, to evaluate the `InternImage-L-22k`:
+
+```bash
+python -m torch.distributed.launch --nproc_per_node <num-of-gpus-to-use> --master_port 12345  main.py \ 
+--cfg configs/internimage_xl_22k_192to384.yaml --data-path <imagenet-path> [--batch-size <batch-size-per-gpu> --output <output-directory>] \
+--resume internimage_xl_22k_192to384.pth --eval
+``` -->
+
+<!-- ### Fine-tuning from a ImageNet-22K pretrained model
+
+For example, to fine-tune a `InternImage-XL-22k` model pretrained on ImageNet-22K:
+
+```bashs
+GPUS=8 sh train_in1k.sh <partition> <job-name> configs/intern_image_.yaml --pretrained intern_image_b.pth --eval
+python -m torch.distributed.launch --nproc_per_node 8 --master_port 12345  main.py \
+--cfg configs/.yaml --pretrained swin_base_patch4_window7_224_22k.pth \
+--data-path <imagenet-path> --batch-size 64 --accumulation-steps 2 [--use-checkpoint]
+``` -->

classification/config.py

@@ -0,0 +1,302 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+import os
+import yaml
+from yacs.config import CfgNode as CN
+
+_C = CN()
+
+# Base config files
+_C.BASE = ['']
+
+# -----------------------------------------------------------------------------
+# Data settings
+# -----------------------------------------------------------------------------
+_C.DATA = CN()
+# Batch size for a single GPU, could be overwritten by command line argument
+_C.DATA.BATCH_SIZE = 128
+# Path to dataset, could be overwritten by command line argument
+_C.DATA.DATA_PATH = ''
+# Dataset name
+_C.DATA.DATASET = 'imagenet'
+# Input image size
+_C.DATA.IMG_SIZE = 224
+# Interpolation to resize image (random, bilinear, bicubic)
+_C.DATA.INTERPOLATION = 'bicubic'
+# Use zipped dataset instead of folder dataset
+# could be overwritten by command line argument
+_C.DATA.ZIP_MODE = False
+# Cache Data in Memory, could be overwritten by command line argument
+_C.DATA.CACHE_MODE = 'part'
+# Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.
+_C.DATA.PIN_MEMORY = True
+# Number of data loading threads
+_C.DATA.NUM_WORKERS = 8
+# Load data to memory
+_C.DATA.IMG_ON_MEMORY = False
+
+# -----------------------------------------------------------------------------
+# Model settings
+# -----------------------------------------------------------------------------
+_C.MODEL = CN()
+# Model type
+_C.MODEL.TYPE = 'INTERN_IMAGE'
+# Model name
+_C.MODEL.NAME = 'intern_image'
+# Pretrained weight from checkpoint, could be imagenet22k pretrained weight
+# could be overwritten by command line argument
+_C.MODEL.PRETRAINED = ''
+# Checkpoint to resume, could be overwritten by command line argument
+_C.MODEL.RESUME = ''
+# Number of classes, overwritten in data preparation
+_C.MODEL.NUM_CLASSES = 1000
+# Dropout rate
+_C.MODEL.DROP_RATE = 0.0
+# Drop path rate
+_C.MODEL.DROP_PATH_RATE = 0.1
+# Drop path type
+_C.MODEL.DROP_PATH_TYPE = 'linear'  # linear, uniform
+# Label Smoothing
+_C.MODEL.LABEL_SMOOTHING = 0.1
+
+# INTERN_IMAGE parameters
+_C.MODEL.INTERN_IMAGE = CN()
+_C.MODEL.INTERN_IMAGE.DEPTHS = [4, 4, 18, 4]
+_C.MODEL.INTERN_IMAGE.GROUPS = [4, 8, 16, 32]
+_C.MODEL.INTERN_IMAGE.CHANNELS = 64
+_C.MODEL.INTERN_IMAGE.LAYER_SCALE = None
+_C.MODEL.INTERN_IMAGE.OFFSET_SCALE = 1.0
+_C.MODEL.INTERN_IMAGE.MLP_RATIO = 4.0
+_C.MODEL.INTERN_IMAGE.CORE_OP = 'DCNv3'
+_C.MODEL.INTERN_IMAGE.POST_NORM = False
+_C.MODEL.INTERN_IMAGE.RES_POST_NORM = False
+_C.MODEL.INTERN_IMAGE.DW_KERNEL_SIZE = None
+_C.MODEL.INTERN_IMAGE.USE_CLIP_PROJECTOR = False
+_C.MODEL.INTERN_IMAGE.LEVEL2_POST_NORM = False
+_C.MODEL.INTERN_IMAGE.LEVEL2_POST_NORM_BLOCK_IDS = None
+_C.MODEL.INTERN_IMAGE.CENTER_FEATURE_SCALE = False
+
+# FLASH_INTERN_IMAGE parameters
+_C.MODEL.FLASH_INTERN_IMAGE = CN()
+_C.MODEL.FLASH_INTERN_IMAGE.DEPTHS = [4, 4, 18, 4]
+_C.MODEL.FLASH_INTERN_IMAGE.GROUPS = [4, 8, 16, 32]
+_C.MODEL.FLASH_INTERN_IMAGE.CHANNELS = 64
+_C.MODEL.FLASH_INTERN_IMAGE.LAYER_SCALE = None
+_C.MODEL.FLASH_INTERN_IMAGE.OFFSET_SCALE = 1.0
+_C.MODEL.FLASH_INTERN_IMAGE.MLP_RATIO = 4.0
+_C.MODEL.FLASH_INTERN_IMAGE.CORE_OP = 'DCNv4'
+_C.MODEL.FLASH_INTERN_IMAGE.POST_NORM = False
+_C.MODEL.FLASH_INTERN_IMAGE.RES_POST_NORM = False
+_C.MODEL.FLASH_INTERN_IMAGE.DW_KERNEL_SIZE = None
+_C.MODEL.FLASH_INTERN_IMAGE.USE_CLIP_PROJECTOR = False
+_C.MODEL.FLASH_INTERN_IMAGE.LEVEL2_POST_NORM = False
+_C.MODEL.FLASH_INTERN_IMAGE.LEVEL2_POST_NORM_BLOCK_IDS = None
+_C.MODEL.FLASH_INTERN_IMAGE.CENTER_FEATURE_SCALE = False
+_C.MODEL.FLASH_INTERN_IMAGE.MLP_FC2_BIAS = False
+_C.MODEL.FLASH_INTERN_IMAGE.DCN_OUTPUT_BIAS = False
+# -----------------------------------------------------------------------------
+# Training settings
+# -----------------------------------------------------------------------------
+_C.TRAIN = CN()
+_C.TRAIN.START_EPOCH = 0
+_C.TRAIN.EPOCHS = 300
+_C.TRAIN.WARMUP_EPOCHS = 20
+_C.TRAIN.WEIGHT_DECAY = 0.05
+_C.TRAIN.BASE_LR = 5e-4
+_C.TRAIN.WARMUP_LR = 5e-7
+_C.TRAIN.MIN_LR = 5e-6
+# Clip gradient norm
+_C.TRAIN.CLIP_GRAD = 5.0
+# Auto resume from latest checkpoint
+_C.TRAIN.AUTO_RESUME = True
+# Gradient accumulation steps
+# could be overwritten by command line argument
+_C.TRAIN.ACCUMULATION_STEPS = 0
+# Whether to use gradient checkpointing to save memory
+# could be overwritten by command line argument
+_C.TRAIN.USE_CHECKPOINT = False
+
+# LR scheduler
+_C.TRAIN.LR_SCHEDULER = CN()
+_C.TRAIN.LR_SCHEDULER.NAME = 'cosine'
+# Epoch interval to decay LR, used in StepLRScheduler
+_C.TRAIN.LR_SCHEDULER.DECAY_EPOCHS = 30
+# LR decay rate, used in StepLRScheduler
+_C.TRAIN.LR_SCHEDULER.DECAY_RATE = 0.1
+
+# Optimizer
+_C.TRAIN.OPTIMIZER = CN()
+_C.TRAIN.OPTIMIZER.NAME = 'adamw'
+# Optimizer Epsilon
+_C.TRAIN.OPTIMIZER.EPS = 1e-8
+# Optimizer Betas
+_C.TRAIN.OPTIMIZER.BETAS = (0.9, 0.999)
+# SGD momentum
+_C.TRAIN.OPTIMIZER.MOMENTUM = 0.9
+# ZeRO
+_C.TRAIN.OPTIMIZER.USE_ZERO = False
+# freeze backbone
+_C.TRAIN.OPTIMIZER.FREEZE_BACKBONE = None
+# dcn lr
+_C.TRAIN.OPTIMIZER.DCN_LR_MUL = None
+
+# EMA
+_C.TRAIN.EMA = CN()
+_C.TRAIN.EMA.ENABLE = False
+_C.TRAIN.EMA.DECAY = 0.9998
+
+# LR_LAYER_DECAY
+_C.TRAIN.LR_LAYER_DECAY = False
+_C.TRAIN.LR_LAYER_DECAY_RATIO = 0.875
+
+# FT head init weights
+_C.TRAIN.RAND_INIT_FT_HEAD = False
+
+# -----------------------------------------------------------------------------
+# Augmentation settings
+# -----------------------------------------------------------------------------
+_C.AUG = CN()
+# Color jitter factor
+_C.AUG.COLOR_JITTER = 0.4
+# Use AutoAugment policy. "v0" or "original"
+_C.AUG.AUTO_AUGMENT = 'rand-m9-mstd0.5-inc1'
+# Random erase prob
+_C.AUG.REPROB = 0.25
+# Random erase mode
+_C.AUG.REMODE = 'pixel'
+# Random erase count
+_C.AUG.RECOUNT = 1
+# Mixup alpha, mixup enabled if > 0
+_C.AUG.MIXUP = 0.8
+# Cutmix alpha, cutmix enabled if > 0
+_C.AUG.CUTMIX = 1.0
+# Cutmix min/max ratio, overrides alpha and enables cutmix if set
+_C.AUG.CUTMIX_MINMAX = None
+# Probability of performing mixup or cutmix when either/both is enabled
+_C.AUG.MIXUP_PROB = 1.0
+# Probability of switching to cutmix when both mixup and cutmix enabled
+_C.AUG.MIXUP_SWITCH_PROB = 0.5
+# How to apply mixup/cutmix params. Per "batch", "pair", or "elem"
+_C.AUG.MIXUP_MODE = 'batch'
+# RandomResizedCrop
+_C.AUG.RANDOM_RESIZED_CROP = False
+_C.AUG.MEAN = (0.485, 0.456, 0.406)
+_C.AUG.STD = (0.229, 0.224, 0.225)
+
+# -----------------------------------------------------------------------------
+# Testing settings
+# -----------------------------------------------------------------------------
+_C.TEST = CN()
+# Whether to use center crop when testing
+_C.TEST.CROP = True
+
+# Whether to use SequentialSampler as validation sampler
+_C.TEST.SEQUENTIAL = False
+
+# -----------------------------------------------------------------------------
+# Misc
+# -----------------------------------------------------------------------------
+# Mixed precision opt level, if O0, no amp is used ('O0', 'O1', 'O2')
+# overwritten by command line argument
+_C.AMP_OPT_LEVEL = ''
+# Path to output folder, overwritten by command line argument
+_C.OUTPUT = ''
+# Tag of experiment, overwritten by command line argument
+_C.TAG = 'default'
+# Frequency to save checkpoint
+_C.SAVE_FREQ = 1
+# Frequency to logging info
+_C.PRINT_FREQ = 10
+# eval freq
+_C.EVAL_FREQ = 1
+# Fixed random seed
+_C.SEED = 0
+# Perform evaluation only, overwritten by command line argument
+_C.EVAL_MODE = False
+# Test throughput only, overwritten by command line argument
+_C.THROUGHPUT_MODE = False
+# local rank for DistributedDataParallel, given by command line argument
+_C.LOCAL_RANK = 0
+_C.EVAL_22K_TO_1K = False
+
+_C.AMP_TYPE = 'float16'
+
+
+def _update_config_from_file(config, cfg_file):
+    config.defrost()
+    with open(cfg_file, 'r') as f:
+        yaml_cfg = yaml.load(f, Loader=yaml.FullLoader)
+
+    for cfg in yaml_cfg.setdefault('BASE', ['']):
+        if cfg:
+            _update_config_from_file(
+                config, os.path.join(os.path.dirname(cfg_file), cfg))
+    print('=> merge config from {}'.format(cfg_file))
+    config.merge_from_file(cfg_file)
+    config.freeze()
+
+
+def update_config(config, args):
+    _update_config_from_file(config, args.cfg)
+
+    config.defrost()
+    if hasattr(args, 'opts') and args.opts:
+        config.merge_from_list(args.opts)
+
+    # merge from specific arguments
+    if hasattr(args, 'batch_size') and args.batch_size:
+        config.DATA.BATCH_SIZE = args.batch_size
+    if hasattr(args, 'dataset') and args.dataset:
+        config.DATA.DATASET = args.dataset
+    if hasattr(args, 'data_path') and args.data_path:
+        config.DATA.DATA_PATH = args.data_path
+    if hasattr(args, 'zip') and args.zip:
+        config.DATA.ZIP_MODE = True
+    if hasattr(args, 'cache_mode') and args.cache_mode:
+        config.DATA.CACHE_MODE = args.cache_mode
+    if hasattr(args, 'pretrained') and args.pretrained:
+        config.MODEL.PRETRAINED = args.pretrained
+    if hasattr(args, 'resume') and args.resume:
+        config.MODEL.RESUME = args.resume
+    if hasattr(args, 'accumulation_steps') and args.accumulation_steps:
+        config.TRAIN.ACCUMULATION_STEPS = args.accumulation_steps
+    if hasattr(args, 'use_checkpoint') and args.use_checkpoint:
+        config.TRAIN.USE_CHECKPOINT = True
+    if hasattr(args, 'amp_opt_level') and args.amp_opt_level:
+        config.AMP_OPT_LEVEL = args.amp_opt_level
+    if hasattr(args, 'output') and args.output:
+        config.OUTPUT = args.output
+    if hasattr(args, 'tag') and args.tag:
+        config.TAG = args.tag
+    if hasattr(args, 'eval') and args.eval:
+        config.EVAL_MODE = True
+    if hasattr(args, 'throughput') and args.throughput:
+        config.THROUGHPUT_MODE = True
+    if hasattr(args, 'save_ckpt_num') and args.save_ckpt_num:
+        config.SAVE_CKPT_NUM = args.save_ckpt_num
+    if hasattr(args, 'use_zero') and args.use_zero:
+        config.TRAIN.OPTIMIZER.USE_ZERO = True
+    # set local rank for distributed training
+    if hasattr(args, 'local_rank') and args.local_rank:
+        config.LOCAL_RANK = args.local_rank
+
+    # output folder
+    config.MODEL.NAME = args.cfg.split('/')[-1].replace('.yaml', '')
+    config.OUTPUT = os.path.join(config.OUTPUT, config.MODEL.NAME)
+    # config.OUTPUT = os.path.join(config.OUTPUT, config.MODEL.NAME, config.TAG)
+
+    config.freeze()
+
+
+def get_config(args):
+    """Get a yacs CfgNode object with default values."""
+    # Return a clone so that the defaults will not be altered
+    # This is for the "local variable" use pattern
+    config = _C.clone()
+    update_config(config, args)
+
+    return config

classification/configs/flash_intern_image_b_1k_224.yaml

@@ -0,0 +1,20 @@
+DATA:
+  IMG_ON_MEMORY: False
+MODEL:
+  TYPE: flash_intern_image
+  DROP_PATH_RATE: 0.5
+  FLASH_INTERN_IMAGE:
+    CORE_OP: 'DCNv4'
+    DEPTHS: [4, 4, 21, 4]
+    GROUPS: [7, 14, 28, 56]
+    CHANNELS: 112
+    LAYER_SCALE: 1e-5
+    OFFSET_SCALE: 0.5
+    MLP_RATIO: 4.0
+    POST_NORM: True
+    DW_KERNEL_SIZE: 3
+TRAIN:
+  EMA:
+    ENABLE: True
+    DECAY: 0.9999
+  BASE_LR: 5e-4

classification/configs/flash_intern_image_l_22kto1k_384.yaml

@@ -0,0 +1,40 @@
+DATA:
+  IMG_SIZE: 384
+  IMG_ON_MEMORY: False
+AUG:
+  MIXUP: 0.0
+  CUTMIX: 0.0
+  REPROB: 0.0
+MODEL:
+  TYPE: flash_intern_image
+  DROP_PATH_RATE: 0.1
+  LABEL_SMOOTHING: 0.3
+  FLASH_INTERN_IMAGE:
+    CORE_OP: 'DCNv4'
+    DEPTHS: [5, 5, 22, 5]
+    GROUPS: [10, 20, 40, 80]
+    CHANNELS: 160
+    LAYER_SCALE: 1e-5
+    OFFSET_SCALE: 2.0
+    MLP_RATIO: 4.0
+    POST_NORM: True
+    DW_KERNEL_SIZE: 3
+    DCN_OUTPUT_BIAS: True
+    MLP_FC2_BIAS: True
+TRAIN:
+  EMA:
+    ENABLE: true
+    DECAY: 0.9999
+  EPOCHS: 20
+  WARMUP_EPOCHS: 2
+  WEIGHT_DECAY: 0.05
+  BASE_LR: 2e-05 # 512
+  WARMUP_LR: .0
+  MIN_LR: .0
+  LR_LAYER_DECAY: true
+  LR_LAYER_DECAY_RATIO: 0.9
+  USE_CHECKPOINT: true
+  OPTIMIZER:
+    DCN_LR_MUL: 0.1
+AMP_OPT_LEVEL: O0
+EVAL_FREQ: 1

classification/configs/flash_intern_image_s_1k_224.yaml

@@ -0,0 +1,20 @@
+DATA:
+  IMG_ON_MEMORY: False
+MODEL:
+  TYPE: flash_intern_image
+  DROP_PATH_RATE: 0.4
+  FLASH_INTERN_IMAGE:
+    CORE_OP: 'DCNv4'
+    DEPTHS: [4, 4, 21, 4]
+    GROUPS: [5, 10, 20, 40]
+    CHANNELS: 80
+    LAYER_SCALE: 1e-5
+    OFFSET_SCALE: 1.0
+    MLP_RATIO: 4.0
+    POST_NORM: True
+    DW_KERNEL_SIZE: 3
+TRAIN:
+  EMA:
+    ENABLE: True
+    DECAY: 0.9999
+  BASE_LR: 5e-4

classification/configs/flash_intern_image_t_1k_224.yaml

@@ -0,0 +1,17 @@
+DATA:
+  IMG_ON_MEMORY: False
+MODEL:
+  TYPE: flash_intern_image
+  DROP_PATH_RATE: 0.1
+  FLASH_INTERN_IMAGE:
+    CORE_OP: 'DCNv4'
+    DEPTHS: [4, 4, 18, 4]
+    GROUPS: [4, 8, 16, 32]
+    CHANNELS: 64
+    OFFSET_SCALE: 1.0
+    MLP_RATIO: 4.0
+TRAIN:
+  EMA:
+    ENABLE: True
+    DECAY: 0.9999
+  BASE_LR: 5e-4

classification/dataset/__init__.py

@@ -0,0 +1,7 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+from .build import build_loader, build_loader2

classification/dataset/build.py

@@ -0,0 +1,286 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+import os
+import torch
+import numpy as np
+import torch.distributed as dist
+from torchvision import transforms
+from timm.data import Mixup
+from timm.data import create_transform
+from .cached_image_folder import ImageCephDataset
+from .samplers import SubsetRandomSampler, NodeDistributedSampler
+
+try:
+    from torchvision.transforms import InterpolationMode
+
+    def _pil_interp(method):
+        if method == 'bicubic':
+            return InterpolationMode.BICUBIC
+        elif method == 'lanczos':
+            return InterpolationMode.LANCZOS
+        elif method == 'hamming':
+            return InterpolationMode.HAMMING
+        else:
+            return InterpolationMode.BILINEAR
+except:
+    from timm.data.transforms import _pil_interp
+
+
+class TTA(torch.nn.Module):
+
+    def __init__(self, size, scales=[1.0, 1.05, 1.1]):
+        super().__init__()
+        self.size = size
+        self.scales = scales
+
+    def forward(self, img):
+        out = []
+        cc = transforms.CenterCrop(self.size)
+        for scale in self.scales:
+            size_ = int(scale * self.size)
+            rs = transforms.Resize(size_, interpolation=_pil_interp('bicubic'))
+            img_ = rs(img)
+            img_ = cc(img_)
+            out.append(img_)
+
+        return out
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(size={self.size}, scale={self.scales})"
+
+
+def build_loader(config):
+    config.defrost()
+    dataset_train, config.MODEL.NUM_CLASSES = build_dataset('train',
+                                                            config=config)
+    config.freeze()
+    print(f"local rank {config.LOCAL_RANK} / global rank {dist.get_rank()}"
+          "successfully build train dataset")
+
+    dataset_val, _ = build_dataset('val', config=config)
+    print(f"local rank {config.LOCAL_RANK} / global rank {dist.get_rank()}"
+          "successfully build val dataset")
+
+    dataset_test, _ = build_dataset('test', config=config)
+    print(f"local rank {config.LOCAL_RANK} / global rank {dist.get_rank()}"
+          "successfully build test dataset")
+
+    num_tasks = dist.get_world_size()
+    global_rank = dist.get_rank()
+
+    if dataset_train is not None:
+        if config.DATA.IMG_ON_MEMORY:
+            sampler_train = NodeDistributedSampler(dataset_train)
+        else:
+            if config.DATA.ZIP_MODE and config.DATA.CACHE_MODE == 'part':
+                indices = np.arange(dist.get_rank(), len(dataset_train),
+                                    dist.get_world_size())
+                sampler_train = SubsetRandomSampler(indices)
+            else:
+                sampler_train = torch.utils.data.DistributedSampler(
+                    dataset_train,
+                    num_replicas=num_tasks,
+                    rank=global_rank,
+                    shuffle=True)
+
+    if dataset_val is not None:
+        if config.TEST.SEQUENTIAL:
+            sampler_val = torch.utils.data.SequentialSampler(dataset_val)
+        else:
+            sampler_val = torch.utils.data.distributed.DistributedSampler(
+                dataset_val, shuffle=False)
+
+    if dataset_test is not None:
+        if config.TEST.SEQUENTIAL:
+            sampler_test = torch.utils.data.SequentialSampler(dataset_test)
+        else:
+            sampler_test = torch.utils.data.distributed.DistributedSampler(
+                dataset_test, shuffle=False)
+
+    data_loader_train = torch.utils.data.DataLoader(
+        dataset_train,
+        sampler=sampler_train,
+        batch_size=config.DATA.BATCH_SIZE,
+        num_workers=config.DATA.NUM_WORKERS,
+        pin_memory=config.DATA.PIN_MEMORY,
+        drop_last=True,
+        persistent_workers=True) if dataset_train is not None else None
+
+    data_loader_val = torch.utils.data.DataLoader(
+        dataset_val,
+        sampler=sampler_val,
+        batch_size=config.DATA.BATCH_SIZE,
+        shuffle=False,
+        num_workers=config.DATA.NUM_WORKERS,
+        pin_memory=config.DATA.PIN_MEMORY,
+        drop_last=False,
+        persistent_workers=True) if dataset_val is not None else None
+
+    data_loader_test = torch.utils.data.DataLoader(
+        dataset_test,
+        sampler=sampler_test,
+        batch_size=config.DATA.BATCH_SIZE,
+        shuffle=False,
+        num_workers=config.DATA.NUM_WORKERS,
+        pin_memory=config.DATA.PIN_MEMORY,
+        drop_last=False,
+        persistent_workers=True) if dataset_test is not None else None
+
+    # setup mixup / cutmix
+    mixup_fn = None
+    mixup_active = config.AUG.MIXUP > 0 or config.AUG.CUTMIX > 0. or config.AUG.CUTMIX_MINMAX is not None
+    if mixup_active:
+        mixup_fn = Mixup(mixup_alpha=config.AUG.MIXUP,
+                         cutmix_alpha=config.AUG.CUTMIX,
+                         cutmix_minmax=config.AUG.CUTMIX_MINMAX,
+                         prob=config.AUG.MIXUP_PROB,
+                         switch_prob=config.AUG.MIXUP_SWITCH_PROB,
+                         mode=config.AUG.MIXUP_MODE,
+                         label_smoothing=config.MODEL.LABEL_SMOOTHING,
+                         num_classes=config.MODEL.NUM_CLASSES)
+
+    return dataset_train, dataset_val, dataset_test, data_loader_train, \
+        data_loader_val, data_loader_test, mixup_fn
+
+
+def build_loader2(config):
+    config.defrost()
+    dataset_train, config.MODEL.NUM_CLASSES = build_dataset('train',
+                                                            config=config)
+    config.freeze()
+    dataset_val, _ = build_dataset('val', config=config)
+    dataset_test, _ = build_dataset('test', config=config)
+
+    data_loader_train = torch.utils.data.DataLoader(
+        dataset_train,
+        shuffle=True,
+        batch_size=config.DATA.BATCH_SIZE,
+        num_workers=config.DATA.NUM_WORKERS,
+        pin_memory=config.DATA.PIN_MEMORY,
+        drop_last=True,
+        persistent_workers=True) if dataset_train is not None else None
+
+    data_loader_val = torch.utils.data.DataLoader(
+        dataset_val,
+        batch_size=config.DATA.BATCH_SIZE,
+        shuffle=False,
+        num_workers=config.DATA.NUM_WORKERS,
+        pin_memory=config.DATA.PIN_MEMORY,
+        drop_last=False,
+        persistent_workers=True) if dataset_val is not None else None
+
+    data_loader_test = torch.utils.data.DataLoader(
+        dataset_test,
+        batch_size=config.DATA.BATCH_SIZE,
+        shuffle=False,
+        num_workers=config.DATA.NUM_WORKERS,
+        pin_memory=config.DATA.PIN_MEMORY,
+        drop_last=False,
+        persistent_workers=True) if dataset_test is not None else None
+
+    # setup mixup / cutmix
+    mixup_fn = None
+    mixup_active = config.AUG.MIXUP > 0 or config.AUG.CUTMIX > 0. or config.AUG.CUTMIX_MINMAX is not None
+    if mixup_active:
+        mixup_fn = Mixup(mixup_alpha=config.AUG.MIXUP,
+                         cutmix_alpha=config.AUG.CUTMIX,
+                         cutmix_minmax=config.AUG.CUTMIX_MINMAX,
+                         prob=config.AUG.MIXUP_PROB,
+                         switch_prob=config.AUG.MIXUP_SWITCH_PROB,
+                         mode=config.AUG.MIXUP_MODE,
+                         label_smoothing=config.MODEL.LABEL_SMOOTHING,
+                         num_classes=config.MODEL.NUM_CLASSES)
+
+    return dataset_train, dataset_val, dataset_test, data_loader_train, \
+        data_loader_val, data_loader_test, mixup_fn
+
+
+def build_dataset(split, config):
+    transform = build_transform(split == 'train', config)
+    dataset = None
+    nb_classes = None
+    prefix = split
+    if config.DATA.DATASET == 'imagenet':
+        if prefix == 'train' and not config.EVAL_MODE:
+            root = os.path.join(config.DATA.DATA_PATH, 'train')
+            dataset = ImageCephDataset(root,
+                                       'train',
+                                       transform=transform,
+                                       on_memory=config.DATA.IMG_ON_MEMORY)
+        elif prefix == 'val':
+            root = os.path.join(config.DATA.DATA_PATH, 'val')
+            dataset = ImageCephDataset(root, 'val', transform=transform)
+        nb_classes = 1000
+    elif config.DATA.DATASET == 'imagenet22K':
+        if prefix == 'train':
+            if not config.EVAL_MODE:
+                root = config.DATA.DATA_PATH
+                dataset = ImageCephDataset(root,
+                                           'train',
+                                           transform=transform,
+                                           on_memory=config.DATA.IMG_ON_MEMORY)
+            nb_classes = 21841
+        elif prefix == 'val':
+            root = os.path.join(config.DATA.DATA_PATH, 'val')
+            dataset = ImageCephDataset(root, 'val', transform=transform)
+            nb_classes = 1000
+    else:
+        raise NotImplementedError(
+            f'build_dataset does support {config.DATA.DATASET}')
+
+    return dataset, nb_classes
+
+
+def build_transform(is_train, config):
+    resize_im = config.DATA.IMG_SIZE > 32
+    if is_train:
+        # this should always dispatch to transforms_imagenet_train
+        transform = create_transform(
+            input_size=config.DATA.IMG_SIZE,
+            is_training=True,
+            color_jitter=config.AUG.COLOR_JITTER
+            if config.AUG.COLOR_JITTER > 0 else None,
+            auto_augment=config.AUG.AUTO_AUGMENT
+            if config.AUG.AUTO_AUGMENT != 'none' else None,
+            re_prob=config.AUG.REPROB,
+            re_mode=config.AUG.REMODE,
+            re_count=config.AUG.RECOUNT,
+            interpolation=config.DATA.INTERPOLATION,
+        )
+        if not resize_im:
+            # replace RandomResizedCropAndInterpolation with
+            # RandomCrop
+            transform.transforms[0] = transforms.RandomCrop(
+                config.DATA.IMG_SIZE, padding=4)
+
+        return transform
+
+    t = []
+    if resize_im:
+        if config.TEST.CROP:
+            size = int(1.0 * config.DATA.IMG_SIZE)
+            t.append(
+                transforms.Resize(size,
+                                  interpolation=_pil_interp(
+                                      config.DATA.INTERPOLATION)),
+                # to maintain same ratio w.r.t. 224 images
+            )
+            t.append(transforms.CenterCrop(config.DATA.IMG_SIZE))
+        elif config.AUG.RANDOM_RESIZED_CROP:
+            t.append(
+                transforms.RandomResizedCrop(
+                    (config.DATA.IMG_SIZE, config.DATA.IMG_SIZE),
+                    interpolation=_pil_interp(config.DATA.INTERPOLATION)))
+        else:
+            t.append(
+                transforms.Resize(
+                    (config.DATA.IMG_SIZE, config.DATA.IMG_SIZE),
+                    interpolation=_pil_interp(config.DATA.INTERPOLATION)))
+    t.append(transforms.ToTensor())
+    t.append(transforms.Normalize(config.AUG.MEAN, config.AUG.STD))
+
+    return transforms.Compose(t)

classification/dataset/cached_image_folder.py

@@ -0,0 +1,538 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+import io
+import os
+import re
+import time
+import json
+import math
+import mmcv
+import torch
+import logging
+import os.path as osp
+from PIL import Image
+from tqdm import tqdm, trange
+from abc import abstractmethod
+import torch.utils.data as data
+import torch.distributed as dist
+from mmcv.fileio import FileClient
+from .zipreader import is_zip_path, ZipReader
+
+_logger = logging.getLogger(__name__)
+
+_ERROR_RETRY = 50
+
+
+def has_file_allowed_extension(filename, extensions):
+    """Checks if a file is an allowed extension.
+    Args:
+        filename (string): path to a file
+    Returns:
+        bool: True if the filename ends with a known image extension
+    """
+    filename_lower = filename.lower()
+    return any(filename_lower.endswith(ext) for ext in extensions)
+
+
+def find_classes(dir):
+    classes = [
+        d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))
+    ]
+    classes.sort()
+    class_to_idx = {classes[i]: i for i in range(len(classes))}
+    return classes, class_to_idx
+
+
+def make_dataset(dir, class_to_idx, extensions):
+    images = []
+    dir = os.path.expanduser(dir)
+    for target in sorted(os.listdir(dir)):
+        d = os.path.join(dir, target)
+        if not os.path.isdir(d):
+            continue
+        for root, _, fnames in sorted(os.walk(d)):
+            for fname in sorted(fnames):
+                if has_file_allowed_extension(fname, extensions):
+                    path = os.path.join(root, fname)
+                    item = (path, class_to_idx[target])
+                    images.append(item)
+
+    return images
+
+
+def make_dataset_with_ann(ann_file, img_prefix, extensions):
+    images = []
+    with open(ann_file, "r") as f:
+        contents = f.readlines()
+        for line_str in contents:
+            path_contents = [c for c in line_str.split('\t')]
+            im_file_name = path_contents[0]
+            class_index = int(path_contents[1])
+            assert str.lower(os.path.splitext(im_file_name)[-1]) in extensions
+            item = (os.path.join(img_prefix, im_file_name), class_index)
+            images.append(item)
+
+    return images
+
+
+class DatasetFolder(data.Dataset):
+    """A generic data loader where the samples are arranged in this way: ::
+        root/class_x/xxx.ext
+        root/class_x/xxy.ext
+        root/class_x/xxz.ext
+        root/class_y/123.ext
+        root/class_y/nsdf3.ext
+        root/class_y/asd932_.ext
+    Args:
+        root (string): Root directory path.
+        loader (callable): A function to load a sample given its path.
+        extensions (list[string]): A list of allowed extensions.
+        transform (callable, optional): A function/transform that takes in
+            a sample and returns a transformed version.
+            E.g, ``transforms.RandomCrop`` for images.
+        target_transform (callable, optional): A function/transform that takes
+            in the target and transforms it.
+     Attributes:
+        samples (list): List of (sample path, class_index) tuples
+    """
+
+    def __init__(self,
+                 root,
+                 loader,
+                 extensions,
+                 ann_file='',
+                 img_prefix='',
+                 transform=None,
+                 target_transform=None,
+                 cache_mode="no"):
+        # image folder mode
+        if ann_file == '':
+            _, class_to_idx = find_classes(root)
+            samples = make_dataset(root, class_to_idx, extensions)
+        # zip mode
+        else:
+            samples = make_dataset_with_ann(os.path.join(root, ann_file),
+                                            os.path.join(root, img_prefix),
+                                            extensions)
+
+        if len(samples) == 0:
+            raise (RuntimeError("Found 0 files in subfolders of: " + root +
+                                "\n" + "Supported extensions are: " +
+                                ",".join(extensions)))
+
+        self.root = root
+        self.loader = loader
+        self.extensions = extensions
+
+        self.samples = samples
+        self.labels = [y_1k for _, y_1k in samples]
+        self.classes = list(set(self.labels))
+
+        self.transform = transform
+        self.target_transform = target_transform
+
+        self.cache_mode = cache_mode
+        if self.cache_mode != "no":
+            self.init_cache()
+
+    def init_cache(self):
+        assert self.cache_mode in ["part", "full"]
+        n_sample = len(self.samples)
+        global_rank = dist.get_rank()
+        world_size = dist.get_world_size()
+
+        samples_bytes = [None for _ in range(n_sample)]
+        start_time = time.time()
+        for index in range(n_sample):
+            if index % (n_sample // 10) == 0:
+                t = time.time() - start_time
+                print(
+                    f'global_rank {dist.get_rank()} cached {index}/{n_sample} takes {t:.2f}s per block'
+                )
+                start_time = time.time()
+            path, target = self.samples[index]
+            if self.cache_mode == "full":
+                samples_bytes[index] = (ZipReader.read(path), target)
+            elif self.cache_mode == "part" and index % world_size == global_rank:
+                samples_bytes[index] = (ZipReader.read(path), target)
+            else:
+                samples_bytes[index] = (path, target)
+        self.samples = samples_bytes
+
+    def __getitem__(self, index):
+        """
+        Args:
+            index (int): Index
+        Returns:
+            tuple: (sample, target) where target is class_index of the target class.
+        """
+        path, target = self.samples[index]
+        sample = self.loader(path)
+        if self.transform is not None:
+            sample = self.transform(sample)
+        if self.target_transform is not None:
+            target = self.target_transform(target)
+
+        return sample, target
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __repr__(self):
+        fmt_str = 'Dataset ' + self.__class__.__name__ + '\n'
+        fmt_str += '    Number of datapoints: {}\n'.format(self.__len__())
+        fmt_str += '    Root Location: {}\n'.format(self.root)
+        tmp = '    Transforms (if any): '
+        fmt_str += '{0}{1}\n'.format(
+            tmp,
+            self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
+        tmp = '    Target Transforms (if any): '
+        fmt_str += '{0}{1}'.format(
+            tmp,
+            self.target_transform.__repr__().replace('\n',
+                                                     '\n' + ' ' * len(tmp)))
+
+        return fmt_str
+
+
+IMG_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif']
+
+
+def pil_loader(path):
+    # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
+    if isinstance(path, bytes):
+        img = Image.open(io.BytesIO(path))
+    elif is_zip_path(path):
+        data = ZipReader.read(path)
+        img = Image.open(io.BytesIO(data))
+    else:
+        with open(path, 'rb') as f:
+            img = Image.open(f)
+            return img.convert('RGB')
+
+    return img.convert('RGB')
+
+
+def accimage_loader(path):
+    import accimage
+    try:
+        return accimage.Image(path)
+    except IOError:
+        # Potentially a decoding problem, fall back to PIL.Image
+        return pil_loader(path)
+
+
+def default_img_loader(path):
+    from torchvision import get_image_backend
+    if get_image_backend() == 'accimage':
+        return accimage_loader(path)
+    else:
+        return pil_loader(path)
+
+
+class CachedImageFolder(DatasetFolder):
+    """A generic data loader where the images are arranged in this way: ::
+        root/dog/xxx.png
+        root/dog/xxy.png
+        root/dog/xxz.png
+        root/cat/123.png
+        root/cat/nsdf3.png
+        root/cat/asd932_.png
+    Args:
+        root (string): Root directory path.
+        transform (callable, optional): A function/transform that  takes in an PIL image
+            and returns a transformed version. E.g, ``transforms.RandomCrop``
+        target_transform (callable, optional): A function/transform that takes in the
+            target and transforms it.
+        loader (callable, optional): A function to load an image given its path.
+     Attributes:
+        imgs (list): List of (image path, class_index) tuples
+    """
+
+    def __init__(self,
+                 root,
+                 ann_file='',
+                 img_prefix='',
+                 transform=None,
+                 target_transform=None,
+                 loader=default_img_loader,
+                 cache_mode="no"):
+        super(CachedImageFolder,
+              self).__init__(root,
+                             loader,
+                             IMG_EXTENSIONS,
+                             ann_file=ann_file,
+                             img_prefix=img_prefix,
+                             transform=transform,
+                             target_transform=target_transform,
+                             cache_mode=cache_mode)
+        self.imgs = self.samples
+
+    def __getitem__(self, index):
+        """
+        Args:
+            index (int): Index
+        Returns:
+            tuple: (image, target) where target is class_index of the target class.
+        """
+        path, target = self.samples[index]
+        image = self.loader(path)
+        if self.transform is not None:
+            img = self.transform(image)
+        else:
+            img = image
+        if self.target_transform is not None:
+            target = self.target_transform(target)
+
+        return img, target
+
+
+class ImageCephDataset(data.Dataset):
+
+    def __init__(self,
+                 root,
+                 split,
+                 parser=None,
+                 transform=None,
+                 target_transform=None,
+                 on_memory=False):
+        if '22k' in root:
+            # Imagenet 22k
+            annotation_root = 'meta/'
+        else:
+            # Imagenet
+            annotation_root = 'meta/'
+        if parser is None or isinstance(parser, str):
+            parser = ParserCephImage(root=root,
+                                     split=split,
+                                     annotation_root=annotation_root,
+                                     on_memory=on_memory)
+        self.parser = parser
+        self.transform = transform
+        self.target_transform = target_transform
+        self._consecutive_errors = 0
+
+    def __getitem__(self, index):
+        img, target = self.parser[index]
+        self._consecutive_errors = 0
+        if self.transform is not None:
+            img = self.transform(img)
+        if target is None:
+            target = -1
+        elif self.target_transform is not None:
+            target = self.target_transform(target)
+        return img, target
+
+    def __len__(self):
+        return len(self.parser)
+
+    def filename(self, index, basename=False, absolute=False):
+        return self.parser.filename(index, basename, absolute)
+
+    def filenames(self, basename=False, absolute=False):
+        return self.parser.filenames(basename, absolute)
+
+
+class Parser:
+
+    def __init__(self):
+        pass
+
+    @abstractmethod
+    def _filename(self, index, basename=False, absolute=False):
+        pass
+
+    def filename(self, index, basename=False, absolute=False):
+        return self._filename(index, basename=basename, absolute=absolute)
+
+    def filenames(self, basename=False, absolute=False):
+        return [
+            self._filename(index, basename=basename, absolute=absolute)
+            for index in range(len(self))
+        ]
+
+
+class ParserCephImage(Parser):
+
+    def __init__(self,
+                 root,
+                 split,
+                 annotation_root,
+                 on_memory=False,
+                 **kwargs):
+        super().__init__()
+
+        self.file_client = None
+        self.kwargs = kwargs
+
+        self.root = root  # dataset:s3://imagenet22k
+        if '22k' in root:
+            self.io_backend = 'petrel'
+            with open(osp.join(annotation_root, '22k_class_to_idx.json'),
+                      'r') as f:
+                self.class_to_idx = json.loads(f.read())
+            with open(osp.join(annotation_root, '22k_label.txt'), 'r') as f:
+                self.samples = f.read().splitlines()
+        else:
+            self.io_backend = 'disk'
+            self.class_to_idx = None
+            with open(osp.join(annotation_root, f'{split}.txt'), 'r') as f:
+                self.samples = f.read().splitlines()
+        local_rank = None
+        local_size = None
+        self._consecutive_errors = 0
+        self.on_memory = on_memory
+        if on_memory:
+            self.holder = {}
+            if local_rank is None:
+                local_rank = int(os.environ.get('LOCAL_RANK', 0))
+            if local_size is None:
+                local_size = int(os.environ.get('LOCAL_SIZE', 1))
+            self.local_rank = local_rank
+            self.local_size = local_size
+            self.rank = int(os.environ["RANK"])
+            self.world_size = int(os.environ['WORLD_SIZE'])
+            self.num_replicas = int(os.environ['WORLD_SIZE'])
+            self.num_parts = local_size
+            self.num_samples = int(
+                math.ceil(len(self.samples) * 1.0 / self.num_replicas))
+            self.total_size = self.num_samples * self.num_replicas
+            self.total_size_parts = self.num_samples * self.num_replicas // self.num_parts
+            self.load_onto_memory_v2()
+
+    def load_onto_memory(self):
+        print("Loading images onto memory...", self.local_rank,
+              self.local_size)
+        if self.file_client is None:
+            self.file_client = FileClient(self.io_backend, **self.kwargs)
+        for index in trange(len(self.samples)):
+            if index % self.local_size != self.local_rank:
+                continue
+            path, _ = self.samples[index].split(' ')
+            path = osp.join(self.root, path)
+            img_bytes = self.file_client.get(path)
+            self.holder[path] = img_bytes
+
+        print("Loading complete!")
+
+    def load_onto_memory_v2(self):
+        # print("Loading images onto memory...", self.local_rank, self.local_size)
+        t = torch.Generator()
+        t.manual_seed(0)
+        indices = torch.randperm(len(self.samples), generator=t).tolist()
+        # indices = range(len(self.samples))
+        indices = [i for i in indices if i % self.num_parts == self.local_rank]
+        # add extra samples to make it evenly divisible
+        indices += indices[:(self.total_size_parts - len(indices))]
+        assert len(indices) == self.total_size_parts
+
+        # subsample
+        indices = indices[self.rank // self.num_parts:self.
+                          total_size_parts:self.num_replicas // self.num_parts]
+        assert len(indices) == self.num_samples
+
+        if self.file_client is None:
+            self.file_client = FileClient(self.io_backend, **self.kwargs)
+        for index in tqdm(indices):
+            if index % self.local_size != self.local_rank:
+                continue
+            path, _ = self.samples[index].split(' ')
+            path = osp.join(self.root, path)
+            img_bytes = self.file_client.get(path)
+
+            self.holder[path] = img_bytes
+
+        print("Loading complete!")
+
+    def __getitem__(self, index):
+        if self.file_client is None:
+            self.file_client = FileClient(self.io_backend, **self.kwargs)
+
+        filepath, target = self.samples[index].split(' ')
+        filepath = osp.join(self.root, filepath)
+
+        try:
+            if self.on_memory:
+                img_bytes = self.holder[filepath]
+            else:
+                # pass
+                img_bytes = self.file_client.get(filepath)
+            img = mmcv.imfrombytes(img_bytes)[:, :, ::-1]
+        except Exception as e:
+            _logger.warning(
+                f'Skipped sample (index {index}, file {filepath}). {str(e)}')
+            self._consecutive_errors += 1
+            if self._consecutive_errors < _ERROR_RETRY:
+                return self.__getitem__((index + 1) % len(self))
+            else:
+                raise e
+        self._consecutive_errors = 0
+
+        img = Image.fromarray(img)
+        try:
+            if self.class_to_idx is not None:
+                target = self.class_to_idx[target]
+            else:
+                target = int(target)
+        except:
+            print('aaaaaaaaaaaa', filepath, target)
+            exit()
+
+        return img, target
+
+    def __len__(self):
+        return len(self.samples)
+
+    def _filename(self, index, basename=False, absolute=False):
+        filename, _ = self.samples[index].split(' ')
+        filename = osp.join(self.root, filename)
+
+        return filename
+
+
+def get_temporal_info(date, miss_hour=False):
+    try:
+        if date:
+            if miss_hour:
+                pattern = re.compile(r'(\d*)-(\d*)-(\d*)', re.I)
+            else:
+                pattern = re.compile(r'(\d*)-(\d*)-(\d*) (\d*):(\d*):(\d*)',
+                                     re.I)
+            m = pattern.match(date.strip())
+
+            if m:
+                year = int(m.group(1))
+                month = int(m.group(2))
+                day = int(m.group(3))
+                x_month = math.sin(2 * math.pi * month / 12)
+                y_month = math.cos(2 * math.pi * month / 12)
+                if miss_hour:
+                    x_hour = 0
+                    y_hour = 0
+                else:
+                    hour = int(m.group(4))
+                    x_hour = math.sin(2 * math.pi * hour / 24)
+                    y_hour = math.cos(2 * math.pi * hour / 24)
+                return [x_month, y_month, x_hour, y_hour]
+            else:
+                return [0, 0, 0, 0]
+        else:
+            return [0, 0, 0, 0]
+    except:
+        return [0, 0, 0, 0]
+
+
+def get_spatial_info(latitude, longitude):
+    if latitude and longitude:
+        latitude = math.radians(latitude)
+        longitude = math.radians(longitude)
+        x = math.cos(latitude) * math.cos(longitude)
+        y = math.cos(latitude) * math.sin(longitude)
+        z = math.sin(latitude)
+        return [x, y, z]
+    else:
+        return [0, 0, 0]

classification/dataset/samplers.py

@@ -0,0 +1,114 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+import torch
+import os
+import math
+from torch.utils.data.sampler import Sampler
+import torch.distributed as dist
+import numpy as np
+
+
+class SubsetRandomSampler(torch.utils.data.Sampler):
+    """Samples elements randomly from a given list of indices, without replacement.
+
+    Arguments:
+        indices (sequence): a sequence of indices
+    """
+
+    def __init__(self, indices):
+        self.epoch = 0
+        self.indices = indices
+
+    def __iter__(self):
+        return (self.indices[i] for i in torch.randperm(len(self.indices)))
+
+    def __len__(self):
+        return len(self.indices)
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch
+
+
+class NodeDistributedSampler(Sampler):
+    """Sampler that restricts data loading to a subset of the dataset.
+    It is especially useful in conjunction with
+    :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each
+    process can pass a DistributedSampler instance as a DataLoader sampler,
+    and load a subset of the original dataset that is exclusive to it.
+    .. note::
+        Dataset is assumed to be of constant size.
+    Arguments:
+        dataset: Dataset used for sampling.
+        num_replicas (optional): Number of processes participating in
+            distributed training.
+        rank (optional): Rank of the current process within num_replicas.
+    """
+
+    def __init__(self,
+                 dataset,
+                 num_replicas=None,
+                 rank=None,
+                 local_rank=None,
+                 local_size=None):
+        if num_replicas is None:
+            if not dist.is_available():
+                raise RuntimeError(
+                    "Requires distributed package to be available")
+            num_replicas = dist.get_world_size()
+        if rank is None:
+            if not dist.is_available():
+                raise RuntimeError(
+                    "Requires distributed package to be available")
+            rank = dist.get_rank()
+        if local_rank is None:
+            local_rank = int(os.environ.get('LOCAL_RANK', 0))
+        if local_size is None:
+            local_size = int(os.environ.get('LOCAL_SIZE', 1))
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.num_parts = local_size
+        self.rank = rank
+        self.local_rank = local_rank
+        self.epoch = 0
+        self.num_samples = int(
+            math.ceil(len(self.dataset) * 1.0 / self.num_replicas))
+        self.total_size = self.num_samples * self.num_replicas
+
+        self.total_size_parts = self.num_samples * self.num_replicas // self.num_parts
+
+    def __iter__(self):
+        # deterministically shuffle based on epoch
+        g = torch.Generator()
+        g.manual_seed(self.epoch)
+
+        t = torch.Generator()
+        t.manual_seed(0)
+
+        indices = torch.randperm(len(self.dataset), generator=t).tolist()
+        # indices = range(len(self.dataset))
+        indices = [i for i in indices if i % self.num_parts == self.local_rank]
+
+        # add extra samples to make it evenly divisible
+        indices += indices[:(self.total_size_parts - len(indices))]
+        assert len(indices) == self.total_size_parts
+
+        # subsample
+        indices = indices[self.rank // self.num_parts:self.
+                          total_size_parts:self.num_replicas // self.num_parts]
+
+        index = torch.randperm(len(indices), generator=g).tolist()
+        indices = list(np.array(indices)[index])
+
+        assert len(indices) == self.num_samples
+
+        return iter(indices)
+
+    def __len__(self):
+        return self.num_samples
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch

classification/dataset/zipreader.py

@@ -0,0 +1,102 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+import os
+import zipfile
+import io
+import numpy as np
+from PIL import Image
+from PIL import ImageFile
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+def is_zip_path(img_or_path):
+    """judge if this is a zip path"""
+    return '.zip@' in img_or_path
+
+
+class ZipReader(object):
+    """A class to read zipped files"""
+    zip_bank = dict()
+
+    def __init__(self):
+        super(ZipReader, self).__init__()
+
+    @staticmethod
+    def get_zipfile(path):
+        zip_bank = ZipReader.zip_bank
+        if path not in zip_bank:
+            zfile = zipfile.ZipFile(path, 'r')
+            zip_bank[path] = zfile
+        return zip_bank[path]
+
+    @staticmethod
+    def split_zip_style_path(path):
+        pos_at = path.index('@')
+        assert pos_at != -1, "character '@' is not found from the given path '%s'" % path
+
+        zip_path = path[0:pos_at]
+        folder_path = path[pos_at + 1:]
+        folder_path = str.strip(folder_path, '/')
+        return zip_path, folder_path
+
+    @staticmethod
+    def list_folder(path):
+        zip_path, folder_path = ZipReader.split_zip_style_path(path)
+
+        zfile = ZipReader.get_zipfile(zip_path)
+        folder_list = []
+        for file_foler_name in zfile.namelist():
+            file_foler_name = str.strip(file_foler_name, '/')
+            if file_foler_name.startswith(folder_path) and \
+                    len(os.path.splitext(file_foler_name)[-1]) == 0 and \
+                    file_foler_name != folder_path:
+                if len(folder_path) == 0:
+                    folder_list.append(file_foler_name)
+                else:
+                    folder_list.append(file_foler_name[len(folder_path) + 1:])
+
+        return folder_list
+
+    @staticmethod
+    def list_files(path, extension=None):
+        if extension is None:
+            extension = ['.*']
+        zip_path, folder_path = ZipReader.split_zip_style_path(path)
+
+        zfile = ZipReader.get_zipfile(zip_path)
+        file_lists = []
+        for file_foler_name in zfile.namelist():
+            file_foler_name = str.strip(file_foler_name, '/')
+            if file_foler_name.startswith(folder_path) and \
+                    str.lower(os.path.splitext(file_foler_name)[-1]) in extension:
+                if len(folder_path) == 0:
+                    file_lists.append(file_foler_name)
+                else:
+                    file_lists.append(file_foler_name[len(folder_path) + 1:])
+
+        return file_lists
+
+    @staticmethod
+    def read(path):
+        zip_path, path_img = ZipReader.split_zip_style_path(path)
+        zfile = ZipReader.get_zipfile(zip_path)
+        data = zfile.read(path_img)
+        return data
+
+    @staticmethod
+    def imread(path):
+        zip_path, path_img = ZipReader.split_zip_style_path(path)
+        zfile = ZipReader.get_zipfile(zip_path)
+        data = zfile.read(path_img)
+        try:
+            im = Image.open(io.BytesIO(data))
+        except:
+            print("ERROR IMG LOADED: ", path_img)
+            random_img = np.random.rand(224, 224, 3) * 255
+            im = Image.fromarray(np.uint8(random_img))
+        return im

classification/ddp_hooks.py

@@ -0,0 +1,183 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+from typing import Any, Callable
+
+import torch
+import torch.distributed as dist
+
+
+def _allreduce_fut(process_group: dist.ProcessGroup,
+                   tensor: torch.Tensor) -> torch.futures.Future[torch.Tensor]:
+    "Averages the input gradient tensor by allreduce and returns a future."
+    group_to_use = process_group if process_group is not None else dist.group.WORLD
+
+    # Apply the division first to avoid overflow, especially for FP16.
+    tensor.div_(group_to_use.size())
+
+    return (dist.all_reduce(
+        tensor, group=group_to_use,
+        async_op=True).get_future().then(lambda fut: fut.value()[0]))
+
+
+def allreduce_hook(
+        process_group: dist.ProcessGroup,
+        bucket: dist.GradBucket) -> torch.futures.Future[torch.Tensor]:
+    """
+    This DDP communication hook just calls ``allreduce`` using ``GradBucket``
+    tensors. Once gradient tensors are aggregated across all workers, its ``then``
+    callback takes the mean and returns the result. If user registers this hook,
+    DDP results is expected to be same as the case where no hook was registered.
+    Hence, this won't change behavior of DDP and user can use this as a reference
+    or modify this hook to log useful information or any other purposes while
+    unaffecting DDP behavior.
+
+    Example::
+        >>> ddp_model.register_comm_hook(process_group, allreduce_hook)
+    """
+    return _allreduce_fut(process_group, bucket.buffer())
+
+
+def fp16_compress_hook(
+        process_group: dist.ProcessGroup,
+        bucket: dist.GradBucket) -> torch.futures.Future[torch.Tensor]:
+    """
+    This DDP communication hook implements a simple gradient compression
+    approach that casts ``GradBucket`` tensor to half-precision floating-point format (``torch.float16``)
+    and then divides it by the process group size.
+    It allreduces those ``float16`` gradient tensors. Once compressed gradient
+    tensors are allreduced, the chained callback ``decompress`` casts it back to the input data type (such as ``float32``).
+
+    Example::
+        >>> ddp_model.register_comm_hook(process_group, fp16_compress_hook)
+    """
+    group_to_use = process_group if process_group is not None else dist.group.WORLD
+    world_size = group_to_use.size()
+
+    compressed_tensor = bucket.buffer().to(torch.float16).div_(world_size)
+
+    fut = dist.all_reduce(compressed_tensor, group=group_to_use,
+                          async_op=True).get_future()
+
+    def decompress(fut):
+        decompressed_tensor = bucket.buffer()
+        # Decompress in place to reduce the peak memory.
+        # See: https://github.com/pytorch/pytorch/issues/45968
+        decompressed_tensor.copy_(fut.value()[0])
+        return decompressed_tensor
+
+    return fut.then(decompress)
+
+
+# TODO: create an internal helper function and extract the duplicate code in FP16_compress and BF16_compress.
+
+
+def bf16_compress_hook(
+        process_group: dist.ProcessGroup,
+        bucket: dist.GradBucket) -> torch.futures.Future[torch.Tensor]:
+    """
+    Warning: This API is experimental, and it requires NCCL version later than 2.9.6.
+
+    This DDP communication hook implements a simple gradient compression
+    approach that casts ``GradBucket`` tensor to half-precision
+    `Brain floating point format <https://en.wikipedia.org/wiki/Bfloat16_floating-point_format>`_ (``torch.bfloat16``)
+    and then divides it by the process group size.
+    It allreduces those ``bfloat16`` gradient tensors. Once compressed gradient
+    tensors are allreduced, the chained callback ``decompress`` casts it back to the input data type (such as ``float32``).
+
+    Example::
+        >>> ddp_model.register_comm_hook(process_group, bf16_compress_hook)
+    """
+    group_to_use = process_group if process_group is not None else dist.group.WORLD
+    world_size = group_to_use.size()
+
+    compressed_tensor = bucket.buffer().to(torch.bfloat16).div_(world_size)
+
+    fut = dist.all_reduce(compressed_tensor, group=group_to_use,
+                          async_op=True).get_future()
+
+    def decompress(fut):
+        decompressed_tensor = bucket.buffer()
+        # Decompress in place to reduce the peak memory.
+        # See: https://github.com/pytorch/pytorch/issues/45968
+        decompressed_tensor.copy_(fut.value()[0])
+        return decompressed_tensor
+
+    return fut.then(decompress)
+
+
+def fp16_compress_wrapper(
+    hook: Callable[[Any, dist.GradBucket], torch.futures.Future[torch.Tensor]]
+) -> Callable[[Any, dist.GradBucket], torch.futures.Future[torch.Tensor]]:
+    """
+    This wrapper casts the input gradient tensor of a given DDP communication hook to half-precision
+    floating point format (``torch.float16``), and casts the resulting tensor of the given hook back to
+    the input data type, such as ``float32``.
+
+    Therefore, ``fp16_compress_hook`` is equivalent to ``fp16_compress_wrapper(allreduce_hook)``.
+
+    Example::
+        >>> state = PowerSGDState(process_group=process_group, matrix_approximation_rank=1, start_powerSGD_iter=10)
+        >>> ddp_model.register_comm_hook(state, fp16_compress_wrapper(powerSGD_hook))
+    """
+
+    def fp16_compress_wrapper_hook(
+            hook_state,
+            bucket: dist.GradBucket) -> torch.futures.Future[torch.Tensor]:
+        # Cast bucket tensor to FP16.
+        bucket.set_buffer(bucket.buffer().to(torch.float16))
+
+        fut = hook(hook_state, bucket)
+
+        def decompress(fut):
+            decompressed_tensor = bucket.buffer()
+            # Decompress in place to reduce the peak memory.
+            # See: https://github.com/pytorch/pytorch/issues/45968
+            decompressed_tensor.copy_(fut.value())
+            return decompressed_tensor
+
+        # Decompress after hook has run.
+        return fut.then(decompress)
+
+    return fp16_compress_wrapper_hook
+
+
+def bf16_compress_wrapper(
+    hook: Callable[[Any, dist.GradBucket], torch.futures.Future[torch.Tensor]]
+) -> Callable[[Any, dist.GradBucket], torch.futures.Future[torch.Tensor]]:
+    """
+    Warning: This API is experimental, and it requires NCCL version later than 2.9.6.
+
+    This wrapper casts the input gradient tensor of a given DDP communication hook to half-precision
+    `Brain floating point format <https://en.wikipedia.org/wiki/Bfloat16_floating-point_format> `_  (``torch.bfloat16``),
+    and casts the resulting tensor of the given hook back to the input data type, such as ``float32``.
+
+    Therefore, ``bf16_compress_hook`` is equivalent to ``bf16_compress_wrapper(allreduce_hook)``.
+
+    Example::
+        >>> state = PowerSGDState(process_group=process_group, matrix_approximation_rank=1, start_powerSGD_iter=10)
+        >>> ddp_model.register_comm_hook(state, bf16_compress_wrapper(powerSGD_hook))
+    """
+
+    def bf16_compress_wrapper_hook(
+            hook_state,
+            bucket: dist.GradBucket) -> torch.futures.Future[torch.Tensor]:
+        # Cast bucket tensor to BF16.
+        bucket.set_buffer(bucket.buffer().to(torch.bfloat16))
+
+        fut = hook(hook_state, bucket)
+
+        def decompress(fut):
+            decompressed_tensor = bucket.buffer()
+            # Decompress in place to reduce the peak memory.
+            # See: https://github.com/pytorch/pytorch/issues/45968
+            decompressed_tensor.copy_(fut.value())
+            return decompressed_tensor
+
+        # Decompress after hook has run.
+        return fut.then(decompress)
+
+    return bf16_compress_wrapper_hook

classification/ema_deepspeed.py

@@ -0,0 +1,99 @@
+import torch
+import torch.nn as nn
+import deepspeed
+from deepspeed.runtime.zero import GatheredParameters
+from contextlib import contextmanager
+
+
+class EMADeepspeed(nn.Module):
+    """ migrated from https://github.com/microsoft/DeepSpeed/issues/2056
+    """
+
+    def __init__(self, model, decay=0.9999, use_num_updates=True):
+        super().__init__()
+        if decay < 0.0 or decay > 1.0:
+            raise ValueError('Decay must be between 0 and 1')
+
+        self.m_name2s_name = {}
+        self.decay = decay
+        self.num_updates = 0 if use_num_updates else -1
+
+        with GatheredParameters(model.parameters(), fwd_module=self):
+            for name, p in model.named_parameters():
+                if p.requires_grad:
+                    # remove as '.'-character is not allowed in buffers
+                    s_name = name.replace('.', '')
+                    self.m_name2s_name.update({name: s_name})
+                    self.register_buffer(s_name, p.clone().detach().data)
+                    # remove as '.'-character is not allowed in buffers
+        self.collected_params = []
+
+    def forward(self, model):
+        decay = self.decay
+
+        if self.num_updates >= 0:
+            self.num_updates += 1
+            decay = min(self.decay, (1 + self.num_updates) / (10 + self.num_updates))
+
+        one_minus_decay = 1.0 - decay
+        shadow_params = dict(self.named_buffers())
+
+        with torch.no_grad():
+            with GatheredParameters(model.parameters()):
+                if deepspeed.comm.get_rank() == 0:
+                    m_param = dict(model.named_parameters())
+
+                    for key in m_param:
+                        if m_param[key].requires_grad:
+                            sname = self.m_name2s_name[key]
+                            shadow_params[sname] = shadow_params[sname].type_as(m_param[key])
+                            shadow_params[sname].sub_(one_minus_decay * (shadow_params[sname] - m_param[key]))
+                        else:
+                            assert not key in self.m_name2s_name
+
+    def copy_to(self, model):
+        shadow_params = dict(self.named_buffers())
+        with GatheredParameters(model.parameters(), modifier_rank=0):
+            if deepspeed.comm.get_rank() == 0:
+                m_param = dict(model.named_parameters())
+                for key in m_param:
+                    if m_param[key].requires_grad:
+                        m_param[key].data.copy_(shadow_params[self.m_name2s_name[key]].data)
+                    else:
+                        assert not key in self.m_name2s_name
+
+    def store(self, model):
+        """
+        Save the current parameters for restoring later.
+        Args:
+          model: A model that parameters will be stored
+        """
+        with GatheredParameters(model.parameters()):
+            if deepspeed.comm.get_rank() == 0:
+                parameters = model.parameters()
+                self.collected_params = [param.clone() for param in parameters]
+
+    def restore(self, model):
+        """
+        Restore the parameters stored with the `store` method.
+        Useful to validate the model with EMA parameters without affecting the
+        original optimization process. Store the parameters before the
+        `copy_to` method. After validation (or model saving), use this to
+        restore the former parameters.
+        Args:
+          model: A model that to restore its parameters.
+        """
+        with GatheredParameters(model.parameters(), modifier_rank=0):
+            if deepspeed.comm.get_rank() == 0:
+                parameters = model.parameters()
+                for c_param, param in zip(self.collected_params, parameters):
+                    param.data.copy_(c_param.data)
+
+    @contextmanager
+    def activate(self, model):
+        try:
+            self.store(model)
+            self.copy_to(model)
+            yield
+        finally:
+            self.restore(model)

classification/eval.sh

@@ -0,0 +1,2 @@
+python -m torch.distributed.launch --nproc_per_node 1 --master_port 12345 main.py --eval \
+--cfg configs/flash_intern_image_l_22k_384.yaml  --data-path /path/to/imagenet1k

classification/export.py

@@ -0,0 +1,122 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+
+import os
+import time
+import argparse
+
+import torch
+from tqdm import tqdm
+
+from config import get_config
+from models import build_model
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model_name', type=str,
+                        default='internimage_t_1k_224')
+    parser.add_argument('--ckpt_dir', type=str,
+                        default='/mnt/petrelfs/share_data/huangzhenhang/code/internimage/checkpoint_dir/new/cls')
+    parser.add_argument('--onnx', default=False, action='store_true')
+    parser.add_argument('--trt', default=False, action='store_true')
+
+    args = parser.parse_args()
+    args.cfg = os.path.join('./configs', f'{args.model_name}.yaml')
+    args.ckpt = os.path.join(args.ckpt_dir, f'{args.model_name}.pth')
+    args.size = int(args.model_name.split('.')[0].split('_')[-1])
+
+    cfg = get_config(args)
+    return args, cfg
+
+def get_model(args, cfg):
+    model = build_model(cfg)
+    ckpt = torch.load(args.ckpt, map_location='cpu')['model']
+
+    model.load_state_dict(ckpt)
+    return model
+
+def speed_test(model, input):
+    # warmup
+    for _ in tqdm(range(100)):
+        _ = model(input)
+
+    # speed test
+    torch.cuda.synchronize()
+    start = time.time()
+    for _ in tqdm(range(100)):
+        _ = model(input)
+    end = time.time()
+    th = 100 / (end - start)
+    print(f"using time: {end - start}, throughput {th}")
+
+def torch2onnx(args, cfg):
+    model = get_model(args, cfg).cuda()
+
+    # speed_test(model)
+
+    onnx_name = f'{args.model_name}.onnx'
+    torch.onnx.export(model,
+                      torch.rand(1, 3, args.size, args.size).cuda(),
+                      onnx_name,
+                      input_names=['input'],
+                      output_names=['output'])
+
+    return model
+
+def onnx2trt(args):
+    from mmdeploy.backend.tensorrt import from_onnx
+
+    onnx_name = f'{args.model_name}.onnx'
+    from_onnx(
+        onnx_name,
+        args.model_name,
+        dict(
+            input=dict(
+                min_shape=[1, 3, args.size, args.size],
+                opt_shape=[1, 3, args.size, args.size],
+                max_shape=[1, 3, args.size, args.size],
+            )
+        ),
+        max_workspace_size=2**30,
+    )
+
+def check(args, cfg):
+    from mmdeploy.backend.tensorrt.wrapper import TRTWrapper
+
+    model = get_model(args, cfg).cuda()
+    model.eval()
+    trt_model = TRTWrapper(f'{args.model_name}.engine',
+                           ['output'])
+
+    x = torch.randn(1, 3, args.size, args.size).cuda()
+
+    torch_out = model(x)
+    trt_out = trt_model(dict(input=x))['output']
+
+    print('torch out shape:', torch_out.shape)
+    print('trt out shape:', trt_out.shape)
+
+    print('max delta:', (torch_out - trt_out).abs().max())
+    print('mean delta:', (torch_out - trt_out).abs().mean())
+
+    speed_test(model, x)
+    speed_test(trt_model, dict(input=x))
+
+def main():
+    args, cfg = get_args()
+
+    if args.onnx or args.trt:
+        torch2onnx(args, cfg)
+        print('torch -> onnx: succeess')
+
+    if args.trt:
+        onnx2trt(args)
+        print('onnx -> trt: success')
+        check(args, cfg)
+
+if __name__ == '__main__':
+    main()

classification/extract_feature.py

@@ -0,0 +1,128 @@
+import functools
+from collections import OrderedDict
+
+
+# using wonder's beautiful simplification:
+# https://stackoverflow.com/questions/31174295/getattr-and-setattr-on-nested-objects/31174427?noredirect=1#comment86638618_31174427
+def rgetattr(obj, attr, *args):
+    def _getattr(obj, attr):
+        return getattr(obj, attr, *args)
+
+    return functools.reduce(_getattr, [obj] + attr.split('.'))
+
+
+class IntermediateLayerGetter:
+    def __init__(self, model, return_layers, keep_output=True):
+        """Wraps a Pytorch module to get intermediate values
+
+        Arguments:
+            model {nn.module} -- The Pytorch module to call
+            return_layers {dict} -- Dictionary with the selected submodules
+            to return the output (format: {[current_module_name]: [desired_output_name]},
+            current_module_name can be a nested submodule, e.g. submodule1.submodule2.submodule3)
+
+        Keyword Arguments:
+            keep_output {bool} -- If True model_output contains the final model's output
+            in the other case model_output is None (default: {True})
+
+        Returns:
+            (mid_outputs {OrderedDict}, model_output {any}) -- mid_outputs keys are
+            your desired_output_name (s) and their values are the returned tensors
+            of those submodules (OrderedDict([(desired_output_name,tensor(...)), ...).
+            See keep_output argument for model_output description.
+            In case a submodule is called more than one time, all it's outputs are
+            stored in a list.
+        """
+        self._model = model
+        self.return_layers = return_layers
+        self.keep_output = keep_output
+
+    def __call__(self, *args, **kwargs):
+        ret = OrderedDict()
+        handles = []
+        for name, new_name in self.return_layers.items():
+            layer = rgetattr(self._model, name)
+
+            def hook(module, input, output, new_name=new_name):
+                if new_name in ret:
+                    if type(ret[new_name]) is list:
+                        ret[new_name].append(output)
+                    else:
+                        ret[new_name] = [ret[new_name], output]
+                else:
+                    ret[new_name] = output
+
+            try:
+                h = layer.register_forward_hook(hook)
+            except AttributeError as e:
+                raise AttributeError(f'Module {name} not found')
+            handles.append(h)
+
+        if self.keep_output:
+            output = self._model(*args, **kwargs)
+        else:
+            self._model(*args, **kwargs)
+            output = None
+
+        for h in handles:
+            h.remove()
+
+        return ret, output
+
+
+def main(args, config):
+    from models import build_model
+    import torchvision.transforms as T
+    from PIL import Image
+
+    model = build_model(config)
+    checkpoint = torch.load(config.MODEL.RESUME, map_location='cpu')
+    model.load_state_dict(checkpoint['model'], strict=False)
+    model.cuda()
+
+    # examples:
+    # return_layers = {
+    #     'patch_embed': 'patch_embed',
+    #     'levels.0.downsample': 'levels.0.downsample',
+    #     'levels.0.blocks.0.dcn': 'levels.0.blocks.0.dcn',
+    # }
+    return_layers = {k: k for k in args.keys}
+    mid_getter = IntermediateLayerGetter(model, return_layers=return_layers, keep_output=True)
+
+    image = Image.open(args.img)
+
+    transforms = T.Compose([
+        T.Resize(config.DATA.IMG_SIZE),
+        T.ToTensor(),
+        T.Normalize(config.AUG.MEAN, config.AUG.STD)
+    ])
+    image = transforms(image)
+    image = image.unsqueeze(0)
+    image = image.cuda()
+
+    mid_outputs, model_output = mid_getter(image)
+
+    for k, v in mid_outputs.items():
+        print(k, v.shape)
+
+    return mid_outputs, model_output
+
+
+if __name__ == '__main__':
+    import argparse
+    import torch
+    from config import get_config
+
+    parser = argparse.ArgumentParser('Get Intermediate Layer Output')
+    parser.add_argument('--cfg', type=str, required=True, metavar="FILE", help='Path to config file')
+    parser.add_argument('--img', type=str, required=True, metavar="FILE", help='Path to img file')
+    parser.add_argument("--keys", default=None, nargs='+', help="The intermediate layer's keys you want to save.")
+    parser.add_argument('--resume', help='resume from checkpoint')
+    parser.add_argument('--save', action='store_true', help='Save the results.')
+    args = parser.parse_args()
+    config = get_config(args)
+
+    mid_outputs, model_output = main(args, config)
+
+    if args.save:
+        torch.save(mid_outputs, args.img[:-3] + '.pth')

classification/logger.py

@@ -0,0 +1,44 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+import os
+import sys
+import logging
+import functools
+from termcolor import colored
+
+
+@functools.lru_cache()
+def create_logger(output_dir, dist_rank=0, name=''):
+    # create logger
+    logger = logging.getLogger(name)
+    logger.setLevel(logging.DEBUG)
+    logger.propagate = False
+
+    # create formatter
+    fmt = '[%(asctime)s %(name)s] (%(filename)s %(lineno)d): %(levelname)s %(message)s'
+    color_fmt = colored('[%(asctime)s %(name)s]', 'green') + \
+        colored('(%(filename)s %(lineno)d)', 'yellow') + \
+        ': %(levelname)s %(message)s'
+
+    # create console handlers for master process
+    if dist_rank == 0:
+        console_handler = logging.StreamHandler(sys.stdout)
+        console_handler.setLevel(logging.DEBUG)
+        console_handler.setFormatter(
+            logging.Formatter(fmt=color_fmt, datefmt='%Y-%m-%d %H:%M:%S'))
+        logger.addHandler(console_handler)
+
+    # create file handlers
+    file_handler = logging.FileHandler(os.path.join(
+        output_dir, f'log_rank{dist_rank}.txt'),
+                                       mode='a')
+    file_handler.setLevel(logging.DEBUG)
+    file_handler.setFormatter(
+        logging.Formatter(fmt=fmt, datefmt='%Y-%m-%d %H:%M:%S'))
+    logger.addHandler(file_handler)
+
+    return logger

classification/lr_scheduler.py

@@ -0,0 +1,112 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+
+import torch
+from timm.scheduler.cosine_lr import CosineLRScheduler
+from timm.scheduler.step_lr import StepLRScheduler
+from timm.scheduler.scheduler import Scheduler
+
+
+def build_scheduler(config, optimizer, n_iter_per_epoch):
+    num_steps = int(config.TRAIN.EPOCHS * n_iter_per_epoch)
+    warmup_steps = int(config.TRAIN.WARMUP_EPOCHS * n_iter_per_epoch)
+    decay_steps = int(config.TRAIN.LR_SCHEDULER.DECAY_EPOCHS *
+                      n_iter_per_epoch)
+
+    lr_scheduler = None
+    if config.TRAIN.LR_SCHEDULER.NAME == 'cosine':
+        lr_scheduler = CosineLRScheduler(
+            optimizer,
+            t_initial=num_steps,
+            # t_mul=1.,
+            lr_min=config.TRAIN.MIN_LR,
+            warmup_lr_init=config.TRAIN.WARMUP_LR,
+            warmup_t=warmup_steps,
+            cycle_limit=1,
+            t_in_epochs=False,
+        )
+    elif config.TRAIN.LR_SCHEDULER.NAME == 'linear':
+        lr_scheduler = LinearLRScheduler(
+            optimizer,
+            t_initial=num_steps,
+            lr_min_rate=0.01,
+            warmup_lr_init=config.TRAIN.WARMUP_LR,
+            warmup_t=warmup_steps,
+            t_in_epochs=False,
+        )
+    elif config.TRAIN.LR_SCHEDULER.NAME == 'step':
+        lr_scheduler = StepLRScheduler(
+            optimizer,
+            decay_t=decay_steps,
+            decay_rate=config.TRAIN.LR_SCHEDULER.DECAY_RATE,
+            warmup_lr_init=config.TRAIN.WARMUP_LR,
+            warmup_t=warmup_steps,
+            t_in_epochs=False,
+        )
+
+    return lr_scheduler
+
+
+class LinearLRScheduler(Scheduler):
+
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        t_initial: int,
+        lr_min_rate: float,
+        warmup_t=0,
+        warmup_lr_init=0.,
+        t_in_epochs=True,
+        noise_range_t=None,
+        noise_pct=0.67,
+        noise_std=1.0,
+        noise_seed=42,
+        initialize=True,
+    ) -> None:
+        super().__init__(optimizer,
+                         param_group_field="lr",
+                         noise_range_t=noise_range_t,
+                         noise_pct=noise_pct,
+                         noise_std=noise_std,
+                         noise_seed=noise_seed,
+                         initialize=initialize)
+
+        self.t_initial = t_initial
+        self.lr_min_rate = lr_min_rate
+        self.warmup_t = warmup_t
+        self.warmup_lr_init = warmup_lr_init
+        self.t_in_epochs = t_in_epochs
+        if self.warmup_t:
+            self.warmup_steps = [(v - warmup_lr_init) / self.warmup_t
+                                 for v in self.base_values]
+            super().update_groups(self.warmup_lr_init)
+        else:
+            self.warmup_steps = [1 for _ in self.base_values]
+
+    def _get_lr(self, t):
+        if t < self.warmup_t:
+            lrs = [self.warmup_lr_init + t * s for s in self.warmup_steps]
+        else:
+            t = t - self.warmup_t
+            total_t = self.t_initial - self.warmup_t
+            lrs = [
+                v - ((v - v * self.lr_min_rate) * (t / total_t))
+                for v in self.base_values
+            ]
+        return lrs
+
+    def get_epoch_values(self, epoch: int):
+        if self.t_in_epochs:
+            return self._get_lr(epoch)
+        else:
+            return None
+
+    def get_update_values(self, num_updates: int):
+        if not self.t_in_epochs:
+            return self._get_lr(num_updates)
+        else:
+            return None

classification/main.py

@@ -0,0 +1,671 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+
+import os
+import time
+import random
+import argparse
+import datetime
+import numpy as np
+import subprocess
+
+import torch
+import torch.backends.cudnn as cudnn
+import torch.distributed as dist
+from timm.utils import ModelEma, ApexScaler
+from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy
+from timm.utils import accuracy, AverageMeter
+
+from config import get_config
+from models import build_model
+from dataset import build_loader
+from lr_scheduler import build_scheduler
+from optimizer import build_optimizer
+from logger import create_logger
+from utils import NativeScalerWithGradNormCount as NativeScaler
+from utils import (load_checkpoint, load_pretrained, save_checkpoint,
+                   get_grad_norm, auto_resume_helper, reduce_tensor,
+                   load_ema_checkpoint, MyAverageMeter)
+
+from contextlib import suppress
+from ddp_hooks import fp16_compress_hook
+
+try:
+    from apex import amp
+    has_apex = True
+except ImportError:
+    has_apex = False
+# assert not has_apex, "The code is modified based on native amp"
+
+has_native_amp = False
+try:
+    if getattr(torch.cuda.amp, 'autocast') is not None:
+        has_native_amp = True
+except AttributeError:
+    pass
+
+TORCH_VERSION = tuple(int(x) for x in torch.__version__.split('.')[:2])
+
+
+def obsolete_torch_version(torch_version, version_threshold):
+    return torch_version == 'parrots' or torch_version <= version_threshold
+
+
+def parse_option():
+    parser = argparse.ArgumentParser(
+        'InternImage training and evaluation script', add_help=False)
+    parser.add_argument('--cfg',
+                        type=str,
+                        required=True,
+                        metavar="FILE",
+                        help='path to config file')
+    parser.add_argument(
+        "--opts",
+        help="Modify config options by adding 'KEY VALUE' pairs. ",
+        default=None,
+        nargs='+')
+
+    # easy config modification
+    parser.add_argument('--batch-size',
+                        type=int,
+                        help="batch size for single GPU")
+    parser.add_argument('--dataset',
+                        type=str,
+                        help='dataset name',
+                        default=None)
+    parser.add_argument('--data-path', type=str, help='path to dataset')
+    parser.add_argument('--zip',
+                        action='store_true',
+                        help='use zipped dataset instead of folder dataset')
+    parser.add_argument(
+        '--cache-mode',
+        type=str,
+        default='part',
+        choices=['no', 'full', 'part'],
+        help='no: no cache, '
+        'full: cache all data, '
+        'part: sharding the dataset into nonoverlapping pieces and only cache one piece'
+    )
+    parser.add_argument(
+        '--pretrained',
+        help=
+        'pretrained weight from checkpoint, could be imagenet22k pretrained weight'
+    )
+    parser.add_argument('--resume', help='resume from checkpoint')
+    parser.add_argument('--accumulation-steps',
+                        type=int,
+                        default=1,
+                        help="gradient accumulation steps")
+    parser.add_argument(
+        '--use-checkpoint',
+        action='store_true',
+        help="whether to use gradient checkpointing to save memory")
+    parser.add_argument(
+        '--amp-opt-level',
+        type=str,
+        default='O1',
+        choices=['O0', 'O1', 'O2'],
+        help='mixed precision opt level, if O0, no amp is used')
+    parser.add_argument(
+        '--output',
+        default='output',
+        type=str,
+        metavar='PATH',
+        help=
+        'root of output folder, the full path is <output>/<model_name>/<tag> (default: output)'
+    )
+    parser.add_argument('--tag', help='tag of experiment')
+    parser.add_argument('--eval',
+                        action='store_true',
+                        help='Perform evaluation only')
+    parser.add_argument('--throughput',
+                        action='store_true',
+                        help='Test throughput only')
+    parser.add_argument('--save-ckpt-num', default=1, type=int)
+    parser.add_argument(
+        '--use-zero',
+        action='store_true',
+        help="whether to use ZeroRedundancyOptimizer (ZeRO) to save memory")
+
+    # distributed training
+    parser.add_argument("--local-rank",
+                        type=int,
+                        default=0,
+                        help='local rank for DistributedDataParallel')
+
+    args, unparsed = parser.parse_known_args()
+  
+    if 'LOCAL_RANK' not in os.environ:
+        os.environ['LOCAL_RANK'] = str(args.local_rank)
+
+
+    config = get_config(args)
+
+    config.defrost()
+    config.LOCAL_RANK = int(os.environ['LOCAL_RANK'])
+    config.freeze()
+    return args, config
+
+
+@torch.no_grad()
+def throughput(data_loader, model, logger):
+    model.eval()
+
+    for idx, (images, _) in enumerate(data_loader):
+        images = images.cuda(non_blocking=True)
+        batch_size = images.shape[0]
+        for i in range(50):
+            model(images)
+        torch.cuda.synchronize()
+        logger.info(f"throughput averaged with 30 times")
+        tic1 = time.time()
+        for i in range(30):
+            model(images)
+        torch.cuda.synchronize()
+        tic2 = time.time()
+        logger.info(
+            f"batch_size {batch_size} throughput {30 * batch_size / (tic2 - tic1)}"
+        )
+        return
+
+
+def main(config):
+    # prepare data loaders
+    dataset_train, dataset_val, dataset_test, data_loader_train, \
+        data_loader_val, data_loader_test, mixup_fn = build_loader(config)
+
+    # build runner
+    logger.info(f"Creating model:{config.MODEL.TYPE}/{config.MODEL.NAME}")
+    model = build_model(config)
+    model.cuda()
+    logger.info(str(model))
+
+    # build optimizer
+    optimizer = build_optimizer(config, model)
+
+    if config.AMP_OPT_LEVEL != "O0":
+        config.defrost()
+        if has_native_amp:
+            config.native_amp = True
+            use_amp = 'native'
+        elif has_apex:
+            config.apex_amp = True
+            use_amp = 'apex'
+        else:
+            use_amp = None
+            logger.warning(
+                "Neither APEX or native Torch AMP is available, using float32. "
+                "Install NVIDA apex or upgrade to PyTorch 1.6")
+        config.freeze()
+
+    # setup automatic mixed-precision (AMP) loss scaling and op casting
+    amp_autocast = suppress  # do nothing
+    loss_scaler = None
+    if config.AMP_OPT_LEVEL != "O0":
+        if use_amp == 'apex':
+            model, optimizer = amp.initialize(model,
+                                              optimizer,
+                                              opt_level=config.AMP_OPT_LEVEL)
+            loss_scaler = ApexScaler()
+            if config.LOCAL_RANK == 0:
+                logger.info(
+                    'Using NVIDIA APEX AMP. Training in mixed precision.')
+        if use_amp == 'native':
+            amp_autocast = torch.cuda.amp.autocast
+            loss_scaler = NativeScaler()
+            if config.LOCAL_RANK == 0:
+                logger.info(
+                    'Using native Torch AMP. Training in mixed precision.')
+        else:
+            if config.LOCAL_RANK == 0:
+                logger.info('AMP not enabled. Training in float32.')
+
+    # put model on gpus
+    model = torch.nn.parallel.DistributedDataParallel(
+        model, device_ids=[config.LOCAL_RANK], broadcast_buffers=False)
+
+    try:
+        model.register_comm_hook(state=None, hook=fp16_compress_hook)
+        logger.info('using fp16_compress_hook!')
+    except:
+        logger.info("cannot register fp16_compress_hook!")
+
+    model_without_ddp = model.module
+
+    n_parameters = sum(p.numel() for p in model.parameters()
+                       if p.requires_grad)
+    logger.info(f"number of params: {n_parameters}")
+    if hasattr(model_without_ddp, 'flops'):
+        flops = model_without_ddp.flops()
+        logger.info(f"number of GFLOPs: {flops / 1e9}")
+
+    # build learning rate scheduler
+    lr_scheduler = build_scheduler(config, optimizer, len(data_loader_train)) \
+        if not config.EVAL_MODE else None
+
+    # build criterion
+    if config.AUG.MIXUP > 0.:
+        # smoothing is handled with mixup label transform
+        criterion = SoftTargetCrossEntropy()
+    elif config.MODEL.LABEL_SMOOTHING > 0.:
+        criterion = LabelSmoothingCrossEntropy(
+            smoothing=config.MODEL.LABEL_SMOOTHING)
+    else:
+        criterion = torch.nn.CrossEntropyLoss()
+
+    max_accuracy = 0.0
+    max_ema_accuracy = 0.0
+    # set auto resume
+    if config.MODEL.RESUME == '' and config.TRAIN.AUTO_RESUME:
+        resume_file = auto_resume_helper(config.OUTPUT)
+        if resume_file:
+            if config.MODEL.RESUME:
+                logger.warning(
+                    f"auto-resume changing resume file from {config.MODEL.RESUME} to {resume_file}"
+                )
+            config.defrost()
+            config.MODEL.RESUME = resume_file
+            config.freeze()
+            logger.info(f'auto resuming from {resume_file}')
+        else:
+            logger.info(
+                f'no checkpoint found in {config.OUTPUT}, ignoring auto resume'
+            )
+
+    # set resume and pretrain
+    if config.MODEL.RESUME:
+        max_accuracy = load_checkpoint(config, model_without_ddp, optimizer,
+                                       lr_scheduler, loss_scaler, logger)
+        if data_loader_val is not None:
+            acc1, acc5, loss = validate(config, data_loader_val, model, amp_autocast=amp_autocast)
+            logger.info(
+                f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%"
+            )
+    elif config.MODEL.PRETRAINED:
+        load_pretrained(config, model_without_ddp, logger)
+        if data_loader_val is not None:
+            acc1, acc5, loss = validate(config, data_loader_val, model, amp_autocast=amp_autocast)
+            logger.info(
+                f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%"
+            )
+
+    # evaluate EMA
+    model_ema = None
+    if config.TRAIN.EMA.ENABLE:
+        # Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
+        model_ema = ModelEma(model, decay=config.TRAIN.EMA.DECAY)
+        print("Using EMA with decay = %.8f" % config.TRAIN.EMA.DECAY)
+        if config.MODEL.RESUME:
+            load_ema_checkpoint(config, model_ema, logger)
+            acc1, acc5, loss = validate(config, data_loader_val, model_ema.ema, amp_autocast=amp_autocast)
+            logger.info(
+                f"Accuracy of the ema network on the {len(dataset_val)} test images: {acc1:.1f}%"
+            )
+
+    if config.THROUGHPUT_MODE:
+        throughput(data_loader_val, model, logger)
+
+    if config.EVAL_MODE:
+        return
+
+    # train
+    logger.info("Start training")
+    start_time = time.time()
+    for epoch in range(config.TRAIN.START_EPOCH, config.TRAIN.EPOCHS):
+        data_loader_train.sampler.set_epoch(epoch)
+
+        train_one_epoch(config,
+                        model,
+                        criterion,
+                        data_loader_train,
+                        optimizer,
+                        epoch,
+                        mixup_fn,
+                        lr_scheduler,
+                        amp_autocast,
+                        loss_scaler,
+                        model_ema=model_ema)
+        if (epoch % config.SAVE_FREQ == 0 or epoch == (config.TRAIN.EPOCHS - 1)) and \
+                config.TRAIN.OPTIMIZER.USE_ZERO:
+            optimizer.consolidate_state_dict(to=0)
+        if dist.get_rank() == 0 and (epoch % config.SAVE_FREQ == 0
+                                     or epoch == (config.TRAIN.EPOCHS - 1)):
+            save_checkpoint(config,
+                            epoch,
+                            model_without_ddp,
+                            max_accuracy,
+                            optimizer,
+                            lr_scheduler,
+                            loss_scaler,
+                            logger,
+                            model_ema=model_ema)
+        if data_loader_val is not None and epoch % config.EVAL_FREQ == 0:
+            acc1, acc5, loss = validate(config, data_loader_val, model, epoch, amp_autocast)
+            logger.info(
+                f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%"
+            )
+            if dist.get_rank() == 0 and acc1 > max_accuracy:
+                save_checkpoint(config,
+                                epoch,
+                                model_without_ddp,
+                                max_accuracy,
+                                optimizer,
+                                lr_scheduler,
+                                loss_scaler,
+                                logger,
+                                model_ema=model_ema,
+                                best='best')
+            max_accuracy = max(max_accuracy, acc1)
+            logger.info(f'Max accuracy: {max_accuracy:.2f}%')
+
+            if config.TRAIN.EMA.ENABLE:
+                acc1, acc5, loss = validate(config, data_loader_val,
+                                            model_ema.ema, epoch, amp_autocast)
+                logger.info(
+                    f"Accuracy of the ema network on the {len(dataset_val)} test images: {acc1:.1f}%"
+                )
+                if dist.get_rank() == 0 and acc1 > max_ema_accuracy:
+                    save_checkpoint(config,
+                                    epoch,
+                                    model_without_ddp,
+                                    max_accuracy,
+                                    optimizer,
+                                    lr_scheduler,
+                                    loss_scaler,
+                                    logger,
+                                    model_ema=model_ema,
+                                    best='ema_best')
+                max_ema_accuracy = max(max_ema_accuracy, acc1)
+                logger.info(f'Max ema accuracy: {max_ema_accuracy:.2f}%')
+
+    total_time = time.time() - start_time
+    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+    logger.info('Training time {}'.format(total_time_str))
+
+
+def train_one_epoch(config,
+                    model,
+                    criterion,
+                    data_loader,
+                    optimizer,
+                    epoch,
+                    mixup_fn,
+                    lr_scheduler,
+                    amp_autocast=suppress,
+                    loss_scaler=None,
+                    model_ema=None):
+    model.train()
+    optimizer.zero_grad()
+
+    num_steps = len(data_loader)
+    batch_time = AverageMeter()
+    model_time = AverageMeter()
+    loss_meter = AverageMeter()
+    norm_meter = MyAverageMeter(300)
+
+    start = time.time()
+    end = time.time()
+
+    amp_type = torch.float16 if config.AMP_TYPE == 'float16' else torch.bfloat16
+    for idx, (samples, targets) in enumerate(data_loader):
+        iter_begin_time = time.time()
+        samples = samples.cuda(non_blocking=True)
+        targets = targets.cuda(non_blocking=True)
+
+        if mixup_fn is not None:
+            samples, targets = mixup_fn(samples, targets)
+
+        if not obsolete_torch_version(TORCH_VERSION,
+                                      (1, 9)) and config.AMP_OPT_LEVEL != "O0":
+            with amp_autocast(dtype=amp_type):
+                outputs = model(samples)
+        else:
+            with amp_autocast():
+                outputs = model(samples)
+
+        if config.TRAIN.ACCUMULATION_STEPS > 1:
+            if not obsolete_torch_version(
+                    TORCH_VERSION, (1, 9)) and config.AMP_OPT_LEVEL != "O0":
+                with amp_autocast(dtype=amp_type):
+                    loss = criterion(outputs, targets)
+                    loss = loss / config.TRAIN.ACCUMULATION_STEPS
+            else:
+                with amp_autocast():
+                    loss = criterion(outputs, targets)
+                    loss = loss / config.TRAIN.ACCUMULATION_STEPS
+            if config.AMP_OPT_LEVEL != "O0":
+                is_second_order = hasattr(optimizer, 'is_second_order') and \
+                    optimizer.is_second_order
+                grad_norm = loss_scaler(loss,
+                                        optimizer,
+                                        clip_grad=config.TRAIN.CLIP_GRAD,
+                                        parameters=model.parameters(),
+                                        create_graph=is_second_order,
+                                        update_grad=(idx + 1) %
+                                        config.TRAIN.ACCUMULATION_STEPS == 0)
+                if (idx + 1) % config.TRAIN.ACCUMULATION_STEPS == 0:
+                    optimizer.zero_grad()
+                    if model_ema is not None:
+                        model_ema.update(model)
+            else:
+                loss.backward()
+                if config.TRAIN.CLIP_GRAD:
+                    grad_norm = torch.nn.utils.clip_grad_norm_(
+                        model.parameters(), config.TRAIN.CLIP_GRAD)
+                else:
+                    grad_norm = get_grad_norm(model.parameters())
+                if (idx + 1) % config.TRAIN.ACCUMULATION_STEPS == 0:
+                    optimizer.step()
+                    optimizer.zero_grad()
+                    if model_ema is not None:
+                        model_ema.update(model)
+            if (idx + 1) % config.TRAIN.ACCUMULATION_STEPS == 0:
+                lr_scheduler.step_update(epoch * num_steps + idx)
+        else:
+            if not obsolete_torch_version(
+                    TORCH_VERSION, (1, 9)) and config.AMP_OPT_LEVEL != "O0":
+                with amp_autocast(dtype=amp_type):
+                    loss = criterion(outputs, targets)
+            else:
+                with amp_autocast():
+                    loss = criterion(outputs, targets)
+            optimizer.zero_grad()
+            if config.AMP_OPT_LEVEL != "O0":
+                is_second_order = hasattr(optimizer, 'is_second_order') and \
+                    optimizer.is_second_order
+                grad_norm = loss_scaler(loss,
+                                        optimizer,
+                                        clip_grad=config.TRAIN.CLIP_GRAD,
+                                        parameters=model.parameters(),
+                                        create_graph=is_second_order,
+                                        update_grad=(idx + 1) %
+                                        config.TRAIN.ACCUMULATION_STEPS == 0)
+                if model_ema is not None:
+                    model_ema.update(model)
+            else:
+                loss.backward()
+                if config.TRAIN.CLIP_GRAD:
+                    grad_norm = torch.nn.utils.clip_grad_norm_(
+                        model.parameters(), config.TRAIN.CLIP_GRAD)
+                else:
+                    grad_norm = get_grad_norm(model.parameters())
+                optimizer.step()
+                if model_ema is not None:
+                    model_ema.update(model)
+
+            lr_scheduler.step_update(epoch * num_steps + idx)
+
+        torch.cuda.synchronize()
+
+        loss_meter.update(loss.item(), targets.size(0))
+        if grad_norm is not None:
+            norm_meter.update(grad_norm.item())
+        batch_time.update(time.time() - end)
+        model_time.update(time.time() - iter_begin_time)
+        end = time.time()
+
+        if idx % config.PRINT_FREQ == 0:
+            lr = optimizer.param_groups[0]['lr']
+            memory_used = torch.cuda.max_memory_allocated() / (1024.0 * 1024.0)
+            etas = batch_time.avg * (num_steps - idx)
+            logger.info(
+                f'Train: [{epoch}/{config.TRAIN.EPOCHS}][{idx}/{num_steps}]\t'
+                f'eta {datetime.timedelta(seconds=int(etas))} lr {lr:.6f}\t'
+                f'time {batch_time.val:.4f} ({batch_time.avg:.4f})\t'
+                f'model_time {model_time.val:.4f} ({model_time.avg:.4f})\t'
+                f'loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t'
+                f'grad_norm {norm_meter.val:.4f} ({norm_meter.avg:.4f}/{norm_meter.var:.4f})\t'
+                f'mem {memory_used:.0f}MB')
+    epoch_time = time.time() - start
+    logger.info(
+        f"EPOCH {epoch} training takes {datetime.timedelta(seconds=int(epoch_time))}"
+    )
+
+
+@torch.no_grad()
+def validate(config, data_loader, model, epoch=None, amp_autocast=None):
+    criterion = torch.nn.CrossEntropyLoss()
+    model.eval()
+
+    batch_time = AverageMeter()
+    loss_meter = AverageMeter()
+    acc1_meter = AverageMeter()
+    acc5_meter = AverageMeter()
+
+    end = time.time()
+    for idx, (images, target) in enumerate(data_loader):
+        images = images.cuda(non_blocking=True)
+        target = target.cuda(non_blocking=True)
+        with amp_autocast():
+            output = model(images)
+
+        # convert 22k to 1k to evaluate
+        if output.size(-1) == 21841:
+            convert_file = './meta_data/map22kto1k.txt'
+            with open(convert_file, 'r') as f:
+                convert_list = [int(line) for line in f.readlines()]
+            output = output[:, convert_list]
+
+        # measure accuracy and record loss
+        loss = criterion(output, target)
+        acc1, acc5 = accuracy(output, target, topk=(1, 5))
+
+        acc1 = reduce_tensor(acc1)
+        acc5 = reduce_tensor(acc5)
+        loss = reduce_tensor(loss)
+
+        loss_meter.update(loss.item(), target.size(0))
+        acc1_meter.update(acc1.item(), target.size(0))
+        acc5_meter.update(acc5.item(), target.size(0))
+
+        # measure elapsed time
+        batch_time.update(time.time() - end)
+        end = time.time()
+
+        if idx % config.PRINT_FREQ == 0:
+            memory_used = torch.cuda.max_memory_allocated() / (1024.0 * 1024.0)
+            logger.info(f'Test: [{idx}/{len(data_loader)}]\t'
+                        f'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
+                        f'Loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t'
+                        f'Acc@1 {acc1_meter.val:.3f} ({acc1_meter.avg:.3f})\t'
+                        f'Acc@5 {acc5_meter.val:.3f} ({acc5_meter.avg:.3f})\t'
+                        f'Mem {memory_used:.0f}MB')
+    if epoch is not None:
+        logger.info(
+            f'[Epoch:{epoch}] * Acc@1 {acc1_meter.avg:.3f} Acc@5 {acc5_meter.avg:.3f}'
+        )
+    else:
+        logger.info(
+            f' * Acc@1 {acc1_meter.avg:.3f} Acc@5 {acc5_meter.avg:.3f}')
+
+    return acc1_meter.avg, acc5_meter.avg, loss_meter.avg
+
+
+if __name__ == '__main__':
+    _, config = parse_option()
+
+    if config.AMP_OPT_LEVEL != "O0":
+        assert has_native_amp, "Please update pytorch(1.6+) to support amp!"
+
+    # init distributed env
+    if 'SLURM_PROCID' in os.environ and int(os.environ['SLURM_NNODES']) != 1:
+        print("\nDist init: SLURM")
+        rank = int(os.environ['SLURM_PROCID'])
+        gpu = rank % torch.cuda.device_count()
+        config.defrost()
+        config.LOCAL_RANK = gpu
+        config.freeze()
+
+        world_size = int(os.environ["SLURM_NTASKS"])
+        if "MASTER_PORT" not in os.environ:
+            os.environ["MASTER_PORT"] = "29501"
+        node_list = os.environ["SLURM_NODELIST"]
+        addr = subprocess.getoutput(
+            f"scontrol show hostname {node_list} | head -n1")
+        if "MASTER_ADDR" not in os.environ:
+            os.environ["MASTER_ADDR"] = addr
+
+        os.environ['RANK'] = str(rank)
+        os.environ['LOCAL_RANK'] = str(gpu)
+        os.environ['LOCAL_SIZE'] = str(torch.cuda.device_count())
+        os.environ['WORLD_SIZE'] = str(world_size)
+    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
+        rank = int(os.environ["RANK"])
+        world_size = int(os.environ['WORLD_SIZE'])
+        print(f"RANK and WORLD_SIZE in environ: {rank}/{world_size}")
+    else:
+        rank = -1
+        world_size = -1
+    torch.cuda.set_device(config.LOCAL_RANK)
+    torch.distributed.init_process_group(backend='nccl',
+                                         init_method='env://',
+                                         world_size=world_size,
+                                         rank=rank)
+    torch.distributed.barrier()
+
+    seed = config.SEED + dist.get_rank()
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    cudnn.benchmark = True
+
+    # linear scale the learning rate according to total batch size, may not be optimal
+    linear_scaled_lr = config.TRAIN.BASE_LR * \
+        config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0
+    linear_scaled_warmup_lr = config.TRAIN.WARMUP_LR * \
+        config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0
+    linear_scaled_min_lr = config.TRAIN.MIN_LR * \
+        config.DATA.BATCH_SIZE * dist.get_world_size() / 512.0
+    # gradient accumulation also need to scale the learning rate
+    if config.TRAIN.ACCUMULATION_STEPS > 1:
+        linear_scaled_lr = linear_scaled_lr * config.TRAIN.ACCUMULATION_STEPS
+        linear_scaled_warmup_lr = linear_scaled_warmup_lr * config.TRAIN.ACCUMULATION_STEPS
+        linear_scaled_min_lr = linear_scaled_min_lr * config.TRAIN.ACCUMULATION_STEPS
+    config.defrost()
+    config.TRAIN.BASE_LR = linear_scaled_lr
+    config.TRAIN.WARMUP_LR = linear_scaled_warmup_lr
+    config.TRAIN.MIN_LR = linear_scaled_min_lr
+    print(config.AMP_OPT_LEVEL, _.amp_opt_level)
+
+    config.freeze()
+
+    os.makedirs(config.OUTPUT, exist_ok=True)
+    logger = create_logger(output_dir=config.OUTPUT,
+                           dist_rank=dist.get_rank(),
+                           name=f"{config.MODEL.NAME}")
+
+    if dist.get_rank() == 0:
+        path = os.path.join(config.OUTPUT, "config.json")
+        with open(path, "w") as f:
+            f.write(config.dump())
+        logger.info(f"Full config saved to {path}")
+
+    # print config
+    logger.info(config.dump())
+
+    main(config)

classification/main_accelerate.py

@@ -0,0 +1,380 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+
+import datetime
+import argparse
+import os
+import time
+import logging
+import random
+
+import torch
+import torch.backends.cudnn as cudnn
+import numpy as np
+from accelerate import Accelerator
+from accelerate import GradScalerKwargs
+from accelerate.logging import get_logger
+from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy
+from timm.utils import AverageMeter, accuracy, ModelEma
+from tqdm import tqdm
+import warnings
+
+from config import get_config
+from models import build_model
+from dataset import build_loader2
+from lr_scheduler import build_scheduler
+from optimizer import build_optimizer
+from utils import load_pretrained, load_ema_checkpoint
+from ddp_hooks import fp16_compress_hook
+
+logger = get_logger(__name__)
+warnings.filterwarnings('ignore')
+
+
+def parse_option():
+    parser = argparse.ArgumentParser(
+        'InternImage training and evaluation script', add_help=False)
+    parser.add_argument('--cfg', type=str, required=True, metavar="FILE", help='path to config file')
+    parser.add_argument("--opts", help="Modify config options by adding 'KEY VALUE' pairs. ", default=None, nargs='+')
+
+    # easy config modification
+    parser.add_argument('--batch-size', type=int, help="batch size for single GPU")
+    parser.add_argument('--dataset', type=str, help='dataset name', default=None)
+    parser.add_argument('--data-path', type=str, help='path to dataset')
+    parser.add_argument('--zip', action='store_true', help='use zipped dataset instead of folder dataset')
+    parser.add_argument('--cache-mode', type=str, default='part', choices=['no', 'full', 'part'],
+                        help='no: no cache, '
+                        'full: cache all data, '
+                        'part: sharding the dataset into nonoverlapping pieces and only cache one piece'
+                        )
+    parser.add_argument('--pretrained', help='pretrained weight from checkpoint, could be imagenet22k pretrained weight')
+    parser.add_argument('--resume', help='resume from checkpoint')
+    parser.add_argument('--output', default='output', type=str, metavar='PATH',
+                        help='root of output folder, the full path is <output>/<model_name>/<tag> (default: output)'
+                        )
+    parser.add_argument('--eval', action='store_true', help='Perform evaluation only')
+    parser.add_argument('--throughput', action='store_true', help='Test throughput only')
+    parser.add_argument('--save-ckpt-num', default=1, type=int)
+    parser.add_argument('--accumulation-steps', type=int, default=1, help="gradient accumulation steps")
+    parser.add_argument('--disable-grad-scalar', action='store_true', help='disable Grad Scalar')
+    parser.add_argument(
+        "--logger",
+        type=str,
+        default="tensorboard",
+        choices=["tensorboard", "wandb"],
+        help=(
+            "Whether to use [tensorboard](https://www.tensorflow.org/tensorboard) or [wandb](https://www.wandb.ai)"
+            " for experiment tracking and logging of model metrics and model checkpoints"
+        ),
+    )
+
+    args, unparsed = parser.parse_known_args()
+    config = get_config(args)
+    config.defrost()
+    config.TRAIN.OPTIMIZER.USE_ZERO = False
+    config.OUTPUT += '_deepspeed'
+    config.DATA.IMG_ON_MEMORY = False
+    config.freeze()
+    return args, config
+
+
+def seed_everything(seed, rank):
+    seed = seed + rank
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    cudnn.benchmark = True
+
+
+def save_config(config):
+    path = os.path.join(config.OUTPUT, "config.json")
+    with open(path, "w") as f:
+        f.write(config.dump())
+    logger.info(f"Full config saved to {path}")
+
+
+def build_criterion(config):
+    if config.AUG.MIXUP > 0.:
+        # smoothing is handled with mixup label transform
+        criterion = SoftTargetCrossEntropy()
+    elif config.MODEL.LABEL_SMOOTHING > 0.:
+        criterion = LabelSmoothingCrossEntropy(
+            smoothing=config.MODEL.LABEL_SMOOTHING)
+    else:
+        criterion = torch.nn.CrossEntropyLoss()
+    return criterion
+
+
+def scale_learning_rate(config, num_processes):
+    # linear scale the learning rate according to total batch size, may not be optimal
+    linear_scaled_lr = config.TRAIN.BASE_LR * \
+        config.DATA.BATCH_SIZE * num_processes / 512.0
+    linear_scaled_warmup_lr = config.TRAIN.WARMUP_LR * \
+        config.DATA.BATCH_SIZE * num_processes / 512.0
+    linear_scaled_min_lr = config.TRAIN.MIN_LR * \
+        config.DATA.BATCH_SIZE * num_processes / 512.0
+    # gradient accumulation also need to scale the learning rate
+    if config.TRAIN.ACCUMULATION_STEPS > 1:
+        linear_scaled_lr = linear_scaled_lr * config.TRAIN.ACCUMULATION_STEPS
+        linear_scaled_warmup_lr = linear_scaled_warmup_lr * config.TRAIN.ACCUMULATION_STEPS
+        linear_scaled_min_lr = linear_scaled_min_lr * config.TRAIN.ACCUMULATION_STEPS
+    config.defrost()
+    config.TRAIN.BASE_LR = linear_scaled_lr
+    config.TRAIN.WARMUP_LR = linear_scaled_warmup_lr
+    config.TRAIN.MIN_LR = linear_scaled_min_lr
+    config.freeze()
+
+    logger.info('BASE_LR={}'.format(config.TRAIN.BASE_LR))
+    logger.info('WARMUP_LR={}'.format(config.TRAIN.WARMUP_LR))
+    logger.info('MIN_LR={}'.format(config.TRAIN.MIN_LR))
+
+
+def setup_autoresume(config):
+    if config.MODEL.RESUME == '' and config.TRAIN.AUTO_RESUME:
+        last_checkpoint = os.path.join(config.OUTPUT, 'last')
+        resume_file = last_checkpoint if os.path.exists(last_checkpoint) else None
+
+        if resume_file:
+            if config.MODEL.RESUME:
+                logger.warning(f"auto-resume changing resume file from {config.MODEL.RESUME} to {resume_file}")
+            config.defrost()
+            config.MODEL.RESUME = resume_file
+            config.freeze()
+            logger.info(f'auto resuming from {resume_file}')
+        else:
+            logger.info(f'no checkpoint found in {config.OUTPUT}, ignoring auto resume')
+
+
+def load_model_checkpoint(config, model, accelerator):
+    if config.MODEL.RESUME:
+        try:
+            checkpoint = torch.load(config.MODEL.RESUME)['model']
+            checkpoint = {k.replace('module.', ''): v for k, v in checkpoint.items()}
+            model.load_state_dict(checkpoint)
+        except:
+            accelerator.load_state(config.MODEL.RESUME)
+    elif config.MODEL.PRETRAINED:
+        try:
+            load_pretrained(config, model, logger)
+        except:
+            accelerator.load_state(config.MODEL.PRETRAINED)
+    return model
+
+
+def save_checkpoint(save_dir, accelerator, epoch, max_acc, config, lr_scheduler=None):
+    # let accelerator handle the model and optimizer state for ddp and deepspeed.
+    accelerator.save_state(save_dir)
+
+    if accelerator.is_main_process:
+        save_state = {
+            'lr_scheduler': lr_scheduler.state_dict(),
+            'max_acc': max_acc,
+            'epoch': epoch,
+            'config': config
+        }
+        torch.save(save_state, os.path.join(save_dir, 'additional_state.pth'))
+
+
+def load_checkpoint_if_needed(accelerator, config, lr_scheduler=None):
+    setup_autoresume(config)
+    save_dir = config.MODEL.RESUME
+    if not save_dir:
+        return 0.0
+    accelerator.load_state(save_dir)
+    checkpoint = torch.load(os.path.join(save_dir, 'additional_state.pth'), map_location='cpu')
+    if lr_scheduler is not None:
+        logger.info('resuming lr_scheduler')
+        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+    config.defrost()
+    config.TRAIN.START_EPOCH = checkpoint['epoch'] + 1
+    config.freeze()
+    max_acc = checkpoint.get('max_acc', 0.0)
+    logger.info(f"=> loaded successfully {config.MODEL.RESUME} (epoch {checkpoint['epoch']})")
+    return max_acc
+
+
+def log_model_statistic(model_wo_ddp):
+    n_parameters = sum(p.numel() for p in model_wo_ddp.parameters()
+                       if p.requires_grad)
+    logger.info(f"number of params: {n_parameters}")
+    if hasattr(model_wo_ddp, 'flops'):
+        flops = model_wo_ddp.flops()
+        logger.info(f"number of GFLOPs: {flops / 1e9}")
+
+
+def train_epoch(*, model, optimizer, data_loader, scheduler, criterion, mixup_fn,
+                accelerator: Accelerator, epoch, config):
+    model.train()
+
+    num_steps = len(data_loader)
+    batch_time = AverageMeter()
+    model_time = AverageMeter()
+    loss_meter = AverageMeter()
+
+    end = time.time()
+
+    gradient_accumulation_steps = config.TRAIN.ACCUMULATION_STEPS
+
+    for step, (samples, targets) in enumerate(data_loader):
+        iter_begin_time = time.time()
+
+        if mixup_fn is not None:
+            samples, targets = mixup_fn(samples, targets)
+
+        with accelerator.accumulate(model):
+            outputs = model(samples)
+            loss = criterion(outputs, targets)
+            accelerator.backward(loss)
+            if accelerator.sync_gradients:
+                accelerator.clip_grad_norm_(model.parameters(), config.TRAIN.CLIP_GRAD)
+            optimizer.step()
+            optimizer.zero_grad()
+
+        accelerator.wait_for_everyone()
+
+        if (step + 1) % gradient_accumulation_steps == 0:
+            if scheduler is not None:
+                scheduler.step_update((epoch * num_steps + step) // gradient_accumulation_steps)
+
+            batch_time.update(time.time() - end)
+            model_time.update(time.time() - iter_begin_time)
+            loss_meter.update(loss.item())
+            end = time.time()
+
+        if accelerator.is_main_process and step % config.PRINT_FREQ == 0:
+            lr = optimizer.param_groups[0]['lr']
+            memory_used = torch.cuda.max_memory_allocated() / (1024.0 * 1024.0)
+            etas = batch_time.avg * (num_steps - step)
+
+            logger.info(
+                f'Train: [{epoch}/{config.TRAIN.EPOCHS}][{step}/{num_steps}]\t'
+                f'eta {datetime.timedelta(seconds=int(etas))} lr {lr:.10f}\t'
+                f'time {batch_time.val:.4f} ({batch_time.avg:.4f})\t'
+                f'model_time {model_time.val:.4f} ({model_time.avg:.4f})\t'
+                f'loss {loss_meter.val:.8f} ({loss_meter.avg:.4f})\t'
+                f'mem {memory_used:.0f}MB')
+
+
+@torch.no_grad()
+def eval_epoch(*, config, data_loader, model, accelerator: Accelerator):
+    model.eval()
+
+    acc1_meter = AverageMeter()
+    acc5_meter = AverageMeter()
+
+    for idx, (images, target) in enumerate(tqdm(data_loader, disable=accelerator.is_main_process)):
+        output = model(images)
+
+        # convert 22k to 1k to evaluate
+        if output.size(-1) == 21841:
+            convert_file = './meta_data/map22kto1k.txt'
+            with open(convert_file, 'r') as f:
+                convert_list = [int(line) for line in f.readlines()]
+            output = output[:, convert_list]
+
+        acc1, acc5 = accuracy(output, target, topk=(1, 5))
+        acc1 = accelerator.gather(acc1).mean(0)
+        acc5 = accelerator.gather(acc5).mean(0)
+
+        acc1_meter.update(acc1.item(), target.size(0))
+        acc5_meter.update(acc5.item(), target.size(0))
+
+        if (idx + 1) % config.PRINT_FREQ == 0 or idx + 1 == len(data_loader):
+            logger.info(f'Test: [{idx+1}/{len(data_loader)}]\t'
+                        f'Acc@1 {acc1_meter.val:.3f} ({acc1_meter.avg:.3f})\t'
+                        f'Acc@5 {acc5_meter.val:.3f} ({acc5_meter.avg:.3f})\t'
+                        )
+    return acc1_meter.avg
+
+
+def eval(config, accelerator: Accelerator):
+    _, _, _, _, validate_dataloader, _, _ = build_loader2(config)
+    model = build_model(config)
+    model, validate_dataloader = accelerator.prepare(model, validate_dataloader)
+    model = load_model_checkpoint(config, model, accelerator)
+    log_model_statistic(accelerator.unwrap_model(model))
+    eval_epoch(config=config, data_loader=validate_dataloader, model=model, accelerator=accelerator)
+
+
+def train(config, accelerator: Accelerator):
+    _, _, _, training_dataloader, validate_dataloader, _, mixup_fn = build_loader2(config)
+    model = build_model(config)
+    optimizer = build_optimizer(config, model)
+    criterion = build_criterion(config)
+
+    model, optimizer, training_dataloader, validate_dataloader = accelerator.prepare(
+        model, optimizer, training_dataloader, validate_dataloader)
+
+    effective_update_steps_per_epoch = len(training_dataloader) // config.TRAIN.ACCUMULATION_STEPS
+    lr_scheduler = build_scheduler(config, optimizer, effective_update_steps_per_epoch)
+
+    try:
+        model.register_comm_hook(state=None, hook=fp16_compress_hook)
+        logger.info('using fp16_compress_hook!')
+    except:
+        logger.info("cannot register fp16_compress_hook!")
+
+    max_acc = load_checkpoint_if_needed(accelerator, config, lr_scheduler)
+
+    logger.info(f"Created model:{config.MODEL.TYPE}/{config.MODEL.NAME}")
+    logger.info(str(model))
+    logger.info("Effective Optimizer Steps: {}".format(effective_update_steps_per_epoch))
+    logger.info("Start training")
+    logger.info("Max accuracy: {}".format(max_acc))
+    log_model_statistic(accelerator.unwrap_model(model))
+
+    for epoch in range(config.TRAIN.START_EPOCH, config.TRAIN.EPOCHS):
+        train_epoch(model=model, optimizer=optimizer, data_loader=training_dataloader,
+                    scheduler=lr_scheduler, criterion=criterion, mixup_fn=mixup_fn,
+                    accelerator=accelerator, epoch=epoch, config=config)
+        acc = eval_epoch(config=config, data_loader=validate_dataloader, model=model,
+                         accelerator=accelerator)
+
+        accelerator.wait_for_everyone()
+        if acc > max_acc:
+            max_acc = acc
+            save_checkpoint(os.path.join(config.OUTPUT, 'best'), accelerator, epoch, max_acc, config, lr_scheduler)
+        logger.info(f'Max Acc@1 {max_acc:.3f}')
+        save_checkpoint(os.path.join(config.OUTPUT, 'last'), accelerator, epoch, max_acc, config, lr_scheduler)
+
+
+def main():
+    args, config = parse_option()
+    os.makedirs(config.OUTPUT, exist_ok=True)
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        filename=os.path.join(config.OUTPUT, 'run.log'),
+        level=logging.INFO,
+    )
+
+    loggers = ['tensorboard']
+    accelerator = Accelerator(
+        log_with=loggers,
+        project_dir=config.OUTPUT,
+        gradient_accumulation_steps=config.TRAIN.ACCUMULATION_STEPS,
+        # When use deepspeed, you could not comment this out
+        # even if you set loss scale to 1.0 in deepspeed config.
+        kwargs_handlers=[GradScalerKwargs(enabled=not args.disable_grad_scalar)],
+    )
+    logger.info(accelerator.state, main_process_only=False)
+
+    scale_learning_rate(config, accelerator.num_processes)
+    seed_everything(config.SEED, accelerator.process_index)
+    save_config(config)
+
+    logger.info(config.dump())
+
+    if config.EVAL_MODE:
+        eval(config, accelerator)
+    else:
+        train(config, accelerator)
+
+
+if __name__ == '__main__':
+    main()

classification/main_deepspeed.py

@@ -0,0 +1,531 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+
+import os
+import time
+import random
+import argparse
+import datetime
+import numpy as np
+import subprocess
+
+import torch
+import torch.backends.cudnn as cudnn
+import torch.distributed as dist
+import deepspeed
+from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy
+from timm.utils import accuracy, AverageMeter
+
+from config import get_config
+from models import build_model
+from dataset import build_loader
+from lr_scheduler import build_scheduler
+from optimizer import set_weight_decay_and_lr
+from logger import create_logger
+from utils import load_pretrained, reduce_tensor, MyAverageMeter
+from ddp_hooks import fp16_compress_hook
+from ema_deepspeed import EMADeepspeed
+
+def parse_option():
+    parser = argparse.ArgumentParser(
+        'InternImage training and evaluation script', add_help=False)
+    parser.add_argument('--cfg', type=str, required=True, metavar="FILE", help='path to config file')
+    parser.add_argument("--opts", help="Modify config options by adding 'KEY VALUE' pairs. ", default=None, nargs='+')
+
+    # easy config modification
+    parser.add_argument('--batch-size', type=int, help="batch size for single GPU")
+    parser.add_argument('--dataset', type=str, help='dataset name', default=None)
+    parser.add_argument('--data-path', type=str, help='path to dataset')
+    parser.add_argument('--zip', action='store_true', help='use zipped dataset instead of folder dataset')
+    parser.add_argument('--cache-mode', type=str, default='part', choices=['no', 'full', 'part'],
+                        help='no: no cache, '
+                        'full: cache all data, '
+                        'part: sharding the dataset into nonoverlapping pieces and only cache one piece'
+                        )
+    parser.add_argument('--pretrained', help='pretrained weight from checkpoint, could be imagenet22k pretrained weight')
+    parser.add_argument('--resume', help='resume from checkpoint')
+    parser.add_argument('--output', default='output', type=str, metavar='PATH',
+                        help='root of output folder, the full path is <output>/<model_name>/<tag> (default: output)'
+                        )
+
+    parser.add_argument('--eval', action='store_true', help='Perform evaluation only')
+    parser.add_argument('--throughput', action='store_true', help='Test throughput only')
+    parser.add_argument('--save-ckpt-num', default=1, type=int)
+    parser.add_argument('--accumulation-steps', type=int, default=1, help="gradient accumulation steps")
+
+    # distributed training
+    parser.add_argument("--local-rank", type=int, required=True, help='local rank for DistributedDataParallel')
+    parser.add_argument('--disable-grad-scalar', action='store_true', help='disable Grad Scalar')
+
+    args, unparsed = parser.parse_known_args()
+    config = get_config(args)
+
+    return args, config
+
+
+def seed_everything(seed, rank):
+    seed = seed + rank
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    cudnn.benchmark = True
+
+
+def save_config(config):
+    path = os.path.join(config.OUTPUT, "config.json")
+    with open(path, "w") as f:
+        f.write(config.dump())
+    logger.info(f"Full config saved to {path}")
+
+
+def build_criterion(config):
+    if config.AUG.MIXUP > 0.:
+        # smoothing is handled with mixup label transform
+        criterion = SoftTargetCrossEntropy()
+    elif config.MODEL.LABEL_SMOOTHING > 0.:
+        criterion = LabelSmoothingCrossEntropy(
+            smoothing=config.MODEL.LABEL_SMOOTHING)
+    else:
+        criterion = torch.nn.CrossEntropyLoss()
+    return criterion
+
+
+def scale_learning_rate(config, num_processes):
+    # linear scale the learning rate according to total batch size, may not be optimal
+    linear_scaled_lr = config.TRAIN.BASE_LR * \
+        config.DATA.BATCH_SIZE * num_processes / 512.0
+    linear_scaled_warmup_lr = config.TRAIN.WARMUP_LR * \
+        config.DATA.BATCH_SIZE * num_processes / 512.0
+    linear_scaled_min_lr = config.TRAIN.MIN_LR * \
+        config.DATA.BATCH_SIZE * num_processes / 512.0
+    # gradient accumulation also need to scale the learning rate
+    if config.TRAIN.ACCUMULATION_STEPS > 1:
+        linear_scaled_lr = linear_scaled_lr * config.TRAIN.ACCUMULATION_STEPS
+        linear_scaled_warmup_lr = linear_scaled_warmup_lr * config.TRAIN.ACCUMULATION_STEPS
+        linear_scaled_min_lr = linear_scaled_min_lr * config.TRAIN.ACCUMULATION_STEPS
+    config.defrost()
+    config.TRAIN.BASE_LR = linear_scaled_lr
+    config.TRAIN.WARMUP_LR = linear_scaled_warmup_lr
+    config.TRAIN.MIN_LR = linear_scaled_min_lr
+    config.freeze()
+
+    logger.info('BASE_LR={}'.format(config.TRAIN.BASE_LR))
+    logger.info('WARMUP_LR={}'.format(config.TRAIN.WARMUP_LR))
+    logger.info('MIN_LR={}'.format(config.TRAIN.MIN_LR))
+
+
+def log_model_statistic(model_wo_ddp):
+    n_parameters = sum(p.numel() for p in model_wo_ddp.parameters()
+                       if p.requires_grad)
+    logger.info(f"number of params: {n_parameters/1e6} M")
+    if hasattr(model_wo_ddp, 'flops'):
+        flops = model_wo_ddp.flops()
+        logger.info(f"number of GFLOPs: {flops / 1e9}")
+
+
+def get_parameter_groups(model, config):
+    skip = {}
+    skip_keywords = {}
+    if hasattr(model, 'no_weight_decay'):
+        skip = model.no_weight_decay()
+    if hasattr(model, 'no_weight_decay_keywords'):
+        skip_keywords = model.no_weight_decay_keywords()
+
+    parameters = set_weight_decay_and_lr(
+        model,
+        config.TRAIN.WEIGHT_DECAY,
+        config.TRAIN.BASE_LR,
+        skip,
+        skip_keywords,
+        lr_layer_decay=config.TRAIN.LR_LAYER_DECAY,
+        lr_layer_decay_ratio=config.TRAIN.LR_LAYER_DECAY_RATIO,
+        freeze_backbone=config.TRAIN.OPTIMIZER.FREEZE_BACKBONE,
+        dcn_lr_mul=config.TRAIN.OPTIMIZER.DCN_LR_MUL,
+    )
+    return parameters
+
+
+def get_optimizer_state_str(optimizer):
+    states = []
+    for param_group in optimizer.param_groups:
+        states.append(f'name={param_group["name"]} lr={param_group["lr"]} weight_decay={param_group["weight_decay"]}')
+    return '\n'.join(states)
+
+
+def build_ds_config(config, args):
+    opt_lower = config.TRAIN.OPTIMIZER.NAME.lower()
+    if opt_lower == 'adamw':
+        optimizer = {
+            "type": "AdamW",
+            "params": {
+                "lr": config.TRAIN.BASE_LR,
+                "eps": config.TRAIN.OPTIMIZER.EPS,
+                "betas": config.TRAIN.OPTIMIZER.BETAS,
+                "weight_decay": config.TRAIN.WEIGHT_DECAY
+            }
+        }
+    else:
+        return NotImplemented
+
+    ds_config = {
+        "train_micro_batch_size_per_gpu": config.DATA.BATCH_SIZE,
+        "optimizer": optimizer,
+        "fp16": {
+            "enabled": True,
+            "auto_cast": True,
+            "loss_scale": 1 if args.disable_grad_scalar else 0
+        },
+        "zero_optimization": {
+            "stage": 1,
+        },
+        "steps_per_print": 1e10,
+        "gradient_accumulation_steps": config.TRAIN.ACCUMULATION_STEPS,
+        "gradient_clipping": config.TRAIN.CLIP_GRAD,
+    }
+    return ds_config
+
+
+@torch.no_grad()
+def throughput(data_loader, model, logger):
+    model.eval()
+
+    for idx, (images, _) in enumerate(data_loader):
+        images = images.cuda(non_blocking=True)
+        batch_size = images.shape[0]
+        for i in range(50):
+            model(images)
+        torch.cuda.synchronize()
+        logger.info(f"throughput averaged with 30 times")
+        tic1 = time.time()
+        for i in range(30):
+            model(images)
+        torch.cuda.synchronize()
+        tic2 = time.time()
+        logger.info(
+            f"batch_size {batch_size} throughput {30 * batch_size / (tic2 - tic1)}"
+        )
+        return
+
+
+def train_epoch(config, model, criterion, data_loader, optimizer, epoch, mixup_fn, lr_scheduler, model_ema=None):
+    model.train()
+
+    num_steps = len(data_loader)
+    batch_time = AverageMeter()
+    model_time = AverageMeter()
+    loss_meter = AverageMeter()
+    norm_meter = MyAverageMeter(300)
+
+    start = time.time()
+    end = time.time()
+
+    for idx, (samples, targets) in enumerate(data_loader):
+        iter_begin_time = time.time()
+        samples = samples.cuda(non_blocking=True)
+        targets = targets.cuda(non_blocking=True)
+
+        if mixup_fn is not None:
+            samples, targets = mixup_fn(samples, targets)
+
+        outputs = model(samples)
+        loss = criterion(outputs, targets)
+
+        model.backward(loss)
+        model.step()
+
+        if model_ema is not None:
+            model_ema(model)
+
+        if (idx + 1) % config.TRAIN.ACCUMULATION_STEPS == 0:
+            lr_scheduler.step_update(epoch * num_steps + idx)
+
+        torch.cuda.synchronize()
+        loss_meter.update(loss.item(), targets.size(0))
+        norm_meter.update(optimizer._global_grad_norm)
+        batch_time.update(time.time() - end)
+        model_time.update(time.time() - iter_begin_time)
+        end = time.time()
+
+        if idx % config.PRINT_FREQ == 0:
+            lr = optimizer.param_groups[0]['lr']
+            memory_used = torch.cuda.max_memory_allocated() / (1024.0 * 1024.0)
+            etas = batch_time.avg * (num_steps - idx)
+            logger.info(
+                f'Train: [{epoch}/{config.TRAIN.EPOCHS}][{idx}/{num_steps}]\t'
+                f'eta {datetime.timedelta(seconds=int(etas))} lr {lr:.6f}\t'
+                f'time {batch_time.val:.4f} ({batch_time.avg:.4f})\t'
+                f'model_time {model_time.val:.4f} ({model_time.avg:.4f})\t'
+                f'loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t'
+                f'grad_norm {norm_meter.val:.4f} ({norm_meter.avg:.4f}/{norm_meter.var:.4f})\t'
+                f'mem {memory_used:.0f}MB')
+
+    epoch_time = time.time() - start
+    logger.info(f"EPOCH {epoch} training takes {datetime.timedelta(seconds=int(epoch_time))}")
+
+
+@torch.no_grad()
+def eval_epoch(config, data_loader, model, epoch=None):
+    criterion = torch.nn.CrossEntropyLoss()
+    model.eval()
+
+    batch_time = AverageMeter()
+    loss_meter = AverageMeter()
+    acc1_meter = AverageMeter()
+    acc5_meter = AverageMeter()
+
+    end = time.time()
+    for idx, (images, target) in enumerate(data_loader):
+        images = images.cuda(non_blocking=True)
+        target = target.cuda(non_blocking=True)
+        output = model(images)
+
+        # convert 22k to 1k to evaluate
+        if output.size(-1) == 21841:
+            convert_file = './meta_data/map22kto1k.txt'
+            with open(convert_file, 'r') as f:
+                convert_list = [int(line) for line in f.readlines()]
+            output = output[:, convert_list]
+
+        # measure accuracy and record loss
+        loss = criterion(output, target)
+        acc1, acc5 = accuracy(output, target, topk=(1, 5))
+
+        acc1 = reduce_tensor(acc1)
+        acc5 = reduce_tensor(acc5)
+        loss = reduce_tensor(loss)
+
+        loss_meter.update(loss.item(), target.size(0))
+        acc1_meter.update(acc1.item(), target.size(0))
+        acc5_meter.update(acc5.item(), target.size(0))
+
+        # measure elapsed time
+        batch_time.update(time.time() - end)
+        end = time.time()
+
+        if idx % config.PRINT_FREQ == 0:
+            memory_used = torch.cuda.max_memory_allocated() / (1024.0 * 1024.0)
+            logger.info(f'Test: [{idx}/{len(data_loader)}]\t'
+                        f'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
+                        f'Loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t'
+                        f'Acc@1 {acc1_meter.val:.3f} ({acc1_meter.avg:.3f})\t'
+                        f'Acc@5 {acc5_meter.val:.3f} ({acc5_meter.avg:.3f})\t'
+                        f'Mem {memory_used:.0f}MB')
+    if epoch is not None:
+        logger.info(f'[Epoch:{epoch}] * Acc@1 {acc1_meter.avg:.3f} Acc@5 {acc5_meter.avg:.3f}')
+    else:
+        logger.info(f' * Acc@1 {acc1_meter.avg:.3f} Acc@5 {acc5_meter.avg:.3f}')
+
+    return acc1_meter.avg, acc5_meter.avg, loss_meter.avg
+
+
+def train(config, ds_config):
+    # -------------- build ---------------- #
+    
+    _, dataset_val, _, data_loader_train, data_loader_val, _, mixup_fn = build_loader(config)
+    model = build_model(config)
+    model.cuda()
+
+    if config.MODEL.PRETRAINED:
+        load_pretrained(config, model, logger)
+
+    logger.info(ds_config)
+    model, optimizer, _, _ = deepspeed.initialize(
+        config=ds_config,
+        model=model,
+        model_parameters=get_parameter_groups(model, config),
+        dist_init_required=False,
+    )
+
+    try:
+        model.register_comm_hook(state=None, hook=fp16_compress_hook)
+        logger.info('using fp16_compress_hook!')
+    except:
+        logger.info("cannot register fp16_compress_hook!")
+
+    model_without_ddp = model.module
+
+    lr_scheduler = build_scheduler(config, optimizer, len(data_loader_train))
+    criterion = build_criterion(config)
+
+    model_ema = None
+    if config.TRAIN.EMA.ENABLE:
+        model_ema = EMADeepspeed(model, config.TRAIN.EMA.DECAY)
+
+    # -------------- resume ---------------- #
+    
+    max_accuracy = 0.0
+    max_accuracy_ema = 0.0
+    client_state = {}
+    if config.MODEL.RESUME == '' and config.TRAIN.AUTO_RESUME:
+        if os.path.exists(os.path.join(config.OUTPUT, 'latest')):
+            config.defrost()
+            config.MODEL.RESUME = config.OUTPUT
+            config.freeze()
+            tag = None
+    elif config.MODEL.RESUME:
+        config.MODEL.RESUME = os.path.dirname(config.MODEL.RESUME)
+        tag = os.path.basename(config.MODEL.RESUME)
+    if config.MODEL.RESUME:
+        logger.info('loading checkpoint from {}'.format(config.MODEL.RESUME))
+        _, client_state = model.load_checkpoint(load_dir=config.MODEL.RESUME, tag=tag)
+        logger.info(f'client_state={client_state.keys()}')
+        lr_scheduler.load_state_dict(client_state['custom_lr_scheduler'])
+        max_accuracy = client_state['max_accuracy']
+
+        if model_ema is not None:
+            max_accuracy_ema = client_state.get('max_accuracy_ema', 0.0)
+            model_ema.load_state_dict((client_state['model_ema']))
+        
+    # -------------- training ---------------- #
+    
+    logger.info(f"Creating model:{config.MODEL.TYPE}/{config.MODEL.NAME}")
+    logger.info(str(model))
+    logger.info(get_optimizer_state_str(optimizer))
+    logger.info("Start training")
+    logger.info('max_accuracy: {}'.format(max_accuracy))
+    log_model_statistic(model_without_ddp)
+
+    start_time = time.time()
+    start_epoch = client_state['epoch'] + 1 if 'epoch' in client_state else config.TRAIN.START_EPOCH
+    for epoch in range(start_epoch, config.TRAIN.EPOCHS):
+        data_loader_train.sampler.set_epoch(epoch)
+        train_epoch(config, model, criterion, data_loader_train, optimizer, epoch, mixup_fn, lr_scheduler,
+                    model_ema=model_ema)
+
+        if epoch % config.SAVE_FREQ == 0 or epoch == config.TRAIN.EPOCHS - 1:
+            model.save_checkpoint(
+                save_dir=config.OUTPUT,
+                tag=f'epoch{epoch}',
+                client_state={
+                    'custom_lr_scheduler': lr_scheduler.state_dict(),
+                    'max_accuracy': max_accuracy,
+                    'epoch': epoch,
+                    'config': config,
+                    'max_accuracy_ema': max_accuracy_ema if model_ema is not None else 0.0,
+                    'model_ema': model_ema.state_dict() if model_ema is not None else None,
+                }
+            )
+
+        if epoch % config.EVAL_FREQ == 0:
+            acc1, _, _ = eval_epoch(config, data_loader_val, model, epoch)
+            logger.info(f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%")
+
+            if acc1 > max_accuracy:
+                model.save_checkpoint(
+                    save_dir=config.OUTPUT,
+                    tag='best',
+                    client_state={
+                        'custom_lr_scheduler': lr_scheduler.state_dict(),
+                        'max_accuracy': max_accuracy,
+                        'epoch': epoch,
+                        'config': config,
+                        'max_accuracy_ema': max_accuracy_ema if model_ema is not None else 0.0,
+                        'model_ema': model_ema.state_dict() if model_ema is not None else None,
+                    }
+                )
+
+            max_accuracy = max(max_accuracy, acc1)
+            logger.info(f'Max accuracy: {max_accuracy:.2f}%')
+
+            if model_ema is not None:
+                with model_ema.activate(model):
+                    acc1_ema, _, _ = eval_epoch(config, data_loader_val, model, epoch)
+                    logger.info(f"[EMA] Accuracy of the network on the {len(dataset_val)} test images: {acc1_ema:.1f}%")
+                    max_accuracy_ema = max(max_accuracy_ema, acc1_ema)
+                    logger.info(f'[EMA] Max accuracy: {max_accuracy_ema:.2f}%')
+
+    total_time = time.time() - start_time
+    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+    logger.info('Training time {}'.format(total_time_str))
+
+
+def eval(config):
+    _, _, _, _, data_loader_val, _, _ = build_loader(config)
+    model = build_model(config)
+    model.cuda()
+    model = torch.nn.parallel.DistributedDataParallel(
+        model, device_ids=[config.LOCAL_RANK], broadcast_buffers=False)
+
+    model_wo_ddp = model.module
+    if config.MODEL.RESUME:
+        try:
+            checkpoint = torch.load(config.MODEL.RESUME, map_location='cpu')
+            msg = model_wo_ddp.load_state_dict(checkpoint['model'], strict=False)
+            logger.info(msg)
+        except:
+            try:
+                from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+                ckpt_dir = os.path.dirname(config.MODEL.RESUME)
+                tag = os.path.basename(config.MODEL.RESUME)
+                state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir=ckpt_dir, tag=tag) 
+                model_wo_ddp.load_state_dict(state_dict)
+            except:
+                checkpoint = torch.load(os.path.join(config.MODEL.RESUME, 'mp_rank_00_model_states.pt'), map_location='cpu')
+                model_wo_ddp.load_state_dict(checkpoint['module'])
+    elif config.MODEL.PRETRAINED:
+        load_pretrained(config, model_wo_ddp, logger)
+
+    if config.THROUGHPUT_MODE:
+        throughput(data_loader_val, model, logger)
+
+    eval_epoch(config, data_loader_val, model)
+
+
+if __name__ == '__main__':
+    args, config = parse_option()
+
+    # init distributed env
+    if 'SLURM_PROCID' in os.environ and int(os.environ['SLURM_TASKS_PER_NODE']) != 1:
+        print("\nDist init: SLURM")
+        rank = int(os.environ['SLURM_PROCID'])
+        gpu = rank % torch.cuda.device_count()
+        config.defrost()
+        config.LOCAL_RANK = gpu
+        config.freeze()
+
+        world_size = int(os.environ["SLURM_NTASKS"])
+        if "MASTER_PORT" not in os.environ:
+            os.environ["MASTER_PORT"] = "29501"
+        node_list = os.environ["SLURM_NODELIST"]
+        addr = subprocess.getoutput(
+            f"scontrol show hostname {node_list} | head -n1")
+        if "MASTER_ADDR" not in os.environ:
+            os.environ["MASTER_ADDR"] = addr
+
+        os.environ['RANK'] = str(rank)
+        os.environ['LOCAL_RANK'] = str(gpu)
+        os.environ['LOCAL_SIZE'] = str(torch.cuda.device_count())
+        os.environ['WORLD_SIZE'] = str(world_size)
+    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
+        rank = int(os.environ["RANK"])
+        world_size = int(os.environ['WORLD_SIZE'])
+        print(f"RANK and WORLD_SIZE in environ: {rank}/{world_size}")
+    else:
+        rank = -1
+        world_size = -1
+    torch.cuda.set_device(config.LOCAL_RANK)
+    torch.distributed.init_process_group(backend='nccl',
+                                         init_method='env://',
+                                         world_size=world_size,
+                                         rank=rank)
+    torch.distributed.barrier()
+
+    os.makedirs(config.OUTPUT, exist_ok=True)
+    logger = create_logger(output_dir=config.OUTPUT,
+                           dist_rank=dist.get_rank(),
+                           name=f"{config.MODEL.NAME}")
+    logger.info(config.dump())
+
+    if dist.get_rank() == 0: save_config(config)
+    scale_learning_rate(config, dist.get_world_size())
+    seed_everything(config.SEED, dist.get_rank())
+
+    if config.EVAL_MODE:
+        eval(config)
+    else:
+        train(config, build_ds_config(config, args))

classification/meta_data/meta

@@ -0,0 +1 @@
+/mnt/petrelfs/share/images/meta/

classification/models/__init__.py

@@ -0,0 +1,7 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+from .build import build_model

classification/models/build.py

@@ -0,0 +1,58 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+
+from .intern_image import InternImage
+from .flash_intern_image import FlashInternImage
+
+def build_model(config):
+    model_type = config.MODEL.TYPE
+    if model_type == 'intern_image':
+        model = InternImage(
+            core_op=config.MODEL.INTERN_IMAGE.CORE_OP,
+            num_classes=config.MODEL.NUM_CLASSES,
+            channels=config.MODEL.INTERN_IMAGE.CHANNELS,
+            depths=config.MODEL.INTERN_IMAGE.DEPTHS,
+            groups=config.MODEL.INTERN_IMAGE.GROUPS,
+            layer_scale=config.MODEL.INTERN_IMAGE.LAYER_SCALE,
+            offset_scale=config.MODEL.INTERN_IMAGE.OFFSET_SCALE,
+            post_norm=config.MODEL.INTERN_IMAGE.POST_NORM,
+            mlp_ratio=config.MODEL.INTERN_IMAGE.MLP_RATIO,
+            with_cp=config.TRAIN.USE_CHECKPOINT,
+            drop_path_rate=config.MODEL.DROP_PATH_RATE,
+            res_post_norm=config.MODEL.INTERN_IMAGE.RES_POST_NORM, # for InternImage-H/G
+            dw_kernel_size=config.MODEL.INTERN_IMAGE.DW_KERNEL_SIZE, # for InternImage-H/G
+            use_clip_projector=config.MODEL.INTERN_IMAGE.USE_CLIP_PROJECTOR, # for InternImage-H/G
+            level2_post_norm=config.MODEL.INTERN_IMAGE.LEVEL2_POST_NORM, # for InternImage-H/G
+            level2_post_norm_block_ids=config.MODEL.INTERN_IMAGE.LEVEL2_POST_NORM_BLOCK_IDS, # for InternImage-H/G
+            center_feature_scale=config.MODEL.INTERN_IMAGE.CENTER_FEATURE_SCALE # for InternImage-H/G
+        )
+    elif model_type == 'flash_intern_image':
+        model = FlashInternImage(
+            core_op=config.MODEL.FLASH_INTERN_IMAGE.CORE_OP,
+            num_classes=config.MODEL.NUM_CLASSES,
+            channels=config.MODEL.FLASH_INTERN_IMAGE.CHANNELS,
+            depths=config.MODEL.FLASH_INTERN_IMAGE.DEPTHS,
+            groups=config.MODEL.FLASH_INTERN_IMAGE.GROUPS,
+            layer_scale=config.MODEL.FLASH_INTERN_IMAGE.LAYER_SCALE,
+            offset_scale=config.MODEL.FLASH_INTERN_IMAGE.OFFSET_SCALE,
+            post_norm=config.MODEL.FLASH_INTERN_IMAGE.POST_NORM,
+            mlp_ratio=config.MODEL.FLASH_INTERN_IMAGE.MLP_RATIO,
+            with_cp=config.TRAIN.USE_CHECKPOINT,
+            drop_path_rate=config.MODEL.DROP_PATH_RATE,
+            mlp_fc2_bias=config.MODEL.FLASH_INTERN_IMAGE.MLP_FC2_BIAS,
+            dcn_output_bias=config.MODEL.FLASH_INTERN_IMAGE.DCN_OUTPUT_BIAS,
+            res_post_norm=config.MODEL.FLASH_INTERN_IMAGE.RES_POST_NORM, # for InternImage-H/G
+            dw_kernel_size=config.MODEL.FLASH_INTERN_IMAGE.DW_KERNEL_SIZE,
+            use_clip_projector=config.MODEL.FLASH_INTERN_IMAGE.USE_CLIP_PROJECTOR, # for InternImage-H/G
+            level2_post_norm=config.MODEL.FLASH_INTERN_IMAGE.LEVEL2_POST_NORM, # for InternImage-H/G
+            level2_post_norm_block_ids=config.MODEL.FLASH_INTERN_IMAGE.LEVEL2_POST_NORM_BLOCK_IDS, # for InternImage-H/G
+            center_feature_scale=config.MODEL.FLASH_INTERN_IMAGE.CENTER_FEATURE_SCALE # for InternImage-H/G
+        )
+    else:
+        raise NotImplementedError(f"Unkown model: {model_type}")
+
+    return model

classification/models/flash_intern_image.py

@@ -0,0 +1,795 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+
+import torch
+import torch.nn as nn
+from collections import OrderedDict
+import torch.utils.checkpoint as checkpoint
+from timm.models.layers import trunc_normal_, DropPath
+import torch.nn.functional as F
+import DCNv4
+
+
+class to_channels_first(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return x.permute(0, 3, 1, 2)
+
+
+class to_channels_last(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return x.permute(0, 2, 3, 1)
+
+
+def build_norm_layer(dim,
+                     norm_layer,
+                     in_format='channels_last',
+                     out_format='channels_last',
+                     eps=1e-6):
+    layers = []
+    if norm_layer == 'BN':
+        if in_format == 'channels_last':
+            layers.append(to_channels_first())
+        layers.append(nn.BatchNorm2d(dim))
+        if out_format == 'channels_last':
+            layers.append(to_channels_last())
+    elif norm_layer == 'LN':
+        if in_format == 'channels_first':
+            layers.append(to_channels_last())
+        layers.append(nn.LayerNorm(dim, eps=eps))
+        if out_format == 'channels_first':
+            layers.append(to_channels_first())
+    else:
+        raise NotImplementedError(
+            f'build_norm_layer does not support {norm_layer}')
+    return nn.Sequential(*layers)
+
+
+def build_act_layer(act_layer):
+    if act_layer == 'ReLU':
+        return nn.ReLU(inplace=True)
+    elif act_layer == 'SiLU':
+        return nn.SiLU(inplace=True)
+    elif act_layer == 'GELU':
+        return nn.GELU()
+
+    raise NotImplementedError(f'build_act_layer does not support {act_layer}')
+
+
+class CrossAttention(nn.Module):
+    r""" Cross Attention Module
+    Args:
+        dim (int): Number of input channels.
+        num_heads (int): Number of attention heads. Default: 8
+        qkv_bias (bool, optional):  If True, add a learnable bias to q, k, v.
+            Default: False.
+        qk_scale (float | None, optional): Override default qk scale of
+            head_dim ** -0.5 if set. Default: None.
+        attn_drop (float, optional): Dropout ratio of attention weight.
+            Default: 0.0
+        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
+        attn_head_dim (int, optional): Dimension of attention head.
+        out_dim (int, optional): Dimension of output.
+    """
+
+    def __init__(self,
+                 dim,
+                 num_heads=8,
+                 qkv_bias=False,
+                 qk_scale=None,
+                 attn_drop=0.,
+                 proj_drop=0.,
+                 attn_head_dim=None,
+                 out_dim=None):
+        super().__init__()
+        if out_dim is None:
+            out_dim = dim
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        if attn_head_dim is not None:
+            head_dim = attn_head_dim
+        all_head_dim = head_dim * self.num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+        assert all_head_dim == dim
+
+        self.q = nn.Linear(dim, all_head_dim, bias=False)
+        self.k = nn.Linear(dim, all_head_dim, bias=False)
+        self.v = nn.Linear(dim, all_head_dim, bias=False)
+
+        if qkv_bias:
+            self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
+            self.k_bias = nn.Parameter(torch.zeros(all_head_dim))
+            self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
+        else:
+            self.q_bias = None
+            self.k_bias = None
+            self.v_bias = None
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(all_head_dim, out_dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, k=None, v=None):
+        B, N, C = x.shape
+        N_k = k.shape[1]
+        N_v = v.shape[1]
+
+        q_bias, k_bias, v_bias = None, None, None
+        if self.q_bias is not None:
+            q_bias = self.q_bias
+            k_bias = self.k_bias
+            v_bias = self.v_bias
+
+        q = F.linear(input=x, weight=self.q.weight, bias=q_bias)
+        q = q.reshape(B, N, 1, self.num_heads,
+                      -1).permute(2, 0, 3, 1,
+                                  4).squeeze(0)  # (B, N_head, N_q, dim)
+
+        k = F.linear(input=k, weight=self.k.weight, bias=k_bias)
+        k = k.reshape(B, N_k, 1, self.num_heads, -1).permute(2, 0, 3, 1,
+                                                             4).squeeze(0)
+
+        v = F.linear(input=v, weight=self.v.weight, bias=v_bias)
+        v = v.reshape(B, N_v, 1, self.num_heads, -1).permute(2, 0, 3, 1,
+                                                             4).squeeze(0)
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))  # (B, N_head, N_q, N_k)
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        return x
+
+
+class AttentiveBlock(nn.Module):
+    r"""Attentive Block
+    Args:
+        dim (int): Number of input channels.
+        num_heads (int): Number of attention heads. Default: 8
+        qkv_bias (bool, optional):  If True, add a learnable bias to q, k, v.
+            Default: False.
+        qk_scale (float | None, optional): Override default qk scale of
+            head_dim ** -0.5 if set. Default: None.
+        drop (float, optional): Dropout rate. Default: 0.0.
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0.
+        drop_path (float | tuple[float], optional): Stochastic depth rate.
+            Default: 0.0.
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm.
+        attn_head_dim (int, optional): Dimension of attention head. Default: None.
+        out_dim (int, optional): Dimension of output. Default: None.
+    """
+
+    def __init__(self,
+                 dim,
+                 num_heads,
+                 qkv_bias=False,
+                 qk_scale=None,
+                 drop=0.,
+                 attn_drop=0.,
+                 drop_path=0.,
+                 norm_layer="LN",
+                 attn_head_dim=None,
+                 out_dim=None):
+        super().__init__()
+
+        self.norm1_q = build_norm_layer(dim, norm_layer, eps=1e-6)
+        self.norm1_k = build_norm_layer(dim, norm_layer, eps=1e-6)
+        self.norm1_v = build_norm_layer(dim, norm_layer, eps=1e-6)
+        self.cross_dcn = CrossAttention(dim,
+                                        num_heads=num_heads,
+                                        qkv_bias=qkv_bias,
+                                        qk_scale=qk_scale,
+                                        attn_drop=attn_drop,
+                                        proj_drop=drop,
+                                        attn_head_dim=attn_head_dim,
+                                        out_dim=out_dim)
+
+        self.drop_path = DropPath(
+            drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(self,
+                x_q,
+                x_kv,
+                pos_q,
+                pos_k,
+                bool_masked_pos,
+                rel_pos_bias=None):
+        x_q = self.norm1_q(x_q + pos_q)
+        x_k = self.norm1_k(x_kv + pos_k)
+        x_v = self.norm1_v(x_kv)
+
+        x = self.cross_dcn(x_q, k=x_k, v=x_v)
+
+        return x
+
+
+class AttentionPoolingBlock(AttentiveBlock):
+
+    def forward(self, x):
+        x_q = x.mean(1, keepdim=True)
+        x_kv = x
+        pos_q, pos_k = 0, 0
+        x = super().forward(x_q, x_kv, pos_q, pos_k,
+                            bool_masked_pos=None,
+                            rel_pos_bias=None)
+        x = x.squeeze(1)
+        return x
+
+
+class StemLayer(nn.Module):
+    r""" Stem layer of InternImage
+    Args:
+        in_chans (int): number of input channels
+        out_chans (int): number of output channels
+        act_layer (str): activation layer
+        norm_layer (str): normalization layer
+    """
+
+    def __init__(self,
+                 in_chans=3,
+                 out_chans=96,
+                 act_layer='GELU',
+                 norm_layer='BN'):
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_chans,
+                               out_chans // 2,
+                               kernel_size=3,
+                               stride=2,
+                               padding=1)
+        self.norm1 = build_norm_layer(out_chans // 2, norm_layer,
+                                      'channels_first', 'channels_first')
+        self.act = build_act_layer(act_layer)
+        self.conv2 = nn.Conv2d(out_chans // 2,
+                               out_chans,
+                               kernel_size=3,
+                               stride=2,
+                               padding=1)
+        self.norm2 = build_norm_layer(out_chans, norm_layer, 'channels_first',
+                                      'channels_last')
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.norm1(x)
+        x = self.act(x)
+        x = self.conv2(x)
+        x = self.norm2(x)
+        return x
+
+
+
+class DownsampleLayer(nn.Module):
+    r""" Downsample layer of InternImage
+    Args:
+        channels (int): number of input channels
+        norm_layer (str): normalization layer
+    """
+
+    def __init__(self, channels, norm_layer='LN'):
+        super().__init__()
+        self.conv = nn.Conv2d(channels,
+                              2 * channels,
+                              kernel_size=3,
+                              stride=2,
+                              padding=1,
+                              bias=False)
+        self.norm = build_norm_layer(2 * channels, norm_layer,
+                                     'channels_first', 'channels_first')
+
+
+    def forward(self, x, shape=None):
+        H, W = shape
+        N, HW, C = x.shape
+        x = x.view(N, H, W, C)
+        x = self.conv(x.permute(0, 3, 1, 2))
+        x = self.norm(x) # B C H W
+        H, W = x.size(2), x.size(3)
+        x = x.flatten(2).permute(0, 2, 1)
+
+        return x, (H, W)
+
+
+
+class MLPLayer(nn.Module):
+    r""" MLP layer of InternImage
+    Args:
+        in_features (int): number of input features
+        hidden_features (int): number of hidden features
+        out_features (int): number of output features
+        act_layer (str): activation layer
+        drop (float): dropout rate
+    """
+
+    def __init__(self,
+                 in_features,
+                 hidden_features=None,
+                 out_features=None,
+                 act_layer='GELU',
+                 mlp_fc2_bias=False,
+                 drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features, bias=True)
+        self.act = build_act_layer(act_layer)
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=mlp_fc2_bias)
+        self.drop = nn.Dropout(drop)
+
+
+    def forward(self, x, shape):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class InternImageLayer(nn.Module):
+    r""" Basic layer of InternImage
+    Args:
+        core_op (nn.Module): core operation of InternImage
+        channels (int): number of input channels
+        groups (list): Groups of each block.
+        mlp_ratio (float): ratio of mlp hidden features to input channels
+        drop (float): dropout rate
+        drop_path (float): drop path rate
+        act_layer (str): activation layer
+        norm_layer (str): normalization layer
+        post_norm (bool): whether to use post normalization
+        layer_scale (float): layer scale
+        offset_scale (float): offset scale
+        with_cp (bool): whether to use checkpoint
+    """
+
+    def __init__(self,
+                 core_op,
+                 channels,
+                 groups,
+                 mlp_ratio=4.,
+                 drop=0.,
+                 drop_path=0.,
+                 act_layer='GELU',
+                 norm_layer='LN',
+                 post_norm=False,
+                 layer_scale=None,
+                 offset_scale=1.0,
+                 with_cp=False,
+                 dcn_output_bias=False,
+                 mlp_fc2_bias=False,
+                 dw_kernel_size=None, # for InternImage-H/G
+                 res_post_norm=False, # for InternImage-H/G
+                 center_feature_scale=False): # for InternImage-H/G
+        super().__init__()
+        self.channels = channels
+        self.groups = groups
+        self.mlp_ratio = mlp_ratio
+        self.with_cp = with_cp
+
+        self.norm1 = build_norm_layer(channels, 'LN')
+        self.post_norm = post_norm
+        self.dcn = core_op(
+            channels=channels,
+            group=groups,
+            offset_scale=offset_scale,
+            dw_kernel_size=dw_kernel_size,
+            output_bias=dcn_output_bias,
+        )
+        self.drop_path = DropPath(drop_path) if drop_path > 0. \
+            else nn.Identity()
+        self.norm2 = build_norm_layer(channels, 'LN')
+        self.mlp = MLPLayer(in_features=channels,
+                            hidden_features=int(channels * mlp_ratio),
+                            act_layer=act_layer,
+                            drop=drop,
+                            mlp_fc2_bias=mlp_fc2_bias
+                            )
+        self.layer_scale = layer_scale is not None
+        if self.layer_scale:
+            self.gamma1 = nn.Parameter(layer_scale * torch.ones(channels),
+                                       requires_grad=True)
+            self.gamma2 = nn.Parameter(layer_scale * torch.ones(channels),
+                                       requires_grad=True)
+        self.res_post_norm = res_post_norm
+        if res_post_norm:
+            self.res_post_norm1 = build_norm_layer(channels, 'LN')
+            self.res_post_norm2 = build_norm_layer(channels, 'LN')
+    def forward(self, x, shape):
+
+        def _inner_forward(x, shape):
+            if not self.layer_scale:
+                if self.post_norm:
+                    x = x + self.drop_path(self.norm1(self.dcn(x, shape)))
+                    x = x + self.drop_path(self.norm2(self.mlp(x, shape)))
+                elif self.res_post_norm: # for InternImage-H/G
+                    x = x + self.drop_path(self.res_post_norm1(self.dcn(self.norm1(x), shape)))
+                    x = x + self.drop_path(self.res_post_norm2(self.mlp(self.norm2(x), shape)))
+
+                else:
+                    x = x + self.drop_path(self.dcn(self.norm1(x), shape,))
+                    x = x + self.drop_path(self.mlp(self.norm2(x), shape))
+                return x
+            if self.post_norm:
+                x = x + self.drop_path(self.gamma1 * self.norm1(self.dcn(x, shape)))
+                x = x + self.drop_path(self.gamma2 * self.norm2(self.mlp(x, shape)))
+            else:
+                x = x + self.drop_path(self.gamma1 * self.dcn(self.norm1(x), shape))
+                x = x + self.drop_path(self.gamma2 * self.mlp(self.norm2(x), shape))
+            return x
+
+        if self.with_cp and x.requires_grad:
+            x = checkpoint.checkpoint(_inner_forward, x, shape)
+        else:
+            x = _inner_forward(x, shape)
+
+        return x
+
+
+class InternImageBlock(nn.Module):
+    r""" Block of InternImage
+    Args:
+        core_op (nn.Module): core operation of InternImage
+        channels (int): number of input channels
+        depths (list): Depth of each block.
+        groups (list): Groups of each block.
+        mlp_ratio (float): ratio of mlp hidden features to input channels
+        drop (float): dropout rate
+        drop_path (float): drop path rate
+        act_layer (str): activation layer
+        norm_layer (str): normalization layer
+        post_norm (bool): whether to use post normalization
+        layer_scale (float): layer scale
+        offset_scale (float): offset scale
+        with_cp (bool): whether to use checkpoint
+    """
+
+    def __init__(self,
+                 core_op,
+                 channels,
+                 depth,
+                 groups,
+                 downsample=True,
+                 downsample_layer=DownsampleLayer,
+                 mlp_ratio=4.,
+                 drop=0.,
+                 drop_path=0.,
+                 act_layer='GELU',
+                 norm_layer='LN',
+                 post_norm=False,
+                 offset_scale=0.5,
+                 layer_scale=None,
+                 with_cp=False,
+                 dcn_output_bias=False,
+                 mlp_fc2_bias=False,
+                 dw_kernel_size=None, # for InternImage-H/G
+                 post_norm_block_ids=None, # for InternImage-H/G
+                 res_post_norm=False, # for InternImage-H/G
+                 center_feature_scale=False): # for InternImage-H/G
+        super().__init__()
+        self.channels = channels
+        self.depth = depth
+        self.post_norm = post_norm
+        self.center_feature_scale = center_feature_scale
+
+        self.blocks = nn.ModuleList([
+            InternImageLayer(
+                core_op=core_op,
+                channels=channels,
+                groups=groups,
+                mlp_ratio=mlp_ratio,
+                drop=drop,
+                drop_path=drop_path[i] if isinstance(
+                    drop_path, list) else drop_path,
+                act_layer=act_layer,
+                norm_layer=norm_layer,
+                post_norm=post_norm,
+                layer_scale=layer_scale,
+                offset_scale=offset_scale,
+                with_cp=with_cp,
+                dcn_output_bias=dcn_output_bias,
+                mlp_fc2_bias=mlp_fc2_bias,
+                dw_kernel_size=dw_kernel_size, # for InternImage-H/G
+                res_post_norm=res_post_norm, # for InternImage-H/G
+                center_feature_scale=center_feature_scale # for InternImage-H/G
+            ) for i in range(depth)
+        ])
+        if not self.post_norm or center_feature_scale:
+            self.norm = build_norm_layer(channels, 'LN')
+        self.post_norm_block_ids = post_norm_block_ids
+        if post_norm_block_ids is not None: # for InternImage-H/G
+            self.post_norms = nn.ModuleList(
+                [build_norm_layer(channels, 'LN', eps=1e-6) for _ in post_norm_block_ids]
+            )
+        self.downsample = downsample_layer(
+            channels=channels, norm_layer=norm_layer) if downsample else None
+
+
+    def forward(self, x, return_wo_downsample=False, shape=None):
+        for i, blk in enumerate(self.blocks):
+            x = blk(x, shape=shape)
+            if (self.post_norm_block_ids is not None) and (i in self.post_norm_block_ids):
+                index = self.post_norm_block_ids.index(i)
+                x = self.post_norms[index](x) # for InternImage-H/G
+        if not self.post_norm or self.center_feature_scale:
+            x = self.norm(x)
+        if return_wo_downsample:
+            x_ = x.clone()
+        if self.downsample is not None:
+            x, shape = self.downsample(x, shape=shape)
+
+        if return_wo_downsample:
+            return x, x_, shape
+        return x, shape
+
+
+class FlashInternImage(nn.Module):
+    r""" FlashInternImage
+        A PyTorch impl based on :
+            `InternImage: Exploring Large-Scale Vision Foundation Models with Deformable Convolutions`  -
+            https://arxiv.org/pdf/2103.14030
+            'DCNv4': TODO: add arxiv
+    Args:
+        core_op (str): Core operator. Default: 'DCNv4'
+        channels (int): Number of the first stage. Default: 64
+        depths (list): Depth of each block. Default: [3, 4, 18, 5]
+        groups (list): Groups of each block. Default: [3, 6, 12, 24]
+        num_classes (int): Number of classes. Default: 1000
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
+        drop_rate (float): Probability of an element to be zeroed. Default: 0.
+        drop_path_rate (float): Stochastic depth rate. Default: 0.2
+        act_layer (str): Activation layer. Default: 'GELU'
+        norm_layer (str): Normalization layer. Default: 'LN'
+        layer_scale (bool): Whether to use layer scale. Default: False
+        cls_scale (bool): Whether to use class scale. Default: False
+        with_cp (bool): Use checkpoint or not. Using checkpoint will save some
+        dw_kernel_size (int): Size of the dwconv. Default: None
+        use_clip_projector (bool): Whether to use clip projector. Default: False
+        level2_post_norm (bool): Whether to use level2 post norm. Default: False
+        level2_post_norm_block_ids (list): Indexes of post norm blocks. Default: None
+        res_post_norm (bool): Whether to use res post norm. Default: False
+        center_feature_scale (bool): Whether to use center feature scale. Default: False
+    """
+
+    def __init__(self,
+                 core_op='DCNv4',
+                 channels=64,
+                 depths=[3, 4, 18, 5],
+                 groups=[3, 6, 12, 24],
+                 num_classes=1000,
+                 mlp_ratio=4.,
+                 drop_rate=0.,
+                 drop_path_rate=0.2,
+                 drop_path_type='linear',
+                 act_layer='GELU',
+                 norm_layer='LN',
+                 layer_scale=None,
+                 offset_scale=0.5,
+                 post_norm=False,
+                 cls_scale=1.5,
+                 with_cp=False,
+                 mlp_fc2_bias=False,
+                 dcn_output_bias=False,
+                 dw_kernel_size=None,
+                 use_clip_projector=False, # for InternImage-H/G
+                 level2_post_norm=False, # for InternImage-H/G
+                 level2_post_norm_block_ids=None, # for InternImage-H/G
+                 res_post_norm=False, # for InternImage-H/G
+                 center_feature_scale=False, # for InternImage-H/G
+                 **kwargs):
+        super().__init__()
+        self.core_op = core_op
+        self.num_classes = num_classes
+        self.num_levels = len(depths)
+        self.depths = depths
+        self.channels = channels
+        self.num_features = int(channels * 2**(self.num_levels - 1))
+        self.post_norm = post_norm
+        self.mlp_ratio = mlp_ratio
+        self.use_clip_projector = use_clip_projector
+
+        self.level2_post_norm_block_ids = level2_post_norm_block_ids
+        print(f'using core type: {core_op}')
+        print(f'using activation layer: {act_layer}')
+        print(f'using main norm layer: {norm_layer}')
+        print(f'using dpr: {drop_path_type}, {drop_path_rate}')
+        print(f"level2_post_norm: {level2_post_norm}")
+        print(f"level2_post_norm_block_ids: {level2_post_norm_block_ids}")
+        print(f"res_post_norm: {res_post_norm}")
+
+        in_chans = 3
+        self.patch_embed = StemLayer(in_chans=in_chans,
+                                     out_chans=channels,
+                                     act_layer=act_layer,
+                                     norm_layer=norm_layer)
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = [
+            x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))
+        ]
+        if drop_path_type == 'uniform':
+            for i in range(len(dpr)):
+                dpr[i] = drop_path_rate
+   
+        self.levels = nn.ModuleList()
+        for i in range(self.num_levels):
+            post_norm_block_ids = level2_post_norm_block_ids if level2_post_norm and (
+                i == 2) else None # for InternImage-H/G
+
+            level = InternImageBlock(
+                core_op=getattr(DCNv4, core_op),
+                channels=int(channels * 2**i),
+                depth=depths[i],
+                groups=groups[i],
+                mlp_ratio=self.mlp_ratio,
+                drop=drop_rate,
+                drop_path=dpr[sum(depths[:i]):sum(depths[:i + 1])],
+                act_layer=act_layer,
+                norm_layer=norm_layer,
+                post_norm=post_norm,
+                downsample=(i < self.num_levels - 1),
+                downsample_layer = DownsampleLayer,
+                layer_scale=layer_scale,
+                offset_scale=offset_scale,
+                with_cp=with_cp,
+                mlp_fc2_bias=mlp_fc2_bias,
+                dcn_output_bias=dcn_output_bias,
+                dw_kernel_size=dw_kernel_size,  # for InternImage-H/G
+                post_norm_block_ids=post_norm_block_ids, # for InternImage-H/G
+                res_post_norm=res_post_norm, # for InternImage-H/G
+                center_feature_scale=center_feature_scale # for InternImage-H/G
+            )
+            self.levels.append(level)
+        
+        if not use_clip_projector: # for InternImage-T/S/B/L/XL
+            self.conv_head = nn.Sequential(
+                nn.Conv2d(self.num_features,
+                          int(self.num_features * cls_scale),
+                          kernel_size=1,
+                          bias=False),
+                build_norm_layer(int(self.num_features * cls_scale), 'BN',
+                                 'channels_first', 'channels_first'),
+                build_act_layer(act_layer))
+            self.head = nn.Linear(int(self.num_features * cls_scale), num_classes) \
+                if num_classes > 0 else nn.Identity()
+        else: # for InternImage-H/G
+            pretrain_embed_dim, _stride, attnpool_num_heads, clip_embed_dim = 1024, 2, 16, 768
+            self.dcnv3_head_x4 = nn.Sequential(
+                nn.Conv2d(in_channels=self.num_features,
+                          out_channels=pretrain_embed_dim * (_stride ** 2),
+                          kernel_size=1), nn.PixelShuffle(_stride))
+            self.dcnv3_head_x3 = nn.Conv2d(in_channels=self.num_features // 2,
+                                           out_channels=pretrain_embed_dim,
+                                           kernel_size=1)
+            self.clip_projector = AttentionPoolingBlock(
+                dim=pretrain_embed_dim,
+                num_heads=attnpool_num_heads,
+                qkv_bias=True,
+                qk_scale=None,
+                drop=0.,
+                attn_drop=0.,
+                norm_layer=norm_layer,
+                out_dim=clip_embed_dim)
+            self.fc_norm = build_norm_layer(clip_embed_dim, norm_layer, eps=1e-6)
+            self.head = nn.Linear(
+                clip_embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+            
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.num_layers = len(depths)
+        self.apply(self._init_weights)
+        self.apply(self._init_deform_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def _init_deform_weights(self, m):
+        if isinstance(m, getattr(DCNv4, self.core_op)):
+            m._reset_parameters()
+
+    @torch.jit.ignore
+    def lr_decay_keywards(self, decay_ratio=0.87):
+        lr_ratios = {}
+
+        # blocks
+        idx = 0
+        for i in range(4):
+            layer_num = 3 - i  # 3 2 1 0
+            for j in range(self.depths[layer_num]):
+                block_num = self.depths[layer_num] - j - 1
+                tag = 'levels.{}.blocks.{}.'.format(layer_num, block_num)
+                decay = 1.0 * (decay_ratio**idx)
+                lr_ratios[tag] = decay
+                idx += 1
+        # patch_embed (before stage-1)
+        lr_ratios["patch_embed"] = lr_ratios['levels.0.blocks.0.']
+        # levels.0.downsample (between stage-1 and stage-2)
+        lr_ratios["levels.0.downsample"] = lr_ratios['levels.1.blocks.0.']
+        lr_ratios["levels.0.norm"] = lr_ratios['levels.1.blocks.0.']
+        # levels.1.downsample (between stage-2 and stage-3)
+        lr_ratios["levels.1.downsample"] = lr_ratios['levels.2.blocks.0.']
+        lr_ratios["levels.1.norm"] = lr_ratios['levels.2.blocks.0.']
+        # levels.2.downsample (between stage-3 and stage-4)
+        lr_ratios["levels.2.downsample"] = lr_ratios['levels.3.blocks.0.']
+        lr_ratios["levels.2.norm"] = lr_ratios['levels.3.blocks.0.']
+        return lr_ratios
+
+    def forward_features(self, x):
+        x = self.patch_embed(x)
+        N, H, W, C = x.shape
+        x = x.view(N, H*W, C)
+
+        shape=(H, W)
+        seq_out = []
+        for level_idx, level in enumerate(self.levels):
+            old_shape = shape
+            x, shape = level(x, shape=shape)   
+        h, w = shape
+        x = x.view(N, h, w, -1)
+        x = self.conv_head(x.permute(0, 3, 1, 2))
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        return x
+
+    def forward_features_seq_out(self, x):
+        x = self.patch_embed(x)
+        N, H, W, C = x.shape
+        x = x.view(N, H*W, C)
+        shape=(H, W)
+        seq_out = []
+        for level_idx, level in enumerate(self.levels):
+            old_shape = shape
+            x, x_ , shape = level(x, return_wo_downsample=True, shape=shape)   
+            h, w= old_shape
+            seq_out.append(x_.reshape(N, h, w, -1).permute(0, 3, 1, 2))
+        return seq_out
+    
+    def forward_clip_projector(self, x): # for InternImage-H/G
+        xs = self.forward_features_seq_out(x)
+        x1, x2, x3, x4 = xs
+        
+        x1 = x1.permute(0, 3, 1, 2) # NHWC -> NCHW
+        x2 = x2.permute(0, 3, 1, 2) # NHWC -> NCHW
+        x3 = x3.permute(0, 3, 1, 2) # NHWC -> NCHW
+        x4 = x4.permute(0, 3, 1, 2) # NHWC -> NCHW
+
+        x4 = self.dcnv3_head_x4(x4)
+        x = x4
+        x3 = self.dcnv3_head_x3(x3)
+        x = x + x3
+
+        x = x.flatten(-2).transpose(1, 2).contiguous()
+        x = self.clip_projector(x)
+        x = self.fc_norm(x)
+        
+        return x
+    
+    def forward(self, x):
+        if self.use_clip_projector: # for InternImage-H/G
+            x = self.forward_clip_projector(x)
+        else: # for InternImage-T/S/B/L/XL
+            x = self.forward_features(x)
+        x = self.head(x)
+        return x
+
+
+

classification/models/intern_image.py

@@ -0,0 +1,759 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint as checkpoint
+from timm.models.layers import trunc_normal_, DropPath
+
+import torch.nn.functional as F
+
+
+class to_channels_first(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return x.permute(0, 3, 1, 2)
+
+
+class to_channels_last(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return x.permute(0, 2, 3, 1)
+
+
+def build_norm_layer(dim,
+                     norm_layer,
+                     in_format='channels_last',
+                     out_format='channels_last',
+                     eps=1e-6):
+    layers = []
+    if norm_layer == 'BN':
+        if in_format == 'channels_last':
+            layers.append(to_channels_first())
+        layers.append(nn.BatchNorm2d(dim))
+        if out_format == 'channels_last':
+            layers.append(to_channels_last())
+    elif norm_layer == 'LN':
+        if in_format == 'channels_first':
+            layers.append(to_channels_last())
+        layers.append(nn.LayerNorm(dim, eps=eps))
+        if out_format == 'channels_first':
+            layers.append(to_channels_first())
+    else:
+        raise NotImplementedError(
+            f'build_norm_layer does not support {norm_layer}')
+    return nn.Sequential(*layers)
+
+
+def build_act_layer(act_layer):
+    if act_layer == 'ReLU':
+        return nn.ReLU(inplace=True)
+    elif act_layer == 'SiLU':
+        return nn.SiLU(inplace=True)
+    elif act_layer == 'GELU':
+        return nn.GELU()
+
+    raise NotImplementedError(f'build_act_layer does not support {act_layer}')
+
+
+class CrossAttention(nn.Module):
+    r""" Cross Attention Module
+    Args:
+        dim (int): Number of input channels.
+        num_heads (int): Number of attention heads. Default: 8
+        qkv_bias (bool, optional):  If True, add a learnable bias to q, k, v.
+            Default: False.
+        qk_scale (float | None, optional): Override default qk scale of
+            head_dim ** -0.5 if set. Default: None.
+        attn_drop (float, optional): Dropout ratio of attention weight.
+            Default: 0.0
+        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
+        attn_head_dim (int, optional): Dimension of attention head.
+        out_dim (int, optional): Dimension of output.
+    """
+    
+    def __init__(self,
+                 dim,
+                 num_heads=8,
+                 qkv_bias=False,
+                 qk_scale=None,
+                 attn_drop=0.,
+                 proj_drop=0.,
+                 attn_head_dim=None,
+                 out_dim=None):
+        super().__init__()
+        if out_dim is None:
+            out_dim = dim
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        if attn_head_dim is not None:
+            head_dim = attn_head_dim
+        all_head_dim = head_dim * self.num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+        assert all_head_dim == dim
+
+        self.q = nn.Linear(dim, all_head_dim, bias=False)
+        self.k = nn.Linear(dim, all_head_dim, bias=False)
+        self.v = nn.Linear(dim, all_head_dim, bias=False)
+
+        if qkv_bias:
+            self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
+            self.k_bias = nn.Parameter(torch.zeros(all_head_dim))
+            self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
+        else:
+            self.q_bias = None
+            self.k_bias = None
+            self.v_bias = None
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(all_head_dim, out_dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, k=None, v=None):
+        B, N, C = x.shape
+        N_k = k.shape[1]
+        N_v = v.shape[1]
+
+        q_bias, k_bias, v_bias = None, None, None
+        if self.q_bias is not None:
+            q_bias = self.q_bias
+            k_bias = self.k_bias
+            v_bias = self.v_bias
+
+        q = F.linear(input=x, weight=self.q.weight, bias=q_bias)
+        q = q.reshape(B, N, 1, self.num_heads,
+                      -1).permute(2, 0, 3, 1,
+                                  4).squeeze(0)  # (B, N_head, N_q, dim)
+
+        k = F.linear(input=k, weight=self.k.weight, bias=k_bias)
+        k = k.reshape(B, N_k, 1, self.num_heads, -1).permute(2, 0, 3, 1,
+                                                             4).squeeze(0)
+
+        v = F.linear(input=v, weight=self.v.weight, bias=v_bias)
+        v = v.reshape(B, N_v, 1, self.num_heads, -1).permute(2, 0, 3, 1,
+                                                             4).squeeze(0)
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))  # (B, N_head, N_q, N_k)
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        return x
+
+
+class AttentiveBlock(nn.Module):
+    r"""Attentive Block
+    Args:
+        dim (int): Number of input channels.
+        num_heads (int): Number of attention heads. Default: 8
+        qkv_bias (bool, optional):  If True, add a learnable bias to q, k, v.
+            Default: False.
+        qk_scale (float | None, optional): Override default qk scale of
+            head_dim ** -0.5 if set. Default: None.
+        drop (float, optional): Dropout rate. Default: 0.0.
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0.
+        drop_path (float | tuple[float], optional): Stochastic depth rate.
+            Default: 0.0.
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm.
+        attn_head_dim (int, optional): Dimension of attention head. Default: None.
+        out_dim (int, optional): Dimension of output. Default: None.
+    """
+    
+    def __init__(self,
+                 dim,
+                 num_heads,
+                 qkv_bias=False,
+                 qk_scale=None,
+                 drop=0.,
+                 attn_drop=0.,
+                 drop_path=0.,
+                 norm_layer="LN",
+                 attn_head_dim=None,
+                 out_dim=None):
+        super().__init__()
+
+        self.norm1_q = build_norm_layer(dim, norm_layer, eps=1e-6)
+        self.norm1_k = build_norm_layer(dim, norm_layer, eps=1e-6)
+        self.norm1_v = build_norm_layer(dim, norm_layer, eps=1e-6)
+        self.cross_dcn = CrossAttention(dim,
+                                        num_heads=num_heads,
+                                        qkv_bias=qkv_bias,
+                                        qk_scale=qk_scale,
+                                        attn_drop=attn_drop,
+                                        proj_drop=drop,
+                                        attn_head_dim=attn_head_dim,
+                                        out_dim=out_dim)
+
+        self.drop_path = DropPath(
+            drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(self,
+                x_q,
+                x_kv,
+                pos_q,
+                pos_k,
+                bool_masked_pos,
+                rel_pos_bias=None):
+        x_q = self.norm1_q(x_q + pos_q)
+        x_k = self.norm1_k(x_kv + pos_k)
+        x_v = self.norm1_v(x_kv)
+
+        x = self.cross_dcn(x_q, k=x_k, v=x_v)
+
+        return x
+
+
+class AttentionPoolingBlock(AttentiveBlock):
+
+    def forward(self, x):
+        x_q = x.mean(1, keepdim=True)
+        x_kv = x
+        pos_q, pos_k = 0, 0
+        x = super().forward(x_q, x_kv, pos_q, pos_k,
+                            bool_masked_pos=None,
+                            rel_pos_bias=None)
+        x = x.squeeze(1)
+        return x
+
+
+class StemLayer(nn.Module):
+    r""" Stem layer of InternImage
+    Args:
+        in_chans (int): number of input channels
+        out_chans (int): number of output channels
+        act_layer (str): activation layer
+        norm_layer (str): normalization layer
+    """
+
+    def __init__(self,
+                 in_chans=3,
+                 out_chans=96,
+                 act_layer='GELU',
+                 norm_layer='BN'):
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_chans,
+                               out_chans // 2,
+                               kernel_size=3,
+                               stride=2,
+                               padding=1)
+        self.norm1 = build_norm_layer(out_chans // 2, norm_layer,
+                                      'channels_first', 'channels_first')
+        self.act = build_act_layer(act_layer)
+        self.conv2 = nn.Conv2d(out_chans // 2,
+                               out_chans,
+                               kernel_size=3,
+                               stride=2,
+                               padding=1)
+        self.norm2 = build_norm_layer(out_chans, norm_layer, 'channels_first',
+                                      'channels_last')
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.norm1(x)
+        x = self.act(x)
+        x = self.conv2(x)
+        x = self.norm2(x)
+        return x
+
+
+class DownsampleLayer(nn.Module):
+    r""" Downsample layer of InternImage
+    Args:
+        channels (int): number of input channels
+        norm_layer (str): normalization layer
+    """
+
+    def __init__(self, channels, norm_layer='LN'):
+        super().__init__()
+        self.conv = nn.Conv2d(channels,
+                              2 * channels,
+                              kernel_size=3,
+                              stride=2,
+                              padding=1,
+                              bias=False)
+        self.norm = build_norm_layer(2 * channels, norm_layer,
+                                     'channels_first', 'channels_last')
+
+    def forward(self, x):
+        x = self.conv(x.permute(0, 3, 1, 2))
+        x = self.norm(x)
+        return x
+
+
+class MLPLayer(nn.Module):
+    r""" MLP layer of InternImage
+    Args:
+        in_features (int): number of input features
+        hidden_features (int): number of hidden features
+        out_features (int): number of output features
+        act_layer (str): activation layer
+        drop (float): dropout rate
+    """
+
+    def __init__(self,
+                 in_features,
+                 hidden_features=None,
+                 out_features=None,
+                 act_layer='GELU',
+                 drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = build_act_layer(act_layer)
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class InternImageLayer(nn.Module):
+    r""" Basic layer of InternImage
+    Args:
+        core_op (nn.Module): core operation of InternImage
+        channels (int): number of input channels
+        groups (list): Groups of each block.
+        mlp_ratio (float): ratio of mlp hidden features to input channels
+        drop (float): dropout rate
+        drop_path (float): drop path rate
+        act_layer (str): activation layer
+        norm_layer (str): normalization layer
+        post_norm (bool): whether to use post normalization
+        layer_scale (float): layer scale
+        offset_scale (float): offset scale
+        with_cp (bool): whether to use checkpoint
+    """
+
+    def __init__(self,
+                 core_op,
+                 channels,
+                 groups,
+                 mlp_ratio=4.,
+                 drop=0.,
+                 drop_path=0.,
+                 act_layer='GELU',
+                 norm_layer='LN',
+                 post_norm=False,
+                 layer_scale=None,
+                 offset_scale=1.0,
+                 with_cp=False,
+                 dw_kernel_size=None, # for InternImage-H/G
+                 res_post_norm=False, # for InternImage-H/G
+                 center_feature_scale=False): # for InternImage-H/G
+        super().__init__()
+        self.channels = channels
+        self.groups = groups
+        self.mlp_ratio = mlp_ratio
+        self.with_cp = with_cp
+
+        self.norm1 = build_norm_layer(channels, 'LN')
+        self.post_norm = post_norm
+        self.dcn = core_op(
+            channels=channels,
+            kernel_size=3,
+            stride=1,
+            pad=1,
+            dilation=1,
+            group=groups,
+            offset_scale=offset_scale,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+            dw_kernel_size=dw_kernel_size, # for InternImage-H/G
+            center_feature_scale=center_feature_scale) # for InternImage-H/G
+        self.drop_path = DropPath(drop_path) if drop_path > 0. \
+            else nn.Identity()
+        self.norm2 = build_norm_layer(channels, 'LN')
+        self.mlp = MLPLayer(in_features=channels,
+                            hidden_features=int(channels * mlp_ratio),
+                            act_layer=act_layer,
+                            drop=drop)
+        self.layer_scale = layer_scale is not None
+        if self.layer_scale:
+            self.gamma1 = nn.Parameter(layer_scale * torch.ones(channels),
+                                       requires_grad=True)
+            self.gamma2 = nn.Parameter(layer_scale * torch.ones(channels),
+                                       requires_grad=True)
+        self.res_post_norm = res_post_norm
+        if res_post_norm:
+            self.res_post_norm1 = build_norm_layer(channels, 'LN')
+            self.res_post_norm2 = build_norm_layer(channels, 'LN')
+
+    def forward(self, x):
+
+        def _inner_forward(x):
+            if not self.layer_scale:
+                if self.post_norm:
+                    x = x + self.drop_path(self.norm1(self.dcn(x)))
+                    x = x + self.drop_path(self.norm2(self.mlp(x)))
+                elif self.res_post_norm: # for InternImage-H/G
+                    x = x + self.drop_path(self.res_post_norm1(self.dcn(self.norm1(x))))
+                    x = x + self.drop_path(self.res_post_norm2(self.mlp(self.norm2(x))))
+                else:
+                    x = x + self.drop_path(self.dcn(self.norm1(x)))
+                    x = x + self.drop_path(self.mlp(self.norm2(x)))
+                return x
+            if self.post_norm:
+                x = x + self.drop_path(self.gamma1 * self.norm1(self.dcn(x)))
+                x = x + self.drop_path(self.gamma2 * self.norm2(self.mlp(x)))
+            else:
+                x = x + self.drop_path(self.gamma1 * self.dcn(self.norm1(x)))
+                x = x + self.drop_path(self.gamma2 * self.mlp(self.norm2(x)))
+            return x
+
+        if self.with_cp and x.requires_grad:
+            x = checkpoint.checkpoint(_inner_forward, x)
+        else:
+            x = _inner_forward(x)
+        return x
+
+
+class InternImageBlock(nn.Module):
+    r""" Block of InternImage
+    Args:
+        core_op (nn.Module): core operation of InternImage
+        channels (int): number of input channels
+        depths (list): Depth of each block.
+        groups (list): Groups of each block.
+        mlp_ratio (float): ratio of mlp hidden features to input channels
+        drop (float): dropout rate
+        drop_path (float): drop path rate
+        act_layer (str): activation layer
+        norm_layer (str): normalization layer
+        post_norm (bool): whether to use post normalization
+        layer_scale (float): layer scale
+        offset_scale (float): offset scale
+        with_cp (bool): whether to use checkpoint
+    """
+
+    def __init__(self,
+                 core_op,
+                 channels,
+                 depth,
+                 groups,
+                 downsample=True,
+                 mlp_ratio=4.,
+                 drop=0.,
+                 drop_path=0.,
+                 act_layer='GELU',
+                 norm_layer='LN',
+                 post_norm=False,
+                 offset_scale=1.0,
+                 layer_scale=None,
+                 with_cp=False,
+                 dw_kernel_size=None, # for InternImage-H/G
+                 post_norm_block_ids=None, # for InternImage-H/G
+                 res_post_norm=False, # for InternImage-H/G
+                 center_feature_scale=False): # for InternImage-H/G
+        super().__init__()
+        self.channels = channels
+        self.depth = depth
+        self.post_norm = post_norm
+        self.center_feature_scale = center_feature_scale
+
+        self.blocks = nn.ModuleList([
+            InternImageLayer(
+                core_op=core_op,
+                channels=channels,
+                groups=groups,
+                mlp_ratio=mlp_ratio,
+                drop=drop,
+                drop_path=drop_path[i] if isinstance(
+                    drop_path, list) else drop_path,
+                act_layer=act_layer,
+                norm_layer=norm_layer,
+                post_norm=post_norm,
+                layer_scale=layer_scale,
+                offset_scale=offset_scale,
+                with_cp=with_cp,
+                dw_kernel_size=dw_kernel_size, # for InternImage-H/G
+                res_post_norm=res_post_norm, # for InternImage-H/G
+                center_feature_scale=center_feature_scale # for InternImage-H/G
+            ) for i in range(depth)
+        ])
+        if not self.post_norm or center_feature_scale:
+            self.norm = build_norm_layer(channels, 'LN')
+        self.post_norm_block_ids = post_norm_block_ids
+        if post_norm_block_ids is not None: # for InternImage-H/G
+            self.post_norms = nn.ModuleList(
+                [build_norm_layer(channels, 'LN', eps=1e-6) for _ in post_norm_block_ids]
+            )
+        self.downsample = DownsampleLayer(
+            channels=channels, norm_layer=norm_layer) if downsample else None
+
+    def forward(self, x, return_wo_downsample=False):
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if (self.post_norm_block_ids is not None) and (i in self.post_norm_block_ids):
+                index = self.post_norm_block_ids.index(i)
+                x = self.post_norms[index](x) # for InternImage-H/G
+        if not self.post_norm or self.center_feature_scale:
+            x = self.norm(x)
+        if return_wo_downsample:
+            x_ = x
+        if self.downsample is not None:
+            x = self.downsample(x)
+
+        if return_wo_downsample:
+            return x, x_
+        return x
+
+
+class InternImage(nn.Module):
+    r""" InternImage
+        A PyTorch impl of : `InternImage: Exploring Large-Scale Vision Foundation Models with Deformable Convolutions`  -
+          https://arxiv.org/pdf/2103.14030
+    Args:
+        core_op (str): Core operator. Default: 'DCNv3'
+        channels (int): Number of the first stage. Default: 64
+        depths (list): Depth of each block. Default: [3, 4, 18, 5]
+        groups (list): Groups of each block. Default: [3, 6, 12, 24]
+        num_classes (int): Number of classes. Default: 1000
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
+        drop_rate (float): Probability of an element to be zeroed. Default: 0.
+        drop_path_rate (float): Stochastic depth rate. Default: 0.
+        act_layer (str): Activation layer. Default: 'GELU'
+        norm_layer (str): Normalization layer. Default: 'LN'
+        layer_scale (bool): Whether to use layer scale. Default: False
+        cls_scale (bool): Whether to use class scale. Default: False
+        with_cp (bool): Use checkpoint or not. Using checkpoint will save some
+        dw_kernel_size (int): Size of the dwconv. Default: None
+        use_clip_projector (bool): Whether to use clip projector. Default: False
+        level2_post_norm (bool): Whether to use level2 post norm. Default: False
+        level2_post_norm_block_ids (list): Indexes of post norm blocks. Default: None
+        res_post_norm (bool): Whether to use res post norm. Default: False
+        center_feature_scale (bool): Whether to use center feature scale. Default: False
+    """
+
+    def __init__(self,
+                 core_op='DCNv3',
+                 channels=64,
+                 depths=[3, 4, 18, 5],
+                 groups=[3, 6, 12, 24],
+                 num_classes=1000,
+                 mlp_ratio=4.,
+                 drop_rate=0.,
+                 drop_path_rate=0.2,
+                 drop_path_type='linear',
+                 act_layer='GELU',
+                 norm_layer='LN',
+                 layer_scale=None,
+                 offset_scale=1.0,
+                 post_norm=False,
+                 cls_scale=1.5,
+                 with_cp=False,
+                 dw_kernel_size=None, # for InternImage-H/G
+                 use_clip_projector=False, # for InternImage-H/G
+                 level2_post_norm=False, # for InternImage-H/G
+                 level2_post_norm_block_ids=None, # for InternImage-H/G
+                 res_post_norm=False, # for InternImage-H/G
+                 center_feature_scale=False, # for InternImage-H/G
+                 **kwargs):
+        super().__init__()
+        self.core_op = core_op
+        self.num_classes = num_classes
+        self.num_levels = len(depths)
+        self.depths = depths
+        self.channels = channels
+        self.num_features = int(channels * 2**(self.num_levels - 1))
+        self.post_norm = post_norm
+        self.mlp_ratio = mlp_ratio
+        self.use_clip_projector = use_clip_projector
+        self.level2_post_norm_block_ids = level2_post_norm_block_ids
+        print(f'using core type: {core_op}')
+        print(f'using activation layer: {act_layer}')
+        print(f'using main norm layer: {norm_layer}')
+        print(f'using dpr: {drop_path_type}, {drop_path_rate}')
+        print(f"level2_post_norm: {level2_post_norm}")
+        print(f"level2_post_norm_block_ids: {level2_post_norm_block_ids}")
+        print(f"res_post_norm: {res_post_norm}")
+
+        in_chans = 3
+        self.patch_embed = StemLayer(in_chans=in_chans,
+                                     out_chans=channels,
+                                     act_layer=act_layer,
+                                     norm_layer=norm_layer)
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = [
+            x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))
+        ]
+        if drop_path_type == 'uniform':
+            for i in range(len(dpr)):
+                dpr[i] = drop_path_rate
+
+        self.levels = nn.ModuleList()
+        for i in range(self.num_levels):
+            post_norm_block_ids = level2_post_norm_block_ids if level2_post_norm and (
+                i == 2) else None # for InternImage-H/G
+            from ops_offset import modules as opsm
+            level = InternImageBlock(
+                core_op=getattr(opsm, core_op),
+                channels=int(channels * 2**i),
+                depth=depths[i],
+                groups=groups[i],
+                mlp_ratio=self.mlp_ratio,
+                drop=drop_rate,
+                drop_path=dpr[sum(depths[:i]):sum(depths[:i + 1])],
+                act_layer=act_layer,
+                norm_layer=norm_layer,
+                post_norm=post_norm,
+                downsample=(i < self.num_levels - 1),
+                layer_scale=layer_scale,
+                offset_scale=offset_scale,
+                with_cp=with_cp,
+                dw_kernel_size=dw_kernel_size,  # for InternImage-H/G
+                post_norm_block_ids=post_norm_block_ids, # for InternImage-H/G
+                res_post_norm=res_post_norm, # for InternImage-H/G
+                center_feature_scale=center_feature_scale # for InternImage-H/G
+            )
+            self.levels.append(level)
+        
+        if not use_clip_projector: # for InternImage-T/S/B/L/XL
+            self.conv_head = nn.Sequential(
+                nn.Conv2d(self.num_features,
+                          int(self.num_features * cls_scale),
+                          kernel_size=1,
+                          bias=False),
+                build_norm_layer(int(self.num_features * cls_scale), 'BN',
+                                 'channels_first', 'channels_first'),
+                build_act_layer(act_layer))
+            self.head = nn.Linear(int(self.num_features * cls_scale), num_classes) \
+                if num_classes > 0 else nn.Identity()
+        else: # for InternImage-H/G
+            pretrain_embed_dim, _stride, attnpool_num_heads, clip_embed_dim = 1024, 2, 16, 768
+            self.dcnv3_head_x4 = nn.Sequential(
+                nn.Conv2d(in_channels=self.num_features,
+                          out_channels=pretrain_embed_dim * (_stride ** 2),
+                          kernel_size=1), nn.PixelShuffle(_stride))
+            self.dcnv3_head_x3 = nn.Conv2d(in_channels=self.num_features // 2,
+                                           out_channels=pretrain_embed_dim,
+                                           kernel_size=1)
+            self.clip_projector = AttentionPoolingBlock(
+                dim=pretrain_embed_dim,
+                num_heads=attnpool_num_heads,
+                qkv_bias=True,
+                qk_scale=None,
+                drop=0.,
+                attn_drop=0.,
+                norm_layer=norm_layer,
+                out_dim=clip_embed_dim)
+            self.fc_norm = build_norm_layer(clip_embed_dim, norm_layer, eps=1e-6)
+            self.head = nn.Linear(
+                clip_embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+            
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.num_layers = len(depths)
+        self.apply(self._init_weights)
+        self.apply(self._init_deform_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def _init_deform_weights(self, m):
+        from ops_offset import modules as opsm
+        if isinstance(m, getattr(opsm, self.core_op)):
+            m._reset_parameters()
+
+    @torch.jit.ignore
+    def lr_decay_keywards(self, decay_ratio=0.87):
+        lr_ratios = {}
+
+        # blocks
+        idx = 0
+        for i in range(4):
+            layer_num = 3 - i  # 3 2 1 0
+            for j in range(self.depths[layer_num]):
+                block_num = self.depths[layer_num] - j - 1
+                tag = 'levels.{}.blocks.{}.'.format(layer_num, block_num)
+                decay = 1.0 * (decay_ratio**idx)
+                lr_ratios[tag] = decay
+                idx += 1
+        # patch_embed (before stage-1)
+        lr_ratios["patch_embed"] = lr_ratios['levels.0.blocks.0.']
+        # levels.0.downsample (between stage-1 and stage-2)
+        lr_ratios["levels.0.downsample"] = lr_ratios['levels.1.blocks.0.']
+        lr_ratios["levels.0.norm"] = lr_ratios['levels.1.blocks.0.']
+        # levels.1.downsample (between stage-2 and stage-3)
+        lr_ratios["levels.1.downsample"] = lr_ratios['levels.2.blocks.0.']
+        lr_ratios["levels.1.norm"] = lr_ratios['levels.2.blocks.0.']
+        # levels.2.downsample (between stage-3 and stage-4)
+        lr_ratios["levels.2.downsample"] = lr_ratios['levels.3.blocks.0.']
+        lr_ratios["levels.2.norm"] = lr_ratios['levels.3.blocks.0.']
+        return lr_ratios
+
+    def forward_features(self, x):
+        x = self.patch_embed(x)
+        x = self.pos_drop(x)
+
+        for level in self.levels:
+            x = level(x)
+
+        x = self.conv_head(x.permute(0, 3, 1, 2))
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        return x
+
+    def forward_features_seq_out(self, x):
+        x = self.patch_embed(x)
+        x = self.pos_drop(x)
+
+        seq_out = []
+        for level in self.levels:
+            x, x_ = level(x, return_wo_downsample=True)
+            seq_out.append(x_)
+        return seq_out
+    
+    def forward_clip_projector(self, x): # for InternImage-H/G
+        xs = self.forward_features_seq_out(x)
+        x1, x2, x3, x4 = xs
+        
+        x1 = x1.permute(0, 3, 1, 2) # NHWC -> NCHW
+        x2 = x2.permute(0, 3, 1, 2) # NHWC -> NCHW
+        x3 = x3.permute(0, 3, 1, 2) # NHWC -> NCHW
+        x4 = x4.permute(0, 3, 1, 2) # NHWC -> NCHW
+
+        x4 = self.dcnv3_head_x4(x4)
+        x = x4
+        x3 = self.dcnv3_head_x3(x3)
+        x = x + x3
+
+        x = x.flatten(-2).transpose(1, 2).contiguous()
+        x = self.clip_projector(x)
+        x = self.fc_norm(x)
+        
+        return x
+    
+    def forward(self, x):
+        if self.use_clip_projector: # for InternImage-H/G
+            x = self.forward_clip_projector(x)
+        else: # for InternImage-T/S/B/L/XL
+            x = self.forward_features(x)
+        x = self.head(x)
+        return x

classification/ops_dcnv3/functions/__init__.py

@@ -0,0 +1,7 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+from .dcnv3_func import DCNv3Function, dcnv3_core_pytorch

classification/ops_dcnv3/functions/dcnv3_func.py

@@ -0,0 +1,220 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import torch
+import torch.nn.functional as F
+from torch.autograd import Function
+from torch.autograd.function import once_differentiable
+from torch.cuda.amp import custom_bwd, custom_fwd
+import DCNv3
+
+import pkg_resources
+dcn_version = float(pkg_resources.get_distribution('DCNv3').version)
+
+
+class DCNv3Function(Function):
+    @staticmethod
+    @custom_fwd
+    def forward(
+            ctx, input, offset, mask,
+            kernel_h, kernel_w, stride_h, stride_w,
+            pad_h, pad_w, dilation_h, dilation_w,
+            group, group_channels, offset_scale, im2col_step, remove_center):
+        ctx.kernel_h = kernel_h
+        ctx.kernel_w = kernel_w
+        ctx.stride_h = stride_h
+        ctx.stride_w = stride_w
+        ctx.pad_h = pad_h
+        ctx.pad_w = pad_w
+        ctx.dilation_h = dilation_h
+        ctx.dilation_w = dilation_w
+        ctx.group = group
+        ctx.group_channels = group_channels
+        ctx.offset_scale = offset_scale
+        ctx.im2col_step = im2col_step
+        ctx.remove_center = remove_center
+
+        args = [
+            input, offset, mask, kernel_h,
+            kernel_w, stride_h, stride_w, pad_h,
+            pad_w, dilation_h, dilation_w, group,
+            group_channels, offset_scale, ctx.im2col_step
+        ]
+        if remove_center or dcn_version > 1.0:
+            args.append(remove_center)
+
+        output = DCNv3.dcnv3_forward(*args)
+        ctx.save_for_backward(input, offset, mask)
+
+        return output
+
+    @staticmethod
+    @once_differentiable
+    @custom_bwd
+    def backward(ctx, grad_output):
+        input, offset, mask = ctx.saved_tensors
+
+        args = [
+            input, offset, mask, ctx.kernel_h,
+            ctx.kernel_w, ctx.stride_h, ctx.stride_w, ctx.pad_h,
+            ctx.pad_w, ctx.dilation_h, ctx.dilation_w, ctx.group,
+            ctx.group_channels, ctx.offset_scale, grad_output.contiguous(), ctx.im2col_step
+        ]
+        if ctx.remove_center or dcn_version > 1.0:
+            args.append(ctx.remove_center)
+
+        grad_input, grad_offset, grad_mask = \
+            DCNv3.dcnv3_backward(*args)
+
+        return grad_input, grad_offset, grad_mask, \
+            None, None, None, None, None, None, None, None, None, None, None, None, None
+
+    @staticmethod
+    def symbolic(g, input, offset, mask, kernel_h, kernel_w, stride_h,
+                 stride_w, pad_h, pad_w, dilation_h, dilation_w, group,
+                 group_channels, offset_scale, im2col_step, remove_center):
+        """Symbolic function for mmdeploy::DCNv3.
+
+        Returns:
+            DCNv3 op for onnx.
+        """
+        return g.op(
+            'mmdeploy::TRTDCNv3',
+            input,
+            offset,
+            mask,
+            kernel_h_i=int(kernel_h),
+            kernel_w_i=int(kernel_w),
+            stride_h_i=int(stride_h),
+            stride_w_i=int(stride_w),
+            pad_h_i=int(pad_h),
+            pad_w_i=int(pad_w),
+            dilation_h_i=int(dilation_h),
+            dilation_w_i=int(dilation_w),
+            group_i=int(group),
+            group_channels_i=int(group_channels),
+            offset_scale_f=float(offset_scale),
+            im2col_step_i=int(im2col_step),
+            remove_center=int(remove_center),
+        )
+
+
+def _get_reference_points(spatial_shapes, device, kernel_h, kernel_w, dilation_h, dilation_w, pad_h=0, pad_w=0, stride_h=1, stride_w=1):
+    _, H_, W_, _ = spatial_shapes
+    H_out = (H_ - (dilation_h * (kernel_h - 1) + 1)) // stride_h + 1
+    W_out = (W_ - (dilation_w * (kernel_w - 1) + 1)) // stride_w + 1
+
+    ref_y, ref_x = torch.meshgrid(
+        torch.linspace(
+            # pad_h + 0.5,
+            # H_ - pad_h - 0.5,
+            (dilation_h * (kernel_h - 1)) // 2 + 0.5,
+            (dilation_h * (kernel_h - 1)) // 2 + 0.5 + (H_out - 1) * stride_h,
+            H_out,
+            dtype=torch.float32,
+            device=device),
+        torch.linspace(
+            # pad_w + 0.5,
+            # W_ - pad_w - 0.5,
+            (dilation_w * (kernel_w - 1)) // 2 + 0.5,
+            (dilation_w * (kernel_w - 1)) // 2 + 0.5 + (W_out - 1) * stride_w,
+            W_out,
+            dtype=torch.float32,
+            device=device))
+    ref_y = ref_y.reshape(-1)[None] / H_
+    ref_x = ref_x.reshape(-1)[None] / W_
+
+    ref = torch.stack((ref_x, ref_y), -1).reshape(
+        1, H_out, W_out, 1, 2)
+
+    return ref
+
+
+def _generate_dilation_grids(spatial_shapes, kernel_h, kernel_w, dilation_h, dilation_w, group, device):
+    _, H_, W_, _ = spatial_shapes
+    points_list = []
+    x, y = torch.meshgrid(
+        torch.linspace(
+            -((dilation_w * (kernel_w - 1)) // 2),
+            -((dilation_w * (kernel_w - 1)) // 2) + (kernel_w - 1) * dilation_w,
+            kernel_w,
+            dtype=torch.float32,
+            device=device),
+        torch.linspace(
+            -((dilation_h * (kernel_h - 1)) // 2),
+            -((dilation_h * (kernel_h - 1)) // 2) + (kernel_h - 1) * dilation_h,
+            kernel_h,
+            dtype=torch.float32,
+            device=device))
+
+    points_list.extend([x / W_, y / H_])
+    grid = torch.stack(points_list, -1).reshape(-1, 1, 2).\
+        repeat(1, group, 1).permute(1, 0, 2)
+    grid = grid.reshape(1, 1, 1, group * kernel_h * kernel_w, 2)
+
+    return grid
+
+
+def remove_center_sampling_locations(sampling_locations, kernel_w, kernel_h):
+    idx = list(range(sampling_locations.shape[-2]))
+    C = (kernel_w * kernel_h - 1)//2
+    idx = [i for i in idx if i != C and (i-C) % (C*2+1) != 0]
+    sampling_locations = sampling_locations[:,:,:,idx, :]
+    return sampling_locations
+
+def dcnv3_core_pytorch(
+        input, offset, mask, kernel_h,
+        kernel_w, stride_h, stride_w, pad_h,
+        pad_w, dilation_h, dilation_w, group,
+        group_channels, offset_scale, remove_center):
+    # for debug and test only,
+    # need to use cuda version instead
+
+    if remove_center and (kernel_h % 2 == 0 or kernel_w % 2 == 0 or kernel_w != kernel_h):
+        raise ValueError('remove_center is only compatible with square odd kernel size.')
+
+    input = F.pad(
+        input,
+        [0, 0, pad_h, pad_h, pad_w, pad_w])
+    N_, H_in, W_in, _ = input.shape
+    _, H_out, W_out, _ = offset.shape
+
+    ref = _get_reference_points(
+        input.shape, input.device, kernel_h, kernel_w, dilation_h, dilation_w, pad_h, pad_w, stride_h, stride_w)
+    grid = _generate_dilation_grids(
+        input.shape, kernel_h, kernel_w, dilation_h, dilation_w, group, input.device)
+    spatial_norm = torch.tensor([W_in, H_in]).reshape(1, 1, 1, 2).\
+        repeat(1, 1, 1, group*(kernel_h*kernel_w-remove_center)).to(input.device)
+
+    sampling_locations = (ref + grid * offset_scale).repeat(N_, 1, 1, 1, 1)
+    if remove_center:
+        sampling_locations = remove_center_sampling_locations(sampling_locations, kernel_w=kernel_w, kernel_h=kernel_h)
+    sampling_locations = sampling_locations.flatten(3, 4)
+    sampling_locations = sampling_locations + offset * offset_scale / spatial_norm
+
+    P_ = kernel_h * kernel_w - remove_center
+    sampling_grids = 2 * sampling_locations - 1
+    # N_, H_in, W_in, group*group_channels -> N_, H_in*W_in, group*group_channels -> N_, group*group_channels, H_in*W_in -> N_*group, group_channels, H_in, W_in
+    input_ = input.view(N_, H_in*W_in, group*group_channels).transpose(1, 2).\
+        reshape(N_*group, group_channels, H_in, W_in)
+    # N_, H_out, W_out, group*P_*2 -> N_, H_out*W_out, group, P_, 2 -> N_, group, H_out*W_out, P_, 2 -> N_*group, H_out*W_out, P_, 2
+    sampling_grid_ = sampling_grids.view(N_, H_out*W_out, group, P_, 2).transpose(1, 2).\
+        flatten(0, 1)
+    # N_*group, group_channels, H_out*W_out, P_
+    sampling_input_ = F.grid_sample(
+        input_, sampling_grid_, mode='bilinear', padding_mode='zeros', align_corners=False)
+
+    # (N_, H_out, W_out, group*P_) -> N_, H_out*W_out, group, P_ -> (N_, group, H_out*W_out, P_) -> (N_*group, 1, H_out*W_out, P_)
+    mask = mask.view(N_, H_out*W_out, group, P_).transpose(1, 2).\
+        reshape(N_*group, 1, H_out*W_out, P_)
+    output = (sampling_input_ * mask).sum(-1).view(N_,
+                                                   group*group_channels, H_out*W_out)
+
+    return output.transpose(1, 2).reshape(N_, H_out, W_out, -1).contiguous()

classification/ops_dcnv3/make.sh

@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+python setup.py build install

classification/ops_dcnv3/modules/__init__.py

@@ -0,0 +1,7 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+from .dcnv3 import DCNv3, DCNv3_pytorch

classification/ops_dcnv3/modules/dcnv3.py

@@ -0,0 +1,357 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import warnings
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn.init import xavier_uniform_, constant_
+from ..functions import DCNv3Function, dcnv3_core_pytorch
+
+
+class to_channels_first(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return x.permute(0, 3, 1, 2)
+
+
+class to_channels_last(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return x.permute(0, 2, 3, 1)
+
+
+def build_norm_layer(dim,
+                     norm_layer,
+                     in_format='channels_last',
+                     out_format='channels_last',
+                     eps=1e-6):
+    layers = []
+    if norm_layer == 'BN':
+        if in_format == 'channels_last':
+            layers.append(to_channels_first())
+        layers.append(nn.BatchNorm2d(dim))
+        if out_format == 'channels_last':
+            layers.append(to_channels_last())
+    elif norm_layer == 'LN':
+        if in_format == 'channels_first':
+            layers.append(to_channels_last())
+        layers.append(nn.LayerNorm(dim, eps=eps))
+        if out_format == 'channels_first':
+            layers.append(to_channels_first())
+    else:
+        raise NotImplementedError(
+            f'build_norm_layer does not support {norm_layer}')
+    return nn.Sequential(*layers)
+
+
+def build_act_layer(act_layer):
+    if act_layer == 'ReLU':
+        return nn.ReLU(inplace=True)
+    elif act_layer == 'SiLU':
+        return nn.SiLU(inplace=True)
+    elif act_layer == 'GELU':
+        return nn.GELU()
+
+    raise NotImplementedError(f'build_act_layer does not support {act_layer}')
+
+
+def _is_power_of_2(n):
+    if (not isinstance(n, int)) or (n < 0):
+        raise ValueError(
+            "invalid input for _is_power_of_2: {} (type: {})".format(n, type(n)))
+
+    return (n & (n - 1) == 0) and n != 0
+
+
+class CenterFeatureScaleModule(nn.Module):
+    def forward(self,
+                query,
+                center_feature_scale_proj_weight,
+                center_feature_scale_proj_bias):
+        center_feature_scale = F.linear(query,
+                                        weight=center_feature_scale_proj_weight,
+                                        bias=center_feature_scale_proj_bias).sigmoid()
+        return center_feature_scale
+
+
+class DCNv3_pytorch(nn.Module):
+    def __init__(
+            self,
+            channels=64,
+            kernel_size=3,
+            dw_kernel_size=None,
+            stride=1,
+            pad=1,
+            dilation=1,
+            group=4,
+            offset_scale=1.0,
+            act_layer='GELU',
+            norm_layer='LN',
+            center_feature_scale=False,
+            remove_center=False,
+    ):
+        """
+        DCNv3 Module
+        :param channels
+        :param kernel_size
+        :param stride
+        :param pad
+        :param dilation
+        :param group
+        :param offset_scale
+        :param act_layer
+        :param norm_layer
+        """
+        super().__init__()
+        if channels % group != 0:
+            raise ValueError(
+                f'channels must be divisible by group, but got {channels} and {group}')
+        _d_per_group = channels // group
+        dw_kernel_size = dw_kernel_size if dw_kernel_size is not None else kernel_size
+        # you'd better set _d_per_group to a power of 2 which is more efficient in our CUDA implementation
+        if not _is_power_of_2(_d_per_group):
+            warnings.warn(
+                "You'd better set channels in DCNv3 to make the dimension of each attention head a power of 2 "
+                "which is more efficient in our CUDA implementation.")
+
+        self.offset_scale = offset_scale
+        self.channels = channels
+        self.kernel_size = kernel_size
+        self.dw_kernel_size = dw_kernel_size
+        self.stride = stride
+        self.dilation = dilation
+        self.pad = pad
+        self.group = group
+        self.group_channels = channels // group
+        self.offset_scale = offset_scale
+        self.center_feature_scale = center_feature_scale
+        self.remove_center = int(remove_center)
+
+        self.dw_conv = nn.Sequential(
+            nn.Conv2d(
+                channels,
+                channels,
+                kernel_size=dw_kernel_size,
+                stride=1,
+                padding=(dw_kernel_size - 1) // 2,
+                groups=channels),
+            build_norm_layer(
+                channels,
+                norm_layer,
+                'channels_first',
+                'channels_last'),
+            build_act_layer(act_layer))
+        self.offset = nn.Linear(
+            channels,
+            group * (kernel_size * kernel_size - remove_center) * 2)
+        self.mask = nn.Linear(
+            channels,
+            group * (kernel_size * kernel_size - remove_center))
+        self.input_proj = nn.Linear(channels, channels)
+        self.output_proj = nn.Linear(channels, channels)
+        self._reset_parameters()
+        
+        if center_feature_scale:
+            self.center_feature_scale_proj_weight = nn.Parameter(
+                torch.zeros((group, channels), dtype=torch.float))
+            self.center_feature_scale_proj_bias = nn.Parameter(
+                torch.tensor(0.0, dtype=torch.float).view((1,)).repeat(group, ))
+            self.center_feature_scale_module = CenterFeatureScaleModule()
+
+    def _reset_parameters(self):
+        constant_(self.offset.weight.data, 0.)
+        constant_(self.offset.bias.data, 0.)
+        constant_(self.mask.weight.data, 0.)
+        constant_(self.mask.bias.data, 0.)
+        xavier_uniform_(self.input_proj.weight.data)
+        constant_(self.input_proj.bias.data, 0.)
+        xavier_uniform_(self.output_proj.weight.data)
+        constant_(self.output_proj.bias.data, 0.)
+
+    def forward(self, input):
+        """
+        :param query                       (N, H, W, C)
+        :return output                     (N, H, W, C)
+        """
+        N, H, W, _ = input.shape
+
+        x = self.input_proj(input)
+        x_proj = x
+
+        x1 = input.permute(0, 3, 1, 2)
+        x1 = self.dw_conv(x1)
+        offset = self.offset(x1)
+        mask = self.mask(x1).reshape(N, H, W, self.group, -1)
+        mask = F.softmax(mask, -1).reshape(N, H, W, -1)
+
+        x = dcnv3_core_pytorch(
+            x, offset, mask,
+            self.kernel_size, self.kernel_size,
+            self.stride, self.stride,
+            self.pad, self.pad,
+            self.dilation, self.dilation,
+            self.group, self.group_channels,
+            self.offset_scale, self.remove_center)
+        if self.center_feature_scale:
+            center_feature_scale = self.center_feature_scale_module(
+                x1, self.center_feature_scale_proj_weight, self.center_feature_scale_proj_bias)
+            # N, H, W, groups -> N, H, W, groups, 1 -> N, H, W, groups, _d_per_group -> N, H, W, channels
+            center_feature_scale = center_feature_scale[..., None].repeat(
+                1, 1, 1, 1, self.channels // self.group).flatten(-2)
+            x = x * (1 - center_feature_scale) + x_proj * center_feature_scale
+        x = self.output_proj(x)
+
+        return x
+
+
+class DCNv3(nn.Module):
+    def __init__(
+        self,
+        channels=64,
+        kernel_size=3,
+        dw_kernel_size=None,
+        stride=1,
+        pad=1,
+        dilation=1,
+        group=4,
+        offset_scale=1.0,
+        act_layer='GELU',
+        norm_layer='LN',
+        center_feature_scale=False,
+        remove_center=False,
+    ):
+        """
+        DCNv3 Module
+        :param channels
+        :param kernel_size
+        :param stride
+        :param pad
+        :param dilation
+        :param group
+        :param offset_scale
+        :param act_layer
+        :param norm_layer
+        """
+        super().__init__()
+        if channels % group != 0:
+            raise ValueError(
+                f'channels must be divisible by group, but got {channels} and {group}')
+        _d_per_group = channels // group
+        dw_kernel_size = dw_kernel_size if dw_kernel_size is not None else kernel_size
+        # you'd better set _d_per_group to a power of 2 which is more efficient in our CUDA implementation
+        if not _is_power_of_2(_d_per_group):
+            warnings.warn(
+                "You'd better set channels in DCNv3 to make the dimension of each attention head a power of 2 "
+                "which is more efficient in our CUDA implementation.")
+
+        self.offset_scale = offset_scale
+        self.channels = channels
+        self.kernel_size = kernel_size
+        self.dw_kernel_size = dw_kernel_size
+        self.stride = stride
+        self.dilation = dilation
+        self.pad = pad
+        self.group = group
+        self.group_channels = channels // group
+        self.offset_scale = offset_scale
+        self.center_feature_scale = center_feature_scale
+        self.remove_center = int(remove_center)
+
+        if self.remove_center and self.kernel_size % 2 == 0:
+            raise ValueError('remove_center is only compatible with odd kernel size.')
+
+        self.dw_conv = nn.Sequential(
+            nn.Conv2d(
+                channels,
+                channels,
+                kernel_size=dw_kernel_size,
+                stride=1,
+                padding=(dw_kernel_size - 1) // 2,
+                groups=channels),
+            build_norm_layer(
+                channels,
+                norm_layer,
+                'channels_first',
+                'channels_last'),
+            build_act_layer(act_layer))
+        self.offset = nn.Linear(
+            channels,
+            group * (kernel_size * kernel_size - remove_center) * 2)
+        self.mask = nn.Linear(
+            channels,
+            group * (kernel_size * kernel_size - remove_center))
+        self.input_proj = nn.Linear(channels, channels)
+        self.output_proj = nn.Linear(channels, channels)
+        self._reset_parameters()
+        
+        if center_feature_scale:
+            self.center_feature_scale_proj_weight = nn.Parameter(
+                torch.zeros((group, channels), dtype=torch.float))
+            self.center_feature_scale_proj_bias = nn.Parameter(
+                torch.tensor(0.0, dtype=torch.float).view((1,)).repeat(group, ))
+            self.center_feature_scale_module = CenterFeatureScaleModule()
+
+    def _reset_parameters(self):
+        constant_(self.offset.weight.data, 0.)
+        constant_(self.offset.bias.data, 0.)
+        constant_(self.mask.weight.data, 0.)
+        constant_(self.mask.bias.data, 0.)
+        xavier_uniform_(self.input_proj.weight.data)
+        constant_(self.input_proj.bias.data, 0.)
+        xavier_uniform_(self.output_proj.weight.data)
+        constant_(self.output_proj.bias.data, 0.)
+
+    def forward(self, input):
+        """
+        :param query                       (N, H, W, C)
+        :return output                     (N, H, W, C)
+        """
+        N, H, W, _ = input.shape
+
+        x = self.input_proj(input)
+        x_proj = x
+        dtype = x.dtype
+
+        x1 = input.permute(0, 3, 1, 2)
+        x1 = self.dw_conv(x1)
+        offset = self.offset(x1)
+        mask = self.mask(x1).reshape(N, H, W, self.group, -1)
+        mask = F.softmax(mask, -1)
+        mask = mask.reshape(N, H, W, -1).type(dtype)
+        
+        x = DCNv3Function.apply(
+            x, offset, mask,
+            self.kernel_size, self.kernel_size,
+            self.stride, self.stride,
+            self.pad, self.pad,
+            self.dilation, self.dilation,
+            self.group, self.group_channels,
+            self.offset_scale,
+            256,
+            self.remove_center)
+        
+        if self.center_feature_scale:
+            center_feature_scale = self.center_feature_scale_module(
+                x1, self.center_feature_scale_proj_weight, self.center_feature_scale_proj_bias)
+            # N, H, W, groups -> N, H, W, groups, 1 -> N, H, W, groups, _d_per_group -> N, H, W, channels
+            center_feature_scale = center_feature_scale[..., None].repeat(
+                1, 1, 1, 1, self.channels // self.group).flatten(-2)
+            x = x * (1 - center_feature_scale) + x_proj * center_feature_scale
+        x = self.output_proj(x)
+
+        return x

classification/ops_dcnv3/setup.py

@@ -0,0 +1,75 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+import os
+import glob
+
+import torch
+
+from torch.utils.cpp_extension import CUDA_HOME
+from torch.utils.cpp_extension import CppExtension
+from torch.utils.cpp_extension import CUDAExtension
+
+from setuptools import find_packages
+from setuptools import setup
+
+requirements = ["torch", "torchvision"]
+
+
+def get_extensions():
+    this_dir = os.path.dirname(os.path.abspath(__file__))
+    extensions_dir = os.path.join(this_dir, "src")
+
+    main_file = glob.glob(os.path.join(extensions_dir, "*.cpp"))
+    source_cpu = glob.glob(os.path.join(extensions_dir, "cpu", "*.cpp"))
+    source_cuda = glob.glob(os.path.join(extensions_dir, "cuda", "*.cu"))
+
+    sources = main_file + source_cpu
+    extension = CppExtension
+    extra_compile_args = {"cxx": []}
+    define_macros = []
+
+    if torch.cuda.is_available() and CUDA_HOME is not None:
+        extension = CUDAExtension
+        sources += source_cuda
+        define_macros += [("WITH_CUDA", None)]
+        extra_compile_args["nvcc"] = [
+            # "-DCUDA_HAS_FP16=1",
+            # "-D__CUDA_NO_HALF_OPERATORS__",
+            # "-D__CUDA_NO_HALF_CONVERSIONS__",
+            # "-D__CUDA_NO_HALF2_OPERATORS__",
+        ]
+    else:
+        raise NotImplementedError('Cuda is not availabel')
+
+    sources = [os.path.join(extensions_dir, s) for s in sources]
+    include_dirs = [extensions_dir]
+    ext_modules = [
+        extension(
+            "DCNv3",
+            sources,
+            include_dirs=include_dirs,
+            define_macros=define_macros,
+            extra_compile_args=extra_compile_args,
+        )
+    ]
+    return ext_modules
+
+
+setup(
+    name="DCNv3",
+    version="1.1",
+    author="InternImage",
+    url="https://github.com/OpenGVLab/InternImage",
+    description=
+    "PyTorch Wrapper for CUDA Functions of DCNv3",
+    packages=find_packages(exclude=(
+        "configs",
+        "tests",
+    )),
+    ext_modules=get_extensions(),
+    cmdclass={"build_ext": torch.utils.cpp_extension.BuildExtension},
+)

classification/ops_dcnv3/src/cpu/dcnv3_cpu.cpp

@@ -0,0 +1,37 @@
+/*!
+**************************************************************************************************
+* InternImage
+* Copyright (c) 2022 OpenGVLab
+* Licensed under The MIT License [see LICENSE for details]
+**************************************************************************************************
+* Modified from
+*https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+#include <vector>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+at::Tensor dcnv3_cpu_forward(const at::Tensor &input, const at::Tensor &offset,
+                             const at::Tensor &mask, const int kernel_h,
+                             const int kernel_w, const int stride_h,
+                             const int stride_w, const int pad_h,
+                             const int pad_w, const int dilation_h,
+                             const int dilation_w, const int group,
+                             const int group_channels, const float offset_scale,
+                             const int im2col_step) {
+    AT_ERROR("Not implement on cpu");
+}
+
+std::vector<at::Tensor>
+dcnv3_cpu_backward(const at::Tensor &input, const at::Tensor &offset,
+                   const at::Tensor &mask, const int kernel_h,
+                   const int kernel_w, const int stride_h, const int stride_w,
+                   const int pad_h, const int pad_w, const int dilation_h,
+                   const int dilation_w, const int group,
+                   const int group_channels, const float offset_scale,
+                   const at::Tensor &grad_output, const int im2col_step) {
+    AT_ERROR("Not implement on cpu");
+}

classification/ops_dcnv3/src/cpu/dcnv3_cpu.h

@@ -0,0 +1,31 @@
+/*!
+**************************************************************************************************
+* InternImage
+* Copyright (c) 2022 OpenGVLab
+* Licensed under The MIT License [see LICENSE for details]
+**************************************************************************************************
+* Modified from
+*https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+#pragma once
+#include <torch/extension.h>
+
+at::Tensor dcnv3_cpu_forward(const at::Tensor &input, const at::Tensor &offset,
+                             const at::Tensor &mask, const int kernel_h,
+                             const int kernel_w, const int stride_h,
+                             const int stride_w, const int pad_h,
+                             const int pad_w, const int dilation_h,
+                             const int dilation_w, const int group,
+                             const int group_channels, const float offset_scale,
+                             const int im2col_step);
+
+std::vector<at::Tensor>
+dcnv3_cpu_backward(const at::Tensor &input, const at::Tensor &offset,
+                   const at::Tensor &mask, const int kernel_h,
+                   const int kernel_w, const int stride_h, const int stride_w,
+                   const int pad_h, const int pad_w, const int dilation_h,
+                   const int dilation_w, const int group,
+                   const int group_channels, const float offset_scale,
+                   const at::Tensor &grad_output, const int im2col_step);

classification/ops_dcnv3/src/cuda/dcnv3_cuda.cu

@@ -0,0 +1,174 @@
+/*!
+**************************************************************************************************
+* InternImage
+* Copyright (c) 2022 OpenGVLab
+* Licensed under The MIT License [see LICENSE for details]
+**************************************************************************************************
+* Modified from
+*https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+#include "cuda/dcnv3_im2col_cuda.cuh"
+#include <vector>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <torch/torch.h>
+
+at::Tensor dcnv3_cuda_forward(const at::Tensor &input, const at::Tensor &offset,
+                              const at::Tensor &mask, const int kernel_h,
+                              const int kernel_w, const int stride_h,
+                              const int stride_w, const int pad_h,
+                              const int pad_w, const int dilation_h,
+                              const int dilation_w, const int group,
+                              const int group_channels,
+                              const float offset_scale, const int im2col_step, const int remove_center) {
+    AT_ASSERTM(input.is_contiguous(), "input tensor has to be contiguous");
+    AT_ASSERTM(offset.is_contiguous(), "offset tensor has to be contiguous");
+    AT_ASSERTM(mask.is_contiguous(), "mask tensor has to be contiguous");
+    AT_ASSERTM(input.type().is_cuda(), "input must be a CUDA tensor");
+    AT_ASSERTM(offset.type().is_cuda(), "offset must be a CUDA tensor");
+    AT_ASSERTM(mask.type().is_cuda(), "mask must be a CUDA tensor");
+
+    const int batch = input.size(0);
+    const int height_in = input.size(1);
+    const int width_in = input.size(2);
+    const int channels = input.size(3);
+    const int height_out =
+        (height_in + 2 * pad_h - (dilation_h * (kernel_h - 1) + 1)) / stride_h +
+        1;
+    const int width_out =
+        (width_in + 2 * pad_w - (dilation_w * (kernel_w - 1) + 1)) / stride_w +
+        1;
+    const int im2col_step_ = std::min(batch, im2col_step);
+
+    AT_ASSERTM(batch % im2col_step_ == 0,
+               "batch(%d) must divide im2col_step(%d)", batch, im2col_step_);
+    AT_ASSERTM(
+        channels == (group * group_channels),
+        "Input channels and group times group channels wont match: (%d vs %d).",
+        channels, group * group_channels);
+
+    auto output =
+        at::zeros({batch, height_out, width_out, group * group_channels},
+                  input.options());
+
+    const int batch_n = im2col_step_;
+    auto output_n = output.view({batch / batch_n, batch_n, height_out,
+                                 width_out, group * group_channels});
+    auto per_input_size = height_in * width_in * group * group_channels;
+    auto per_offset_size =
+        height_out * width_out * group * (kernel_h * kernel_w - remove_center) * 2;
+    auto per_mask_size = height_out * width_out * group * (kernel_h * kernel_w - remove_center);
+    for (int n = 0; n < batch / im2col_step_; ++n) {
+        auto columns = output_n.select(0, n);
+        // AT_DISPATCH_FLOATING_TYPES(
+        AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+            input.type(), "ms_deform_attn_forward_cuda", ([&] {
+                dcnv3_im2col_cuda(
+                    at::cuda::getCurrentCUDAStream(),
+                    input.data<scalar_t>() + n * im2col_step_ * per_input_size,
+                    offset.data<scalar_t>() +
+                        n * im2col_step_ * per_offset_size,
+                    mask.data<scalar_t>() + n * im2col_step_ * per_mask_size,
+                    columns.data<scalar_t>(), kernel_h, kernel_w, stride_h,
+                    stride_w, pad_h, pad_w, dilation_h, dilation_w, group,
+                    group_channels, batch_n, height_in, width_in, height_out,
+                    width_out, offset_scale, remove_center);
+            }));
+    }
+
+    return output;
+}
+
+std::vector<at::Tensor>
+dcnv3_cuda_backward(const at::Tensor &input, const at::Tensor &offset,
+                    const at::Tensor &mask, const int kernel_h,
+                    const int kernel_w, const int stride_h, const int stride_w,
+                    const int pad_h, const int pad_w, const int dilation_h,
+                    const int dilation_w, const int group,
+                    const int group_channels, const float offset_scale,
+                    const at::Tensor &grad_output, const int im2col_step, const int remove_center) {
+
+    AT_ASSERTM(input.is_contiguous(), "input tensor has to be contiguous");
+    AT_ASSERTM(offset.is_contiguous(), "offset tensor has to be contiguous");
+    AT_ASSERTM(mask.is_contiguous(), "mask tensor has to be contiguous");
+    AT_ASSERTM(grad_output.is_contiguous(),
+               "grad_output tensor has to be contiguous");
+    AT_ASSERTM(input.type().is_cuda(), "input must be a CUDA tensor");
+    AT_ASSERTM(offset.type().is_cuda(), "offset must be a CUDA tensor");
+    AT_ASSERTM(mask.type().is_cuda(), "mask must be a CUDA tensor");
+    AT_ASSERTM(grad_output.type().is_cuda(),
+               "grad_output must be a CUDA tensor");
+
+    const int batch = input.size(0);
+    const int height_in = input.size(1);
+    const int width_in = input.size(2);
+    const int channels = input.size(3);
+    const int height_out =
+        (height_in + 2 * pad_h - (dilation_h * (kernel_h - 1) + 1)) / stride_h +
+        1;
+    const int width_out =
+        (width_in + 2 * pad_w - (dilation_w * (kernel_w - 1) + 1)) / stride_w +
+        1;
+    const int im2col_step_ = std::min(batch, im2col_step);
+
+    AT_ASSERTM(batch % im2col_step_ == 0,
+               "batch(%d) must divide im2col_step(%d)", batch, im2col_step_);
+    AT_ASSERTM(
+        channels == (group * group_channels),
+        "Input channels and group times group channels wont match: (%d vs %d).",
+        channels, group * group_channels);
+
+    auto dtype = input.dtype();
+    if (dtype == at::kHalf) {
+        dtype = at::kFloat;
+    }
+
+    auto grad_input = at::zeros_like(input, dtype);
+    auto grad_offset = at::zeros_like(offset, dtype);
+    auto grad_mask = at::zeros_like(mask, dtype);
+
+    const int batch_n = im2col_step_;
+    auto per_input_size = height_in * width_in * group * group_channels;
+    auto per_offset_size =
+        height_out * width_out * group * (kernel_h * kernel_w - remove_center) * 2;
+    auto per_mask_size = height_out * width_out * group * (kernel_h * kernel_w - remove_center);
+    auto grad_output_n =
+        grad_output.view({batch / im2col_step_, batch_n, height_out * width_out,
+                          group, group_channels});
+
+    for (int n = 0; n < batch / im2col_step_; ++n) {
+        auto grad_output_g = grad_output_n.select(0, n);
+        // AT_DISPATCH_FLOATING_TYPES(
+        AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+            input.type(), "ms_deform_attn_backward_cuda", ([&] {
+                dcnv3_col2im_cuda(
+                    at::cuda::getCurrentCUDAStream(),
+                    grad_output_g.data<scalar_t>(),
+                    input.data<scalar_t>() + n * im2col_step_ * per_input_size,
+                    offset.data<scalar_t>() +
+                        n * im2col_step_ * per_offset_size,
+                    mask.data<scalar_t>() + n * im2col_step_ * per_mask_size,
+                    kernel_h, kernel_w, stride_h, stride_w, pad_h, pad_w,
+                    dilation_h, dilation_w, group, group_channels, batch_n,
+                    height_in, width_in, height_out, width_out, offset_scale, remove_center,
+                    grad_input.data<opmath_t>() +
+                        n * im2col_step_ * per_input_size,
+                    grad_offset.data<opmath_t>() +
+                        n * im2col_step_ * per_offset_size,
+                    grad_mask.data<opmath_t>() +
+                        n * im2col_step_ * per_mask_size);
+            }));
+    }
+
+    if (input.dtype() == torch::kHalf) {
+        return {grad_input.to(torch::kHalf), grad_offset.to(torch::kHalf),
+                grad_mask.to(torch::kHalf)};
+    } else {
+        return {grad_input, grad_offset, grad_mask};
+    }
+}

classification/ops_dcnv3/src/cuda/dcnv3_cuda.h

@@ -0,0 +1,31 @@
+/*!
+**************************************************************************************************
+* InternImage
+* Copyright (c) 2022 OpenGVLab
+* Licensed under The MIT License [see LICENSE for details]
+**************************************************************************************************
+* Modified from
+*https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+#pragma once
+#include <torch/extension.h>
+
+at::Tensor dcnv3_cuda_forward(const at::Tensor &input, const at::Tensor &offset,
+                              const at::Tensor &mask, const int kernel_h,
+                              const int kernel_w, const int stride_h,
+                              const int stride_w, const int pad_h,
+                              const int pad_w, const int dilation_h,
+                              const int dilation_w, const int group,
+                              const int group_channels,
+                              const float offset_scale, const int im2col_step, const int remove_center);
+
+std::vector<at::Tensor>
+dcnv3_cuda_backward(const at::Tensor &input, const at::Tensor &offset,
+                    const at::Tensor &mask, const int kernel_h,
+                    const int kernel_w, const int stride_h, const int stride_w,
+                    const int pad_h, const int pad_w, const int dilation_h,
+                    const int dilation_w, const int group,
+                    const int group_channels, const float offset_scale,
+                    const at::Tensor &grad_output, const int im2col_step, const int remove_center);

classification/ops_dcnv3/src/dcnv3.h

@@ -0,0 +1,59 @@
+/*!
+**************************************************************************************************
+* InternImage
+* Copyright (c) 2022 OpenGVLab
+* Licensed under The MIT License [see LICENSE for details]
+**************************************************************************************************
+* Modified from
+*https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+#pragma once
+
+#include "cpu/dcnv3_cpu.h"
+
+#ifdef WITH_CUDA
+#include "cuda/dcnv3_cuda.h"
+#endif
+
+at::Tensor dcnv3_forward(const at::Tensor &input, const at::Tensor &offset,
+                         const at::Tensor &mask, const int kernel_h,
+                         const int kernel_w, const int stride_h,
+                         const int stride_w, const int pad_h, const int pad_w,
+                         const int dilation_h, const int dilation_w,
+                         const int group, const int group_channels,
+                         const float offset_scale, const int im2col_step, const int remove_center) {
+    if (input.type().is_cuda()) {
+#ifdef WITH_CUDA
+        return dcnv3_cuda_forward(input, offset, mask, kernel_h, kernel_w,
+                                  stride_h, stride_w, pad_h, pad_w, dilation_h,
+                                  dilation_w, group, group_channels,
+                                  offset_scale, im2col_step, remove_center);
+#else
+        AT_ERROR("Not compiled with GPU support");
+#endif
+    }
+    AT_ERROR("Not implemented on the CPU");
+}
+
+std::vector<at::Tensor>
+dcnv3_backward(const at::Tensor &input, const at::Tensor &offset,
+               const at::Tensor &mask, const int kernel_h, const int kernel_w,
+               const int stride_h, const int stride_w, const int pad_h,
+               const int pad_w, const int dilation_h, const int dilation_w,
+               const int group, const int group_channels,
+               const float offset_scale, const at::Tensor &grad_output,
+               const int im2col_step, const int remove_center) {
+    if (input.type().is_cuda()) {
+#ifdef WITH_CUDA
+        return dcnv3_cuda_backward(input, offset, mask, kernel_h, kernel_w,
+                                   stride_h, stride_w, pad_h, pad_w, dilation_h,
+                                   dilation_w, group, group_channels,
+                                   offset_scale, grad_output, im2col_step, remove_center);
+#else
+        AT_ERROR("Not compiled with GPU support");
+#endif
+    }
+    AT_ERROR("Not implemented on the CPU");
+}

classification/ops_dcnv3/src/vision.cpp

@@ -0,0 +1,17 @@
+/*!
+**************************************************************************************************
+* InternImage
+* Copyright (c) 2022 OpenGVLab
+* Licensed under The MIT License [see LICENSE for details]
+**************************************************************************************************
+* Modified from
+*https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+#include "dcnv3.h"
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("dcnv3_forward", &dcnv3_forward, "dcnv3_forward");
+    m.def("dcnv3_backward", &dcnv3_backward, "dcnv3_backward");
+}

classification/ops_dcnv3/test.py

@@ -0,0 +1,264 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import time
+import torch
+import torch.nn as nn
+import math
+from torch.autograd import gradcheck
+
+from functions.dcnv3_func import DCNv3Function, dcnv3_core_pytorch
+
+H_in, W_in = 8, 8
+N, M, D = 2, 4, 16
+Kh, Kw = 3, 3
+remove_center = False
+P = Kh * Kw - remove_center
+offset_scale = 2.0
+pad = 1
+dilation = 1
+stride = 1
+H_out = (H_in + 2 * pad - (dilation * (Kh - 1) + 1)) // stride + 1
+W_out = (W_in + 2 * pad - (dilation * (Kw - 1) + 1)) // stride + 1
+
+torch.manual_seed(3)
+
+
+@torch.no_grad()
+def check_forward_equal_with_pytorch_double():
+    input = torch.rand(N, H_in, W_in, M*D).cuda() * 0.01
+    offset = torch.rand(N, H_out, W_out, M*P*2).cuda() * 10
+    mask = torch.rand(N, H_out, W_out, M, P).cuda() + 1e-5
+    mask /= mask.sum(-1, keepdim=True)
+    mask = mask.reshape(N, H_out, W_out, M*P)
+
+    output_pytorch = dcnv3_core_pytorch(
+        input.double(),
+        offset.double(),
+        mask.double(),
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale, remove_center).detach().cpu()
+
+    im2col_step = 2
+    output_cuda = DCNv3Function.apply(
+        input.double(),
+        offset.double(),
+        mask.double(),
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center).detach().cpu()
+
+    fwdok = torch.allclose(output_cuda, output_pytorch)
+    max_abs_err = (output_cuda - output_pytorch).abs().max()
+    max_rel_err = ((output_cuda - output_pytorch).abs() /
+                   output_pytorch.abs()).max()
+    print('>>> forward double')
+    print(f'* {fwdok} check_forward_equal_with_pytorch_double: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}')
+
+
+@torch.no_grad()
+def check_forward_equal_with_pytorch_float():
+    input = torch.rand(N, H_in, W_in, M*D).cuda() * 0.01
+    offset = torch.rand(N, H_out, W_out, M*P*2).cuda() * 10
+    mask = torch.rand(N, H_out, W_out, M, P).cuda() + 1e-5
+    mask /= mask.sum(-1, keepdim=True)
+    mask = mask.reshape(N, H_out, W_out, M*P)
+
+    output_pytorch = dcnv3_core_pytorch(
+        input,
+        offset,
+        mask,
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale, remove_center).detach().cpu()
+
+    im2col_step = 2
+    output_cuda = DCNv3Function.apply(
+        input,
+        offset,
+        mask,
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center).detach().cpu()
+
+    fwdok = torch.allclose(output_cuda, output_pytorch, rtol=1e-2, atol=1e-3)
+    max_abs_err = (output_cuda - output_pytorch).abs().max()
+    max_rel_err = ((output_cuda - output_pytorch).abs() /
+                   output_pytorch.abs()).max()
+    print('>>> forward float')
+    print(f'* {fwdok} check_forward_equal_with_pytorch_float: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}')
+
+
+def check_backward_equal_with_pytorch_double(channels=4, grad_input=True, grad_offset=True, grad_mask=True):
+    # H_in, W_in = 4, 4
+    N = 2
+    M = 2
+    H_out = (H_in + 2 * pad - (dilation * (Kh - 1) + 1)) // stride + 1
+    W_out = (W_in + 2 * pad - (dilation * (Kw - 1) + 1)) // stride + 1
+
+    D = channels
+    input0 = torch.rand(N, H_in, W_in, M*D).cuda() * 0.01
+    offset0 = torch.rand(N, H_out, W_out, M*P*2).cuda() * 10
+    mask0 = torch.rand(N, H_out, W_out, M, P).cuda() + 1e-5
+    mask0 /= mask0.sum(-1, keepdim=True)
+    mask0 = mask0.reshape(N, H_out, W_out, M*P)
+    input0.requires_grad = grad_input
+    offset0.requires_grad = grad_offset
+    mask0.requires_grad = grad_mask
+
+    output_pytorch = dcnv3_core_pytorch(
+        input0.double(),
+        offset0.double(),
+        mask0.double(),
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale, remove_center)
+    output_pytorch.sum().backward()
+
+    input1 = input0.detach()
+    offset1 = offset0.detach()
+    mask1 = mask0.detach()
+    input1.requires_grad = grad_input
+    offset1.requires_grad = grad_offset
+    mask1.requires_grad = grad_mask
+
+    im2col_step = 2
+    output_cuda = DCNv3Function.apply(
+        input1.double(),
+        offset1.double(),
+        mask1.double(),
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center)
+    output_cuda.sum().backward()
+
+    print(f'>>> backward double: channels {D}')
+    bwdok = torch.allclose(input0.grad, input1.grad, rtol=1e-2, atol=1e-3)
+    max_abs_err = (input0.grad - input1.grad).abs().max()
+    max_rel_err = ((input0.grad - input1.grad).abs() /
+                   input0.grad.abs()).max()
+    print(
+        f'* {bwdok} input_grad check_backward_equal_with_pytorch_double: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}')
+
+    bwdok = torch.allclose(offset0.grad, offset1.grad, rtol=1e-2, atol=1e-3)
+    max_abs_err = (offset0.grad - offset1.grad).abs().max()
+    max_rel_err = ((offset0.grad - offset1.grad).abs() /
+                   offset0.grad.abs()).max()
+    print(
+        f'* {bwdok} offset_grad check_backward_equal_with_pytorch_double: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}')
+
+    bwdok = torch.allclose(mask0.grad, mask1.grad, rtol=1e-2, atol=1e-3)
+    max_abs_err = (mask0.grad - mask1.grad).abs().max()
+    max_rel_err = ((mask0.grad - mask1.grad).abs() /
+                   mask0.grad.abs()).max()
+    print(
+        f'* {bwdok} mask_grad check_backward_equal_with_pytorch_double: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}')
+
+
+def check_backward_equal_with_pytorch_float(channels=4, grad_input=True, grad_offset=True, grad_mask=True):
+    # H_in, W_in = 4, 4
+    N = 2
+    M = 2
+    H_out = (H_in + 2 * pad - (dilation * (Kh - 1) + 1)) // stride + 1
+    W_out = (W_in + 2 * pad - (dilation * (Kw - 1) + 1)) // stride + 1
+
+    D = channels
+    input0 = torch.rand(N, H_in, W_in, M*D).cuda() * 0.01
+    offset0 = torch.rand(N, H_out, W_out, M*P*2).cuda() * 10
+    mask0 = torch.rand(N, H_out, W_out, M, P).cuda() + 1e-5
+    mask0 /= mask0.sum(-1, keepdim=True)
+    mask0 = mask0.reshape(N, H_out, W_out, M*P)
+    input0.requires_grad = grad_input
+    offset0.requires_grad = grad_offset
+    mask0.requires_grad = grad_mask
+
+    output_pytorch = dcnv3_core_pytorch(
+        input0,
+        offset0,
+        mask0,
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale, remove_center)
+    output_pytorch.sum().backward()
+
+    input1 = input0.detach()
+    offset1 = offset0.detach()
+    mask1 = mask0.detach()
+    input1.requires_grad = grad_input
+    offset1.requires_grad = grad_offset
+    mask1.requires_grad = grad_mask
+
+    im2col_step = 2
+    output_cuda = DCNv3Function.apply(
+        input1,
+        offset1,
+        mask1,
+        Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, offset_scale,
+        im2col_step, remove_center)
+    output_cuda.sum().backward()
+
+    print(f'>>> backward float: channels {D}')
+    bwdok = torch.allclose(input0.grad, input1.grad, rtol=1e-2, atol=1e-3)
+    max_abs_err = (input0.grad - input1.grad).abs().max()
+    max_rel_err = ((input0.grad - input1.grad).abs() /
+                   input0.grad.abs()).max()
+    print(
+        f'* {bwdok} input_grad check_backward_equal_with_pytorch_float: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}')
+
+    bwdok = torch.allclose(offset0.grad, offset1.grad, rtol=1e-2, atol=1e-3)
+    max_abs_err = (offset0.grad - offset1.grad).abs().max()
+    max_rel_err = ((offset0.grad - offset1.grad).abs() /
+                   offset0.grad.abs()).max()
+    print(
+        f'* {bwdok} offset_grad check_backward_equal_with_pytorch_float: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}')
+
+    bwdok = torch.allclose(mask0.grad, mask1.grad, rtol=1e-2, atol=1e-3)
+    max_abs_err = (mask0.grad - mask1.grad).abs().max()
+    max_rel_err = ((mask0.grad - mask1.grad).abs() /
+                   mask0.grad.abs()).max()
+    print(
+        f'* {bwdok} mask_grad check_backward_equal_with_pytorch_float: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}')
+
+
+@torch.no_grad()
+def check_time_cost(im2col_step=128):
+    N = 512
+    H_in, W_in = 64, 64
+    H_out = (H_in + 2 * pad - (dilation * (Kh - 1) + 1)) // stride + 1
+    W_out = (W_in + 2 * pad - (dilation * (Kw - 1) + 1)) // stride + 1
+
+    input = torch.rand(N, H_in, W_in, M*D).cuda() * 0.01
+    offset = torch.rand(N, H_out, W_out, M*P*2).cuda() * 10
+    mask = torch.rand(N, H_out, W_out, M, P).cuda() + 1e-5
+    mask /= mask.sum(-1, keepdim=True)
+    mask = mask.reshape(N, H_out, W_out, M*P)
+    print(
+        f'>>> time cost: im2col_step {im2col_step}; input {input.shape}; points {P} ')
+    repeat = 100
+    for i in range(repeat):
+        output_cuda = DCNv3Function.apply(
+            input,
+            offset,
+            mask,
+            Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, 1.0,
+            im2col_step, remove_center)
+    torch.cuda.synchronize()
+    start = time.time()
+    for i in range(repeat):
+        output_cuda = DCNv3Function.apply(
+            input,
+            offset,
+            mask,
+            Kh, Kw, stride, stride, Kh // 2, Kw // 2, dilation, dilation, M, D, 1.0,
+            im2col_step, remove_center)
+    torch.cuda.synchronize()
+    print(f'foward time cost: {(time.time() - start) / repeat}')
+
+
+if __name__ == '__main__':
+    check_forward_equal_with_pytorch_double()
+    check_forward_equal_with_pytorch_float()
+    for channels in [1, 16, 30, 32, 64, 71, 1025]:
+        check_backward_equal_with_pytorch_double(channels, True, True, True)
+    for channels in [1, 16, 30, 32, 64, 71, 1025]:
+        check_backward_equal_with_pytorch_float(channels, True, True, True)
+    for i in range(3):
+        im2col_step = 128 * (2 ** i)
+        check_time_cost(im2col_step)

classification/optimizer.py

@@ -0,0 +1,159 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+
+from torch import optim as optim
+from torch.distributed.optim import ZeroRedundancyOptimizer
+
+
+def build_optimizer(config, model):
+    """
+    Build optimizer, set weight decay of normalization to 0 by default.
+    """
+    skip = {}
+    skip_keywords = {}
+    if hasattr(model, 'no_weight_decay'):
+        skip = model.no_weight_decay()
+    if hasattr(model, 'no_weight_decay_keywords'):
+        skip_keywords = model.no_weight_decay_keywords()
+
+    parameters = set_weight_decay_and_lr(
+        model,
+        config.TRAIN.WEIGHT_DECAY,
+        config.TRAIN.BASE_LR,
+        skip,
+        skip_keywords,
+        lr_layer_decay=config.TRAIN.LR_LAYER_DECAY,
+        lr_layer_decay_ratio=config.TRAIN.LR_LAYER_DECAY_RATIO,
+        freeze_backbone=config.TRAIN.OPTIMIZER.FREEZE_BACKBONE,
+        dcn_lr_mul=config.TRAIN.OPTIMIZER.DCN_LR_MUL,
+    )
+
+    opt_lower = config.TRAIN.OPTIMIZER.NAME.lower()
+    optimizer = None
+    use_zero = config.TRAIN.OPTIMIZER.USE_ZERO
+    if use_zero:
+        print(f"\nUse Zero!")
+        if opt_lower == 'sgd':
+            # an ugly implementation
+            # this problem is fixed after torch 1.12
+            # https://github.com/pytorch/pytorch/issues/71347
+
+            # before 1.12, we could only pass list to zero optimizer, so we first pass parameters[0] with its lr and weight decay,
+            # then we add other parameter via parameter group.
+
+            optimizer = ZeroRedundancyOptimizer(
+                parameters[0]['params'],
+                optimizer_class=optim.SGD,
+                momentum=config.TRAIN.OPTIMIZER.MOMENTUM, nesterov=True,
+                lr=parameters[0]['lr'], weight_decay=parameters[0]['weight_decay']
+            )
+            if len(parameters) > 1:
+                for param_group in parameters[1:]:
+                    optimizer.add_param_group(param_group)
+        elif opt_lower == 'adamw':
+            optimizer = ZeroRedundancyOptimizer(
+                parameters[0]['params'],
+                optimizer_class=optim.AdamW,
+                eps=config.TRAIN.OPTIMIZER.EPS, betas=config.TRAIN.OPTIMIZER.BETAS,
+                lr=parameters[0]['lr'], weight_decay=parameters[0]['weight_decay']
+            )
+            if len(parameters) > 1:
+                for param_group in parameters[1:]:
+                    optimizer.add_param_group(param_group)
+    else:
+        if opt_lower == 'sgd':
+            optimizer = optim.SGD(parameters,
+                                  momentum=config.TRAIN.OPTIMIZER.MOMENTUM,
+                                  nesterov=True,
+                                  lr=config.TRAIN.BASE_LR,
+                                  weight_decay=config.TRAIN.WEIGHT_DECAY)
+        elif opt_lower == 'adamw':
+            optimizer = optim.AdamW(parameters,
+                                    eps=config.TRAIN.OPTIMIZER.EPS,
+                                    betas=config.TRAIN.OPTIMIZER.BETAS,
+                                    lr=config.TRAIN.BASE_LR,
+                                    weight_decay=config.TRAIN.WEIGHT_DECAY)
+
+    return optimizer
+
+
+def check_keywords_in_name(name, keywords=()):
+    isin = False
+    for keyword in keywords:
+        if keyword in name:
+            isin = True
+    return isin
+
+
+def check_keywords_in_dict(name, keywords_dict):
+    for k, v in keywords_dict.items():
+        if k in name:
+            return v
+    return None
+
+
+def set_weight_decay_and_lr(
+    model,
+    weight_decay,
+    base_lr,
+    skip_list=(),
+    skip_keywords=(),
+    lr_layer_decay=None,
+    lr_layer_decay_ratio=None,
+    freeze_backbone=None,
+    dcn_lr_mul=None,
+    layerwise_lr=True,
+):
+    parameters = []
+    no_decay_name = []
+    lr_ratio_log = {}
+
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue  # frozen weights
+        if freeze_backbone:
+            for i in freeze_backbone:
+                if f'levels.{i}' in name:
+                    param.requires_grad = False
+        # 1. check wd
+        if len(param.shape) == 1 or name.endswith(".bias") or (
+                name in skip_list) or check_keywords_in_name(
+                    name, skip_keywords):
+            wd = 0.
+            no_decay_name.append(name)
+        else:
+            wd = weight_decay
+
+        if lr_layer_decay:
+            print('layer-wise lr decay is used !')
+            assert hasattr(model, 'lr_decay_keywards')
+            lr_ratio_keywards = model.lr_decay_keywards(lr_layer_decay_ratio)
+
+            # 2. check lr
+            ratio = check_keywords_in_dict(name, lr_ratio_keywards)
+            if ratio is not None:
+                lr = ratio * base_lr
+            else:
+                lr = base_lr
+
+            # dcn lr
+            if dcn_lr_mul is not None:
+                if 'offset' in name or 'attention_weights' in name or 'center_feature_scale_proj' in name or 'alpha_beta' in name:
+                    lr = dcn_lr_mul * lr
+
+            lr_ratio_log[name] = (base_lr, ratio, wd, param.requires_grad)
+        else:
+            lr = base_lr
+        parameters.append({'params': [param], 'weight_decay': wd, 'lr': lr, 'name': name})
+
+    print('no decay params: {no_decay_name}')
+    if layerwise_lr:
+        print('lr_ratio_params:')
+        for k, v in lr_ratio_log.items():
+            print(k, v)
+
+    return parameters

classification/train_in1k.sh

@@ -0,0 +1,31 @@
+#!/usr/bin/env bash
+
+set -x
+
+PARTITION=$1
+JOB_NAME=$2
+CONFIG=$3
+WORK_DIR=$4
+GPUS=${GPUS:-1}
+GPUS_PER_NODE=${GPUS_PER_NODE:-1}
+CPUS_PER_TASK=${CPUS_PER_TASK:-10}
+SRUN_ARGS=${SRUN_ARGS:-""}
+PY_ARGS=${@:5}
+
+PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \
+srun -p ${PARTITION} \
+    --job-name=${JOB_NAME} \
+    --gres=gpu:${GPUS_PER_NODE} \
+    --ntasks=${GPUS} \
+    --ntasks-per-node=${GPUS_PER_NODE} \
+    --cpus-per-task=${CPUS_PER_TASK} \
+    --kill-on-bad-exit=1 \
+    --quotatype=reserved \
+    ${SRUN_ARGS} \
+python -u main.py \
+    --cfg ${CONFIG} \
+    --accumulation-steps 1 \
+    --local-rank 0 \
+    --batch-size 128 \
+    --data-path /mnt/petrelfs/share/images \
+    --output work_dirs ${@:4} --launcher="slurm" 

classification/train_in1k_deepspeed.sh

@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+
+set -x
+
+PARTITION=$1
+JOB_NAME=$2
+CONFIG=$3
+GPUS=${GPUS:-8}
+GPUS_PER_NODE=${GPUS_PER_NODE:-8}
+CPUS_PER_TASK=${CPUS_PER_TASK:-12}
+SRUN_ARGS=${SRUN_ARGS:-""}
+
+PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \
+    srun -p ${PARTITION} \
+    --job-name=${JOB_NAME} \
+    --gres=gpu:${GPUS_PER_NODE} \
+    --ntasks=${GPUS} \
+    --ntasks-per-node=${GPUS_PER_NODE} \
+    --cpus-per-task=${CPUS_PER_TASK} \
+    --kill-on-bad-exit=1 \
+    --quotatype=spot \
+    ${SRUN_ARGS} \
+    python -u main_deepspeed.py \
+    --cfg ${CONFIG} \
+    --local-rank 0 \
+    --data-path /mnt/lustre/share/images \
+    --output work_dirs_deepspeed ${@:4}

classification/utils.py

@@ -0,0 +1,423 @@
+# --------------------------------------------------------
+# DCNv4
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+import os
+import math
+import torch
+import numpy as np
+import torch.distributed as dist
+from collections import OrderedDict
+from timm.utils import get_state_dict
+try:
+    # noinspection PyUnresolvedReferences
+    from apex import amp
+except ImportError:
+    amp = None
+
+
+def load_ema_checkpoint(config, model_ema, logger):
+    logger.info(
+        f'==============> Resuming form {config.MODEL.RESUME}....................'
+    )
+    if config.MODEL.RESUME.startswith('https'):
+        checkpoint = torch.hub.load_state_dict_from_url(config.MODEL.RESUME,
+                                                        map_location='cpu',
+                                                        check_hash=True)
+    else:
+        checkpoint = torch.load(config.MODEL.RESUME, map_location='cpu')
+
+    assert isinstance(checkpoint, dict)
+    if 'model_ema' in checkpoint:
+        new_state_dict = OrderedDict()
+        for k, v in checkpoint['model_ema'].items():
+            if model_ema.ema_has_module:
+                name = 'module.' + k if not k.startswith('module') else k
+            else:
+                name = k
+            new_state_dict[name] = v
+        msg = model_ema.ema.load_state_dict(new_state_dict, strict=False)
+        logger.info(msg)
+        logger.info('Loaded state_dict_ema')
+    else:
+        logger.warning(
+            'Failed to find state_dict_ema, starting from loaded model weights'
+        )
+
+    max_accuracy_ema = 0
+    if 'max_accuracy_ema' in checkpoint:
+        max_accuracy_ema = checkpoint['max_accuracy_ema']
+    if 'ema_decay' in checkpoint:
+        model_ema.decay = checkpoint['ema_decay']
+    return max_accuracy_ema
+
+
+def load_checkpoint(config, model, optimizer, lr_scheduler, scaler, logger):
+    logger.info(
+        f'==============> Resuming form {config.MODEL.RESUME}....................'
+    )
+    if config.MODEL.RESUME.startswith('https'):
+        checkpoint = torch.hub.load_state_dict_from_url(config.MODEL.RESUME,
+                                                        map_location='cpu',
+                                                        check_hash=True)
+    else:
+        checkpoint = torch.load(config.MODEL.RESUME, map_location='cpu')
+
+    print('resuming model')
+    msg = model.load_state_dict(checkpoint['model'], strict=False)
+    logger.info(msg)
+    max_accuracy = 0.0
+    if not config.EVAL_MODE and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:
+        if optimizer is not None:
+            print('resuming optimizer')
+            try:
+                optimizer.load_state_dict(checkpoint['optimizer'])
+            except:
+                print('resume optimizer failed')
+        if lr_scheduler is not None:
+            print('resuming lr_scheduler')
+            lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+        config.defrost()
+        config.TRAIN.START_EPOCH = checkpoint['epoch'] + 1
+        config.freeze()
+        if 'amp' in checkpoint and config.AMP_OPT_LEVEL != 'O0' and checkpoint[
+                'config'].AMP_OPT_LEVEL != 'O0':
+            scaler.load_state_dict(checkpoint['amp'])
+        logger.info(
+            f"=> loaded successfully {config.MODEL.RESUME} (epoch {checkpoint['epoch']})"
+        )
+        if 'max_accuracy' in checkpoint:
+            max_accuracy = checkpoint['max_accuracy']
+
+    del checkpoint
+    torch.cuda.empty_cache()
+
+    return max_accuracy
+
+
+def load_pretrained(config, model, logger):
+    logger.info(
+        f'==============> Loading weight {config.MODEL.PRETRAINED} for fine-tuning......'
+    )
+    checkpoint = torch.load(config.MODEL.PRETRAINED, map_location='cpu')
+
+    state_dict = checkpoint
+    if 'model' in checkpoint:
+        state_dict = checkpoint['model']
+    elif 'module' in checkpoint:
+        state_dict = checkpoint['module']
+
+    first_key = list(state_dict.keys())[0]
+    # delete teacher weights
+    if 'student' in first_key or 'teacher' in first_key:
+        new_state_dict = OrderedDict()
+        for k, v in state_dict.items():
+            if 'student_proj' in k:
+                continue
+            if 'student' in k:
+                new_k = k.replace('student.', '')
+                new_state_dict[new_k] = v
+        state_dict = new_state_dict
+
+    # weights from sim
+    if 'mask_token' in first_key:
+        new_state_dict = OrderedDict()
+        for k, v in state_dict.items():
+            if 'mm_dcnv3' in k:
+                continue
+            if 'dcnv3' not in k and 'clip_projector' not in k:
+                continue
+            new_k = k.replace('dcnv3.', '')
+            new_state_dict[new_k] = v
+        new_state_dict['fc_norm.weight'] = state_dict[
+            'clip.classifier_ln.weight']
+        new_state_dict['fc_norm.bias'] = state_dict['clip.classifier_ln.bias']
+        new_state_dict['head.weight'] = state_dict['clip.classifier.weight']
+        new_state_dict['head.bias'] = state_dict['clip.classifier.bias']
+        state_dict = new_state_dict
+
+    # delete relative_position_index since we always re-init it
+    relative_position_index_keys = [
+        k for k in state_dict.keys() if 'relative_position_index' in k
+    ]
+    for k in relative_position_index_keys:
+        del state_dict[k]
+
+    # delete relative_coords_table since we always re-init it
+    relative_position_index_keys = [
+        k for k in state_dict.keys() if 'relative_coords_table' in k
+    ]
+    for k in relative_position_index_keys:
+        del state_dict[k]
+
+    # delete attn_mask since we always re-init it
+    attn_mask_keys = [k for k in state_dict.keys() if 'attn_mask' in k]
+    for k in attn_mask_keys:
+        del state_dict[k]
+
+    # bicubic interpolate relative_position_bias_table if not match
+    relative_position_bias_table_keys = [
+        k for k in state_dict.keys() if 'relative_position_bias_table' in k
+    ]
+    for k in relative_position_bias_table_keys:
+        relative_position_bias_table_pretrained = state_dict[k]
+        relative_position_bias_table_current = model.state_dict()[k]
+        L1, nH1 = relative_position_bias_table_pretrained.size()
+        L2, nH2 = relative_position_bias_table_current.size()
+        if nH1 != nH2:
+            logger.warning(f'Error in loading {k}, passing......')
+        else:
+            if L1 != L2:
+                # bicubic interpolate relative_position_bias_table if not match
+                S1 = int(L1**0.5)
+                S2 = int(L2**0.5)
+                relative_position_bias_table_pretrained_resized = torch.nn.functional.interpolate(
+                    relative_position_bias_table_pretrained.permute(1, 0).view(
+                        1, nH1, S1, S1),
+                    size=(S2, S2),
+                    mode='bicubic')
+                state_dict[
+                    k] = relative_position_bias_table_pretrained_resized.view(
+                        nH2, L2).permute(1, 0)
+
+    # bicubic interpolate absolute_pos_embed if not match
+    absolute_pos_embed_keys = [
+        k for k in state_dict.keys() if 'absolute_pos_embed' in k
+    ]
+    for k in absolute_pos_embed_keys:
+        # dpe
+        absolute_pos_embed_pretrained = state_dict[k]
+        absolute_pos_embed_current = model.state_dict()[k]
+        _, L1, C1 = absolute_pos_embed_pretrained.size()
+        _, L2, C2 = absolute_pos_embed_current.size()
+        if C1 != C1:
+            logger.warning(f'Error in loading {k}, passing......')
+        else:
+            if L1 != L2:
+                S1 = int(L1**0.5)
+                S2 = int(L2**0.5)
+                absolute_pos_embed_pretrained = absolute_pos_embed_pretrained.reshape(
+                    -1, S1, S1, C1)
+                absolute_pos_embed_pretrained = absolute_pos_embed_pretrained.permute(
+                    0, 3, 1, 2)
+                absolute_pos_embed_pretrained_resized = torch.nn.functional.interpolate(
+                    absolute_pos_embed_pretrained,
+                    size=(S2, S2),
+                    mode='bicubic')
+                absolute_pos_embed_pretrained_resized = absolute_pos_embed_pretrained_resized.permute(
+                    0, 2, 3, 1)
+                absolute_pos_embed_pretrained_resized = absolute_pos_embed_pretrained_resized.flatten(
+                    1, 2)
+                state_dict[k] = absolute_pos_embed_pretrained_resized
+
+    # check classifier, if not match, then re-init classifier to zero
+    if 'head.bias' in state_dict:
+        head_bias_pretrained = state_dict['head.bias']
+        Nc1 = head_bias_pretrained.shape[0]
+        Nc2 = model.head.bias.shape[0]
+
+        if (Nc1 != Nc2):
+            if config.TRAIN.RAND_INIT_FT_HEAD:
+                model.head.weight.data = model.head.weight.data * 0.001
+                model.head.bias.data = model.head.bias.data * 0.001
+                del state_dict['head.weight']
+                del state_dict['head.bias']
+                logger.warning(
+                    f'Error in loading classifier head, re-init classifier head to 0'
+                )
+            elif Nc1 == 21841 and Nc2 == 1000:
+                logger.info(
+                    'loading ImageNet-22K weight to ImageNet-1K ......')
+                map22kto1k_path = 'meta_data/map22kto1k.txt'
+                logger.info(map22kto1k_path)
+                with open(map22kto1k_path) as f:
+                    map22kto1k = f.readlines()
+                map22kto1k = [int(id22k.strip()) for id22k in map22kto1k]
+                state_dict['head.weight'] = state_dict['head.weight'][
+                    map22kto1k, :]
+                state_dict['head.bias'] = state_dict['head.bias'][map22kto1k]
+
+    msg = model.load_state_dict(state_dict, strict=False)
+    logger.warning(msg)
+
+    # from IPython import embed
+    # embed()
+
+    logger.info(f'=> loaded successfully {config.MODEL.PRETRAINED}')
+
+    del checkpoint
+    torch.cuda.empty_cache()
+
+
+def convert_22k_head_to_1k(model, logger):
+    head_weight = model.module.head.weight
+    head_bias = model.module.head.bias
+    Nc1 = head_bias.shape[0]
+
+    if Nc1 == 21841:
+        logger.info('converting ImageNet-22K head to ImageNet-1K ......')
+        map22kto1k_path = 'meta_data/map22kto1k.txt'
+        logger.info(map22kto1k_path)
+        with open(map22kto1k_path) as f:
+            map22kto1k = f.readlines()
+        map22kto1k = [int(id22k.strip()) for id22k in map22kto1k]
+        model.module.head.weight = torch.nn.Parameter(
+            head_weight[map22kto1k, :])
+        model.module.head.bias = torch.nn.Parameter(head_bias[map22kto1k])
+    else:
+        logger.warning(f'Error in converting classifier head')
+
+    return model
+
+
+def save_checkpoint(config,
+                    epoch,
+                    model,
+                    max_accuracy,
+                    optimizer,
+                    lr_scheduler,
+                    scaler,
+                    logger,
+                    model_ema=None,
+                    max_accuracy_ema=None,
+                    ema_decay=None,
+                    model_ems=None,
+                    max_accuracy_ems=None,
+                    ems_model_num=None,
+                    best=None):
+
+    save_state = {
+        'model': model.state_dict(),
+        'optimizer': optimizer.state_dict(),
+        'lr_scheduler': lr_scheduler.state_dict(),
+        'max_accuracy': max_accuracy,
+        'epoch': epoch,
+        'config': config
+    }
+    if model_ema is not None:
+        save_state['model_ema'] = get_state_dict(model_ema)
+    if max_accuracy_ema is not None:
+        save_state['max_accuracy_ema'] = max_accuracy_ema
+    if ema_decay is not None:
+        save_state['ema_decay'] = ema_decay
+    if model_ems is not None:
+        save_state['model_ems'] = get_state_dict(model_ems)
+    if max_accuracy_ems is not None:
+        save_state['max_accuracy_ems'] = max_accuracy_ems
+    if ems_model_num is not None:
+        save_state['ems_model_num'] = ems_model_num
+    if config.AMP_OPT_LEVEL != 'O0':
+        # save_state['amp'] = amp.state_dict()
+        save_state['amp'] = scaler.state_dict()
+    if best is None:
+        save_path = os.path.join(config.OUTPUT, f'ckpt_epoch_{epoch}.pth')
+    else:
+        save_path = os.path.join(config.OUTPUT, f'ckpt_epoch_{best}.pth')
+    logger.info(f'{save_path} saving......')
+    torch.save(save_state, save_path)
+    logger.info(f'{save_path} saved !!!')
+
+    if dist.get_rank() == 0 and isinstance(epoch, int):
+        to_del = epoch - config.SAVE_CKPT_NUM * config.SAVE_FREQ
+        old_ckpt = os.path.join(config.OUTPUT, f'ckpt_epoch_{to_del}.pth')
+        if os.path.exists(old_ckpt):
+            os.remove(old_ckpt)
+
+
+def get_grad_norm(parameters, norm_type=2):
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = list(filter(lambda p: p.grad is not None, parameters))
+    norm_type = float(norm_type)
+    total_norm = 0
+    for p in parameters:
+        param_norm = p.grad.data.norm(norm_type)
+        total_norm += param_norm.item()**norm_type
+    total_norm = total_norm**(1. / norm_type)
+    return total_norm
+
+
+def auto_resume_helper(output_dir):
+    checkpoints = os.listdir(output_dir)
+    checkpoints = [ckpt for ckpt in checkpoints if ckpt.endswith('pth')]
+    print(f'All checkpoints founded in {output_dir}: {checkpoints}')
+    if len(checkpoints) > 0:
+        latest_checkpoint = max(
+            [os.path.join(output_dir, d) for d in checkpoints],
+            key=os.path.getmtime)
+        print(f'The latest checkpoint founded: {latest_checkpoint}')
+        resume_file = latest_checkpoint
+    else:
+        resume_file = None
+    return resume_file
+
+
+def reduce_tensor(tensor):
+    rt = tensor.clone()
+    dist.all_reduce(rt, op=dist.ReduceOp.SUM)
+    rt /= dist.get_world_size()
+    return rt
+
+
+# https://github.com/facebookresearch/ConvNeXt/blob/main/utils.py
+class NativeScalerWithGradNormCount:
+    state_dict_key = 'amp_scaler'
+
+    def __init__(self):
+        self._scaler = torch.cuda.amp.GradScaler()
+
+    def __call__(self,
+                 loss,
+                 optimizer,
+                 clip_grad=None,
+                 parameters=None,
+                 create_graph=False,
+                 update_grad=True):
+        self._scaler.scale(loss).backward(create_graph=create_graph)
+        if update_grad:
+            if clip_grad is not None:
+                assert parameters is not None
+                self._scaler.unscale_(
+                    optimizer
+                )  # unscale the gradients of optimizer's assigned params in-place
+                norm = torch.nn.utils.clip_grad_norm_(parameters, clip_grad)
+            else:
+                self._scaler.unscale_(optimizer)
+                norm = get_grad_norm(parameters)
+            self._scaler.step(optimizer)
+            self._scaler.update()
+        else:
+            norm = None
+        return norm
+
+    def state_dict(self):
+        return self._scaler.state_dict()
+
+    def load_state_dict(self, state_dict):
+        self._scaler.load_state_dict(state_dict)
+
+
+class MyAverageMeter(object):
+    """Computes and stores the average and current value."""
+
+    def __init__(self, max_len=-1):
+        self.val_list = []
+        self.count = []
+        self.max_len = max_len
+        self.val = 0
+        self.avg = 0
+        self.var = 0
+
+    def update(self, val):
+        self.val = val
+        self.avg = 0
+        self.var = 0
+        if not math.isnan(val) and not math.isinf(val):
+            self.val_list.append(val)
+        if self.max_len > 0 and len(self.val_list) > self.max_len:
+            self.val_list = self.val_list[-self.max_len:]
+        if len(self.val_list) > 0:
+            self.avg = np.mean(np.array(self.val_list))
+            self.var = np.std(np.array(self.val_list))

detection/README.md

@@ -0,0 +1,93 @@
+# FlashInternImage for Object Detection
+
+This folder contains the implementation of the FlashInternImage for object detection. 
+
+Our detection code is developed on top of [MMDetection v2.28.1](https://github.com/open-mmlab/mmdetection/tree/v2.28.1).
+
+
+## Usage
+
+### Install
+
+- Clone this repo:
+
+```bash
+git clone https://github.com/OpenGVLab/DCNv4.git
+cd DCNv4
+```
+
+- Create a conda virtual environment and activate it:
+
+```bash
+conda create -n dcnv4 python=3.7 -y
+conda activate dcnv4
+```
+
+- Install `CUDA>=10.2` with `cudnn>=7` following
+  the [official installation instructions](https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html)
+- Install `PyTorch>=1.10.0` and `torchvision>=0.9.0` with `CUDA>=10.2`:
+
+For examples, to install torch==1.11 with CUDA==11.3:
+```bash
+pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113  -f https://download.pytorch.org/whl/torch_stable.html
+```
+
+- Install `timm==0.6.11` and `mmcv-full==1.5.0`:
+
+```bash
+pip install -U openmim
+mim install mmcv-full==1.5.0
+pip install timm==0.6.11 mmdet==2.28.1
+```
+
+- Install other requirements:
+
+```bash
+pip install opencv-python termcolor yacs pyyaml scipy
+```
+
+- Install DCNv4
+```bash
+pip install DCNv4==latest
+```
+
+
+### Data Preparation
+
+Prepare COCO according to the guidelines in [MMDetection v2.28.1](https://github.com/open-mmlab/mmdetection/resolve/master/docs/en/1_exist_data_model.md).
+
+
+### Evaluation
+
+To evaluate our `FlashInternImage` on COCO val, run:
+
+```bash
+sh dist_test.sh <config-file> <checkpoint> <gpu-num> --eval bbox segm
+```
+
+For example, to evaluate the `FlashInternImage-T` with a single GPU:
+
+```bash
+python test.py configs/coco/mask_rcnn_flash_intern_image_t_fpn_1x_coco.py checkpoint_dir/det/mask_rcnn_flash_internimage_t_fpn_1x_coco.pth --eval bbox segm
+```
+
+For example, to evaluate the `FlashInternImage-B` with a single node with 8 GPUs:
+
+```bash
+sh dist_test.sh configs/coco/mask_rcnn_flash_intern_image_b_fpn_1x_coco.py checkpoint_dir/det/mask_rcnn_flash_internimage_b_fpn_1x_coco.py 8 --eval bbox segm
+```
+
+### Training on COCO
+
+To train an `FlashInternImage` on COCO, run:
+
+```bash
+sh dist_train.sh <config-file> <gpu-num>
+```
+
+For example, to train `FlashInternImage-T` with 8 GPU on 1 node, run:
+
+```bash
+sh dist_train.sh configs/coco/mask_rcnn_flash_intern_image_t_fpn_1x_coco.py 8
+```
+

detection/configs/_base_/datasets/coco_detection.py

@@ -0,0 +1,49 @@
+# dataset settings
+dataset_type = 'CocoDataset'
+data_root = 'data/coco/'
+img_norm_cfg = dict(
+    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='LoadAnnotations', with_bbox=True),
+    dict(type='Resize', img_scale=(1333, 800), keep_ratio=True),
+    dict(type='RandomFlip', flip_ratio=0.5),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='Pad', size_divisor=32),
+    dict(type='DefaultFormatBundle'),
+    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels']),
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='MultiScaleFlipAug',
+        img_scale=(1333, 800),
+        flip=False,
+        transforms=[
+            dict(type='Resize', keep_ratio=True),
+            dict(type='RandomFlip'),
+            dict(type='Normalize', **img_norm_cfg),
+            dict(type='Pad', size_divisor=32),
+            dict(type='ImageToTensor', keys=['img']),
+            dict(type='Collect', keys=['img']),
+        ])
+]
+data = dict(
+    samples_per_gpu=2,
+    workers_per_gpu=2,
+    train=dict(
+        type=dataset_type,
+        ann_file=data_root + 'annotations/instances_train2017.json',
+        img_prefix=data_root + 'train2017/',
+        pipeline=train_pipeline),
+    val=dict(
+        type=dataset_type,
+        ann_file=data_root + 'annotations/instances_val2017.json',
+        img_prefix=data_root + 'val2017/',
+        pipeline=test_pipeline),
+    test=dict(
+        type=dataset_type,
+        ann_file=data_root + 'annotations/instances_val2017.json',
+        img_prefix=data_root + 'val2017/',
+        pipeline=test_pipeline))
+evaluation = dict(interval=1, metric='bbox', classwise=True)

detection/configs/_base_/datasets/coco_instance.py

@@ -0,0 +1,49 @@
+# dataset settings
+dataset_type = 'CocoDataset'
+data_root = 'data/coco/'
+img_norm_cfg = dict(
+    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='LoadAnnotations', with_bbox=True, with_mask=True),
+    dict(type='Resize', img_scale=(1333, 800), keep_ratio=True),
+    dict(type='RandomFlip', flip_ratio=0.5),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='Pad', size_divisor=32),
+    dict(type='DefaultFormatBundle'),
+    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels', 'gt_masks']),
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='MultiScaleFlipAug',
+        img_scale=(1333, 800),
+        flip=False,
+        transforms=[
+            dict(type='Resize', keep_ratio=True),
+            dict(type='RandomFlip'),
+            dict(type='Normalize', **img_norm_cfg),
+            dict(type='Pad', size_divisor=32),
+            dict(type='ImageToTensor', keys=['img']),
+            dict(type='Collect', keys=['img']),
+        ])
+]
+data = dict(
+    samples_per_gpu=2,
+    workers_per_gpu=2,
+    train=dict(
+        type=dataset_type,
+        ann_file=data_root + 'annotations/instances_train2017.json',
+        img_prefix=data_root + 'train2017/',
+        pipeline=train_pipeline),
+    val=dict(
+        type=dataset_type,
+        ann_file=data_root + 'annotations/instances_val2017.json',
+        img_prefix=data_root + 'val2017/',
+        pipeline=test_pipeline),
+    test=dict(
+        type=dataset_type,
+        ann_file=data_root + 'annotations/instances_val2017.json',
+        img_prefix=data_root + 'val2017/',
+        pipeline=test_pipeline))
+evaluation = dict(metric=['bbox', 'segm'], classwise=True)

detection/configs/_base_/datasets/crowd_human.py

@@ -0,0 +1,54 @@
+# dataset settings
+dataset_type = 'CrowdHumanDataset'
+data_root = 'data/CrowdHuman/'
+classes = ('person',)
+img_norm_cfg = dict(
+    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='LoadAnnotations', with_bbox=True),
+    dict(type='Resize', img_scale=(1333, 800), keep_ratio=True),
+    dict(type='RandomFlip', flip_ratio=0.5),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='Pad', size_divisor=32),
+    dict(type='DefaultFormatBundle'),
+    dict(type='Collect', keys=['img', 'gt_bboxes', 'gt_labels']),
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='MultiScaleFlipAug',
+        img_scale=(1333, 800),
+        flip=False,
+        transforms=[
+            dict(type='Resize', keep_ratio=True),
+            dict(type='RandomFlip'),
+            dict(type='Normalize', **img_norm_cfg),
+            dict(type='Pad', size_divisor=32),
+            dict(type='ImageToTensor', keys=['img']),
+            dict(type='Collect', keys=['img']),
+        ])
+]
+data = dict(
+    samples_per_gpu=2,
+    workers_per_gpu=2,
+    train=dict(
+        type=dataset_type,
+        classes=classes,
+        filter_empty_gt=True,
+        ann_file=data_root + 'annotations/annotation_train.json',
+        img_prefix=data_root + 'Images',
+        pipeline=train_pipeline),
+    val=dict(
+        type=dataset_type,
+        classes=classes,
+        ann_file=data_root + 'annotations/annotation_val.json',
+        img_prefix=data_root + 'Images',
+        pipeline=test_pipeline),
+    test=dict(
+        type=dataset_type,
+        classes=classes,
+        ann_file=data_root + 'annotations/annotation_val.json',
+        img_prefix=data_root + 'Images',
+        pipeline=test_pipeline))
+evaluation = dict(interval=100, metric='bbox')

detection/configs/_base_/default_runtime.py

@@ -0,0 +1,16 @@
+checkpoint_config = dict(interval=1)
+# yapf:disable
+log_config = dict(
+    interval=50,
+    hooks=[
+        dict(type='TextLoggerHook'),
+        # dict(type='TensorboardLoggerHook')
+    ])
+# yapf:enable
+custom_hooks = [dict(type='NumClassCheckHook')]
+
+dist_params = dict(backend='nccl')
+log_level = 'INFO'
+load_from = None
+resume_from = None
+workflow = [('train', 1)]

detection/configs/_base_/models/cascade_mask_rcnn_convnext_fpn.py

@@ -0,0 +1,208 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+# model settings
+model = dict(
+    type='CascadeRCNN',
+    pretrained=None,
+    backbone=dict(
+        type='ConvNeXt_speed',
+        in_chans=3,
+        depths=[3, 3, 9, 3], 
+        dims=[96, 192, 384, 768], 
+        drop_path_rate=0.2,
+        layer_scale_init_value=1e-6,
+        out_indices=[0, 1, 2, 3],
+    ),
+    neck=dict(
+        type='FPN',
+        in_channels=[128, 256, 512, 1024],
+        out_channels=256,
+        num_outs=5),
+    rpn_head=dict(
+        type='RPNHead',
+        in_channels=256,
+        feat_channels=256,
+        anchor_generator=dict(
+            type='AnchorGenerator',
+            scales=[8],
+            ratios=[0.5, 1.0, 2.0],
+            strides=[4, 8, 16, 32, 64]),
+        bbox_coder=dict(
+            type='DeltaXYWHBBoxCoder',
+            target_means=[.0, .0, .0, .0],
+            target_stds=[1.0, 1.0, 1.0, 1.0]),
+        loss_cls=dict(
+            type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0),
+        loss_bbox=dict(type='SmoothL1Loss', beta=1.0 / 9.0, loss_weight=1.0)),
+    roi_head=dict(
+        type='CascadeRoIHead',
+        num_stages=3,
+        stage_loss_weights=[1, 0.5, 0.25],
+        bbox_roi_extractor=dict(
+            type='SingleRoIExtractor',
+            roi_layer=dict(type='RoIAlign', output_size=7, sampling_ratio=0),
+            out_channels=256,
+            featmap_strides=[4, 8, 16, 32]),
+        bbox_head=[
+            dict(
+                type='Shared2FCBBoxHead',
+                in_channels=256,
+                fc_out_channels=1024,
+                roi_feat_size=7,
+                num_classes=80,
+                bbox_coder=dict(
+                    type='DeltaXYWHBBoxCoder',
+                    target_means=[0., 0., 0., 0.],
+                    target_stds=[0.1, 0.1, 0.2, 0.2]),
+                reg_class_agnostic=True,
+                loss_cls=dict(
+                    type='CrossEntropyLoss',
+                    use_sigmoid=False,
+                    loss_weight=1.0),
+                loss_bbox=dict(type='SmoothL1Loss', beta=1.0,
+                               loss_weight=1.0)),
+            dict(
+                type='Shared2FCBBoxHead',
+                in_channels=256,
+                fc_out_channels=1024,
+                roi_feat_size=7,
+                num_classes=80,
+                bbox_coder=dict(
+                    type='DeltaXYWHBBoxCoder',
+                    target_means=[0., 0., 0., 0.],
+                    target_stds=[0.05, 0.05, 0.1, 0.1]),
+                reg_class_agnostic=True,
+                loss_cls=dict(
+                    type='CrossEntropyLoss',
+                    use_sigmoid=False,
+                    loss_weight=1.0),
+                loss_bbox=dict(type='SmoothL1Loss', beta=1.0,
+                               loss_weight=1.0)),
+            dict(
+                type='Shared2FCBBoxHead',
+                in_channels=256,
+                fc_out_channels=1024,
+                roi_feat_size=7,
+                num_classes=80,
+                bbox_coder=dict(
+                    type='DeltaXYWHBBoxCoder',
+                    target_means=[0., 0., 0., 0.],
+                    target_stds=[0.033, 0.033, 0.067, 0.067]),
+                reg_class_agnostic=True,
+                loss_cls=dict(
+                    type='CrossEntropyLoss',
+                    use_sigmoid=False,
+                    loss_weight=1.0),
+                loss_bbox=dict(type='SmoothL1Loss', beta=1.0, loss_weight=1.0))
+        ],
+        mask_roi_extractor=dict(
+            type='SingleRoIExtractor',
+            roi_layer=dict(type='RoIAlign', output_size=14, sampling_ratio=0),
+            out_channels=256,
+            featmap_strides=[4, 8, 16, 32]),
+        mask_head=dict(
+            type='FCNMaskHead',
+            num_convs=4,
+            in_channels=256,
+            conv_out_channels=256,
+            num_classes=80,
+            loss_mask=dict(
+                type='CrossEntropyLoss', use_mask=True, loss_weight=1.0))),
+    # model training and testing settings
+    train_cfg = dict(
+        rpn=dict(
+            assigner=dict(
+                type='MaxIoUAssigner',
+                pos_iou_thr=0.7,
+                neg_iou_thr=0.3,
+                min_pos_iou=0.3,
+                match_low_quality=True,
+                ignore_iof_thr=-1),
+            sampler=dict(
+                type='RandomSampler',
+                num=256,
+                pos_fraction=0.5,
+                neg_pos_ub=-1,
+                add_gt_as_proposals=False),
+            allowed_border=0,
+            pos_weight=-1,
+            debug=False),
+        rpn_proposal=dict(
+            nms_across_levels=False,
+            nms_pre=2000,
+            nms_post=2000,
+            max_per_img=2000,
+            nms=dict(type='nms', iou_threshold=0.7),
+            min_bbox_size=0),
+        rcnn=[
+            dict(
+                assigner=dict(
+                    type='MaxIoUAssigner',
+                    pos_iou_thr=0.5,
+                    neg_iou_thr=0.5,
+                    min_pos_iou=0.5,
+                    match_low_quality=False,
+                    ignore_iof_thr=-1),
+                sampler=dict(
+                    type='RandomSampler',
+                    num=512,
+                    pos_fraction=0.25,
+                    neg_pos_ub=-1,
+                    add_gt_as_proposals=True),
+                mask_size=28,
+                pos_weight=-1,
+                debug=False),
+            dict(
+                assigner=dict(
+                    type='MaxIoUAssigner',
+                    pos_iou_thr=0.6,
+                    neg_iou_thr=0.6,
+                    min_pos_iou=0.6,
+                    match_low_quality=False,
+                    ignore_iof_thr=-1),
+                sampler=dict(
+                    type='RandomSampler',
+                    num=512,
+                    pos_fraction=0.25,
+                    neg_pos_ub=-1,
+                    add_gt_as_proposals=True),
+                mask_size=28,
+                pos_weight=-1,
+                debug=False),
+            dict(
+                assigner=dict(
+                    type='MaxIoUAssigner',
+                    pos_iou_thr=0.7,
+                    neg_iou_thr=0.7,
+                    min_pos_iou=0.7,
+                    match_low_quality=False,
+                    ignore_iof_thr=-1),
+                sampler=dict(
+                    type='RandomSampler',
+                    num=512,
+                    pos_fraction=0.25,
+                    neg_pos_ub=-1,
+                    add_gt_as_proposals=True),
+                mask_size=28,
+                pos_weight=-1,
+                debug=False)
+        ]),
+    test_cfg = dict(
+        rpn=dict(
+            nms_across_levels=False,
+            nms_pre=1000,
+            nms_post=1000,
+            max_per_img=1000,
+            nms=dict(type='nms', iou_threshold=0.7),
+            min_bbox_size=0),
+        rcnn=dict(
+            score_thr=0.05,
+            nms=dict(type='nms', iou_threshold=0.5),
+            max_per_img=100,
+            mask_thr_binary=0.5)))

detection/configs/_base_/models/cascade_mask_rcnn_r50_fpn.py

@@ -0,0 +1,196 @@
+# model settings
+model = dict(
+    type='CascadeRCNN',
+    backbone=dict(
+        type='ResNet',
+        depth=50,
+        num_stages=4,
+        out_indices=(0, 1, 2, 3),
+        frozen_stages=1,
+        norm_cfg=dict(type='BN', requires_grad=True),
+        norm_eval=True,
+        style='pytorch',
+        init_cfg=dict(type='Pretrained', checkpoint='torchvision://resnet50')),
+    neck=dict(
+        type='FPN',
+        in_channels=[256, 512, 1024, 2048],
+        out_channels=256,
+        num_outs=5),
+    rpn_head=dict(
+        type='RPNHead',
+        in_channels=256,
+        feat_channels=256,
+        anchor_generator=dict(
+            type='AnchorGenerator',
+            scales=[8],
+            ratios=[0.5, 1.0, 2.0],
+            strides=[4, 8, 16, 32, 64]),
+        bbox_coder=dict(
+            type='DeltaXYWHBBoxCoder',
+            target_means=[.0, .0, .0, .0],
+            target_stds=[1.0, 1.0, 1.0, 1.0]),
+        loss_cls=dict(
+            type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0),
+        loss_bbox=dict(type='SmoothL1Loss', beta=1.0 / 9.0, loss_weight=1.0)),
+    roi_head=dict(
+        type='CascadeRoIHead',
+        num_stages=3,
+        stage_loss_weights=[1, 0.5, 0.25],
+        bbox_roi_extractor=dict(
+            type='SingleRoIExtractor',
+            roi_layer=dict(type='RoIAlign', output_size=7, sampling_ratio=0),
+            out_channels=256,
+            featmap_strides=[4, 8, 16, 32]),
+        bbox_head=[
+            dict(
+                type='Shared2FCBBoxHead',
+                in_channels=256,
+                fc_out_channels=1024,
+                roi_feat_size=7,
+                num_classes=80,
+                bbox_coder=dict(
+                    type='DeltaXYWHBBoxCoder',
+                    target_means=[0., 0., 0., 0.],
+                    target_stds=[0.1, 0.1, 0.2, 0.2]),
+                reg_class_agnostic=True,
+                loss_cls=dict(
+                    type='CrossEntropyLoss',
+                    use_sigmoid=False,
+                    loss_weight=1.0),
+                loss_bbox=dict(type='SmoothL1Loss', beta=1.0,
+                               loss_weight=1.0)),
+            dict(
+                type='Shared2FCBBoxHead',
+                in_channels=256,
+                fc_out_channels=1024,
+                roi_feat_size=7,
+                num_classes=80,
+                bbox_coder=dict(
+                    type='DeltaXYWHBBoxCoder',
+                    target_means=[0., 0., 0., 0.],
+                    target_stds=[0.05, 0.05, 0.1, 0.1]),
+                reg_class_agnostic=True,
+                loss_cls=dict(
+                    type='CrossEntropyLoss',
+                    use_sigmoid=False,
+                    loss_weight=1.0),
+                loss_bbox=dict(type='SmoothL1Loss', beta=1.0,
+                               loss_weight=1.0)),
+            dict(
+                type='Shared2FCBBoxHead',
+                in_channels=256,
+                fc_out_channels=1024,
+                roi_feat_size=7,
+                num_classes=80,
+                bbox_coder=dict(
+                    type='DeltaXYWHBBoxCoder',
+                    target_means=[0., 0., 0., 0.],
+                    target_stds=[0.033, 0.033, 0.067, 0.067]),
+                reg_class_agnostic=True,
+                loss_cls=dict(
+                    type='CrossEntropyLoss',
+                    use_sigmoid=False,
+                    loss_weight=1.0),
+                loss_bbox=dict(type='SmoothL1Loss', beta=1.0, loss_weight=1.0))
+        ],
+        mask_roi_extractor=dict(
+            type='SingleRoIExtractor',
+            roi_layer=dict(type='RoIAlign', output_size=14, sampling_ratio=0),
+            out_channels=256,
+            featmap_strides=[4, 8, 16, 32]),
+        mask_head=dict(
+            type='FCNMaskHead',
+            num_convs=4,
+            in_channels=256,
+            conv_out_channels=256,
+            num_classes=80,
+            loss_mask=dict(
+                type='CrossEntropyLoss', use_mask=True, loss_weight=1.0))),
+    # model training and testing settings
+    train_cfg=dict(
+        rpn=dict(
+            assigner=dict(
+                type='MaxIoUAssigner',
+                pos_iou_thr=0.7,
+                neg_iou_thr=0.3,
+                min_pos_iou=0.3,
+                match_low_quality=True,
+                ignore_iof_thr=-1),
+            sampler=dict(
+                type='RandomSampler',
+                num=256,
+                pos_fraction=0.5,
+                neg_pos_ub=-1,
+                add_gt_as_proposals=False),
+            allowed_border=0,
+            pos_weight=-1,
+            debug=False),
+        rpn_proposal=dict(
+            nms_pre=2000,
+            max_per_img=2000,
+            nms=dict(type='nms', iou_threshold=0.7),
+            min_bbox_size=0),
+        rcnn=[
+            dict(
+                assigner=dict(
+                    type='MaxIoUAssigner',
+                    pos_iou_thr=0.5,
+                    neg_iou_thr=0.5,
+                    min_pos_iou=0.5,
+                    match_low_quality=False,
+                    ignore_iof_thr=-1),
+                sampler=dict(
+                    type='RandomSampler',
+                    num=512,
+                    pos_fraction=0.25,
+                    neg_pos_ub=-1,
+                    add_gt_as_proposals=True),
+                mask_size=28,
+                pos_weight=-1,
+                debug=False),
+            dict(
+                assigner=dict(
+                    type='MaxIoUAssigner',
+                    pos_iou_thr=0.6,
+                    neg_iou_thr=0.6,
+                    min_pos_iou=0.6,
+                    match_low_quality=False,
+                    ignore_iof_thr=-1),
+                sampler=dict(
+                    type='RandomSampler',
+                    num=512,
+                    pos_fraction=0.25,
+                    neg_pos_ub=-1,
+                    add_gt_as_proposals=True),
+                mask_size=28,
+                pos_weight=-1,
+                debug=False),
+            dict(
+                assigner=dict(
+                    type='MaxIoUAssigner',
+                    pos_iou_thr=0.7,
+                    neg_iou_thr=0.7,
+                    min_pos_iou=0.7,
+                    match_low_quality=False,
+                    ignore_iof_thr=-1),
+                sampler=dict(
+                    type='RandomSampler',
+                    num=512,
+                    pos_fraction=0.25,
+                    neg_pos_ub=-1,
+                    add_gt_as_proposals=True),
+                mask_size=28,
+                pos_weight=-1,
+                debug=False)
+        ]),
+    test_cfg=dict(
+        rpn=dict(
+            nms_pre=1000,
+            max_per_img=1000,
+            nms=dict(type='nms', iou_threshold=0.7),
+            min_bbox_size=0),
+        rcnn=dict(
+            score_thr=0.05,
+            nms=dict(type='nms', iou_threshold=0.5),
+            max_per_img=100,
+            mask_thr_binary=0.5)))

detection/configs/_base_/models/cascade_mask_rcnn_r50_fpn_crowdhuman.py

@@ -0,0 +1,183 @@
+# model settings
+model = dict(
+    type='CascadeRCNN',
+    backbone=dict(
+        type='ResNet',
+        depth=50,
+        num_stages=4,
+        out_indices=(0, 1, 2, 3),
+        frozen_stages=1,
+        norm_cfg=dict(type='BN', requires_grad=True),
+        norm_eval=True,
+        style='pytorch',
+        init_cfg=dict(type='Pretrained', checkpoint='torchvision://resnet50')),
+    neck=dict(
+        type='FPN',
+        in_channels=[256, 512, 1024, 2048],
+        out_channels=256,
+        num_outs=5),
+    rpn_head=dict(
+        type='RPNHead',
+        in_channels=256,
+        feat_channels=256,
+        anchor_generator=dict(
+            type='AnchorGenerator',
+            scales=[8],
+            ratios=[0.5, 1.0, 2.0],
+            strides=[4, 8, 16, 32, 64]),
+        bbox_coder=dict(
+            type='DeltaXYWHBBoxCoder',
+            target_means=[.0, .0, .0, .0],
+            target_stds=[1.0, 1.0, 1.0, 1.0]),
+        loss_cls=dict(
+            type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0),
+        loss_bbox=dict(type='SmoothL1Loss', beta=1.0 / 9.0, loss_weight=1.0)),
+    roi_head=dict(
+        type='CascadeRoIHead',
+        num_stages=3,
+        stage_loss_weights=[1, 0.5, 0.25],
+        bbox_roi_extractor=dict(
+            type='SingleRoIExtractor',
+            roi_layer=dict(type='RoIAlign', output_size=7, sampling_ratio=0),
+            out_channels=256,
+            featmap_strides=[4, 8, 16, 32]),
+        bbox_head=[
+            dict(
+                type='Shared2FCBBoxHead',
+                in_channels=256,
+                fc_out_channels=1024,
+                roi_feat_size=7,
+                num_classes=80,
+                bbox_coder=dict(
+                    type='DeltaXYWHBBoxCoder',
+                    target_means=[0., 0., 0., 0.],
+                    target_stds=[0.1, 0.1, 0.2, 0.2]),
+                reg_class_agnostic=True,
+                loss_cls=dict(
+                    type='CrossEntropyLoss',
+                    use_sigmoid=False,
+                    loss_weight=1.0),
+                loss_bbox=dict(type='SmoothL1Loss', beta=1.0,
+                               loss_weight=1.0)),
+            dict(
+                type='Shared2FCBBoxHead',
+                in_channels=256,
+                fc_out_channels=1024,
+                roi_feat_size=7,
+                num_classes=80,
+                bbox_coder=dict(
+                    type='DeltaXYWHBBoxCoder',
+                    target_means=[0., 0., 0., 0.],
+                    target_stds=[0.05, 0.05, 0.1, 0.1]),
+                reg_class_agnostic=True,
+                loss_cls=dict(
+                    type='CrossEntropyLoss',
+                    use_sigmoid=False,
+                    loss_weight=1.0),
+                loss_bbox=dict(type='SmoothL1Loss', beta=1.0,
+                               loss_weight=1.0)),
+            dict(
+                type='Shared2FCBBoxHead',
+                in_channels=256,
+                fc_out_channels=1024,
+                roi_feat_size=7,
+                num_classes=80,
+                bbox_coder=dict(
+                    type='DeltaXYWHBBoxCoder',
+                    target_means=[0., 0., 0., 0.],
+                    target_stds=[0.033, 0.033, 0.067, 0.067]),
+                reg_class_agnostic=True,
+                loss_cls=dict(
+                    type='CrossEntropyLoss',
+                    use_sigmoid=False,
+                    loss_weight=1.0),
+                loss_bbox=dict(type='SmoothL1Loss', beta=1.0, loss_weight=1.0))
+        ],),
+    # model training and testing settings
+    train_cfg=dict(
+        rpn=dict(
+            assigner=dict(
+                type='MaxIoUAssigner',
+                pos_iou_thr=0.7,
+                neg_iou_thr=0.3,
+                min_pos_iou=0.3,
+                match_low_quality=True,
+                ignore_iof_thr=-1),
+            sampler=dict(
+                type='RandomSampler',
+                num=256,
+                pos_fraction=0.5,
+                neg_pos_ub=-1,
+                add_gt_as_proposals=False),
+            allowed_border=0,
+            pos_weight=-1,
+            debug=False),
+        rpn_proposal=dict(
+            nms_pre=2000,
+            max_per_img=2000,
+            nms=dict(type='nms', iou_threshold=0.7),
+            min_bbox_size=0),
+        rcnn=[
+            dict(
+                assigner=dict(
+                    type='MaxIoUAssigner',
+                    pos_iou_thr=0.5,
+                    neg_iou_thr=0.5,
+                    min_pos_iou=0.5,
+                    match_low_quality=False,
+                    ignore_iof_thr=-1),
+                sampler=dict(
+                    type='RandomSampler',
+                    num=512,
+                    pos_fraction=0.25,
+                    neg_pos_ub=-1,
+                    add_gt_as_proposals=True),
+                mask_size=28,
+                pos_weight=-1,
+                debug=False),
+            dict(
+                assigner=dict(
+                    type='MaxIoUAssigner',
+                    pos_iou_thr=0.6,
+                    neg_iou_thr=0.6,
+                    min_pos_iou=0.6,
+                    match_low_quality=False,
+                    ignore_iof_thr=-1),
+                sampler=dict(
+                    type='RandomSampler',
+                    num=512,
+                    pos_fraction=0.25,
+                    neg_pos_ub=-1,
+                    add_gt_as_proposals=True),
+                mask_size=28,
+                pos_weight=-1,
+                debug=False),
+            dict(
+                assigner=dict(
+                    type='MaxIoUAssigner',
+                    pos_iou_thr=0.7,
+                    neg_iou_thr=0.7,
+                    min_pos_iou=0.7,
+                    match_low_quality=False,
+                    ignore_iof_thr=-1),
+                sampler=dict(
+                    type='RandomSampler',
+                    num=512,
+                    pos_fraction=0.25,
+                    neg_pos_ub=-1,
+                    add_gt_as_proposals=True),
+                mask_size=28,
+                pos_weight=-1,
+                debug=False)
+        ]),
+    test_cfg=dict(
+        rpn=dict(
+            nms_pre=1000,
+            max_per_img=1000,
+            nms=dict(type='nms', iou_threshold=0.7),
+            min_bbox_size=0),
+        rcnn=dict(
+            score_thr=0.05,
+            nms=dict(type='nms', iou_threshold=0.5),
+            max_per_img=100,
+            mask_thr_binary=0.5)))

detection/configs/_base_/models/cascade_rcnn_r50_fpn.py

@@ -0,0 +1,179 @@
+# model settings
+model = dict(
+    type='CascadeRCNN',
+    backbone=dict(
+        type='ResNet',
+        depth=50,
+        num_stages=4,
+        out_indices=(0, 1, 2, 3),
+        frozen_stages=1,
+        norm_cfg=dict(type='BN', requires_grad=True),
+        norm_eval=True,
+        style='pytorch',
+        init_cfg=dict(type='Pretrained', checkpoint='torchvision://resnet50')),
+    neck=dict(
+        type='FPN',
+        in_channels=[256, 512, 1024, 2048],
+        out_channels=256,
+        num_outs=5),
+    rpn_head=dict(
+        type='RPNHead',
+        in_channels=256,
+        feat_channels=256,
+        anchor_generator=dict(
+            type='AnchorGenerator',
+            scales=[8],
+            ratios=[0.5, 1.0, 2.0],
+            strides=[4, 8, 16, 32, 64]),
+        bbox_coder=dict(
+            type='DeltaXYWHBBoxCoder',
+            target_means=[.0, .0, .0, .0],
+            target_stds=[1.0, 1.0, 1.0, 1.0]),
+        loss_cls=dict(
+            type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0),
+        loss_bbox=dict(type='SmoothL1Loss', beta=1.0 / 9.0, loss_weight=1.0)),
+    roi_head=dict(
+        type='CascadeRoIHead',
+        num_stages=3,
+        stage_loss_weights=[1, 0.5, 0.25],
+        bbox_roi_extractor=dict(
+            type='SingleRoIExtractor',
+            roi_layer=dict(type='RoIAlign', output_size=7, sampling_ratio=0),
+            out_channels=256,
+            featmap_strides=[4, 8, 16, 32]),
+        bbox_head=[
+            dict(
+                type='Shared2FCBBoxHead',
+                in_channels=256,
+                fc_out_channels=1024,
+                roi_feat_size=7,
+                num_classes=80,
+                bbox_coder=dict(
+                    type='DeltaXYWHBBoxCoder',
+                    target_means=[0., 0., 0., 0.],
+                    target_stds=[0.1, 0.1, 0.2, 0.2]),
+                reg_class_agnostic=True,
+                loss_cls=dict(
+                    type='CrossEntropyLoss',
+                    use_sigmoid=False,
+                    loss_weight=1.0),
+                loss_bbox=dict(type='SmoothL1Loss', beta=1.0,
+                               loss_weight=1.0)),
+            dict(
+                type='Shared2FCBBoxHead',
+                in_channels=256,
+                fc_out_channels=1024,
+                roi_feat_size=7,
+                num_classes=80,
+                bbox_coder=dict(
+                    type='DeltaXYWHBBoxCoder',
+                    target_means=[0., 0., 0., 0.],
+                    target_stds=[0.05, 0.05, 0.1, 0.1]),
+                reg_class_agnostic=True,
+                loss_cls=dict(
+                    type='CrossEntropyLoss',
+                    use_sigmoid=False,
+                    loss_weight=1.0),
+                loss_bbox=dict(type='SmoothL1Loss', beta=1.0,
+                               loss_weight=1.0)),
+            dict(
+                type='Shared2FCBBoxHead',
+                in_channels=256,
+                fc_out_channels=1024,
+                roi_feat_size=7,
+                num_classes=80,
+                bbox_coder=dict(
+                    type='DeltaXYWHBBoxCoder',
+                    target_means=[0., 0., 0., 0.],
+                    target_stds=[0.033, 0.033, 0.067, 0.067]),
+                reg_class_agnostic=True,
+                loss_cls=dict(
+                    type='CrossEntropyLoss',
+                    use_sigmoid=False,
+                    loss_weight=1.0),
+                loss_bbox=dict(type='SmoothL1Loss', beta=1.0, loss_weight=1.0))
+        ]),
+    # model training and testing settings
+    train_cfg=dict(
+        rpn=dict(
+            assigner=dict(
+                type='MaxIoUAssigner',
+                pos_iou_thr=0.7,
+                neg_iou_thr=0.3,
+                min_pos_iou=0.3,
+                match_low_quality=True,
+                ignore_iof_thr=-1),
+            sampler=dict(
+                type='RandomSampler',
+                num=256,
+                pos_fraction=0.5,
+                neg_pos_ub=-1,
+                add_gt_as_proposals=False),
+            allowed_border=0,
+            pos_weight=-1,
+            debug=False),
+        rpn_proposal=dict(
+            nms_pre=2000,
+            max_per_img=2000,
+            nms=dict(type='nms', iou_threshold=0.7),
+            min_bbox_size=0),
+        rcnn=[
+            dict(
+                assigner=dict(
+                    type='MaxIoUAssigner',
+                    pos_iou_thr=0.5,
+                    neg_iou_thr=0.5,
+                    min_pos_iou=0.5,
+                    match_low_quality=False,
+                    ignore_iof_thr=-1),
+                sampler=dict(
+                    type='RandomSampler',
+                    num=512,
+                    pos_fraction=0.25,
+                    neg_pos_ub=-1,
+                    add_gt_as_proposals=True),
+                pos_weight=-1,
+                debug=False),
+            dict(
+                assigner=dict(
+                    type='MaxIoUAssigner',
+                    pos_iou_thr=0.6,
+                    neg_iou_thr=0.6,
+                    min_pos_iou=0.6,
+                    match_low_quality=False,
+                    ignore_iof_thr=-1),
+                sampler=dict(
+                    type='RandomSampler',
+                    num=512,
+                    pos_fraction=0.25,
+                    neg_pos_ub=-1,
+                    add_gt_as_proposals=True),
+                pos_weight=-1,
+                debug=False),
+            dict(
+                assigner=dict(
+                    type='MaxIoUAssigner',
+                    pos_iou_thr=0.7,
+                    neg_iou_thr=0.7,
+                    min_pos_iou=0.7,
+                    match_low_quality=False,
+                    ignore_iof_thr=-1),
+                sampler=dict(
+                    type='RandomSampler',
+                    num=512,
+                    pos_fraction=0.25,
+                    neg_pos_ub=-1,
+                    add_gt_as_proposals=True),
+                pos_weight=-1,
+                debug=False)
+        ]),
+    test_cfg=dict(
+        rpn=dict(
+            nms_pre=1000,
+            max_per_img=1000,
+            nms=dict(type='nms', iou_threshold=0.7),
+            min_bbox_size=0),
+        rcnn=dict(
+            score_thr=0.05,
+            nms=dict(type='nms', iou_threshold=0.5),
+            max_per_img=100)))

detection/configs/_base_/models/fast_rcnn_r50_fpn.py

@@ -0,0 +1,62 @@
+# model settings
+model = dict(
+    type='FastRCNN',
+    backbone=dict(
+        type='ResNet',
+        depth=50,
+        num_stages=4,
+        out_indices=(0, 1, 2, 3),
+        frozen_stages=1,
+        norm_cfg=dict(type='BN', requires_grad=True),
+        norm_eval=True,
+        style='pytorch',
+        init_cfg=dict(type='Pretrained', checkpoint='torchvision://resnet50')),
+    neck=dict(
+        type='FPN',
+        in_channels=[256, 512, 1024, 2048],
+        out_channels=256,
+        num_outs=5),
+    roi_head=dict(
+        type='StandardRoIHead',
+        bbox_roi_extractor=dict(
+            type='SingleRoIExtractor',
+            roi_layer=dict(type='RoIAlign', output_size=7, sampling_ratio=0),
+            out_channels=256,
+            featmap_strides=[4, 8, 16, 32]),
+        bbox_head=dict(
+            type='Shared2FCBBoxHead',
+            in_channels=256,
+            fc_out_channels=1024,
+            roi_feat_size=7,
+            num_classes=80,
+            bbox_coder=dict(
+                type='DeltaXYWHBBoxCoder',
+                target_means=[0., 0., 0., 0.],
+                target_stds=[0.1, 0.1, 0.2, 0.2]),
+            reg_class_agnostic=False,
+            loss_cls=dict(
+                type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0),
+            loss_bbox=dict(type='L1Loss', loss_weight=1.0))),
+    # model training and testing settings
+    train_cfg=dict(
+        rcnn=dict(
+            assigner=dict(
+                type='MaxIoUAssigner',
+                pos_iou_thr=0.5,
+                neg_iou_thr=0.5,
+                min_pos_iou=0.5,
+                match_low_quality=False,
+                ignore_iof_thr=-1),
+            sampler=dict(
+                type='RandomSampler',
+                num=512,
+                pos_fraction=0.25,
+                neg_pos_ub=-1,
+                add_gt_as_proposals=True),
+            pos_weight=-1,
+            debug=False)),
+    test_cfg=dict(
+        rcnn=dict(
+            score_thr=0.05,
+            nms=dict(type='nms', iou_threshold=0.5),
+            max_per_img=100)))

detection/configs/_base_/models/faster_rcnn_r50_caffe_c4.py

@@ -0,0 +1,114 @@
+# model settings
+norm_cfg = dict(type='BN', requires_grad=False)
+model = dict(
+    type='FasterRCNN',
+    backbone=dict(
+        type='ResNet',
+        depth=50,
+        num_stages=3,
+        strides=(1, 2, 2),
+        dilations=(1, 1, 1),
+        out_indices=(2, ),
+        frozen_stages=1,
+        norm_cfg=norm_cfg,
+        norm_eval=True,
+        style='caffe',
+        init_cfg=dict(
+            type='Pretrained',
+            checkpoint='open-mmlab://detectron2/resnet50_caffe')),
+    rpn_head=dict(
+        type='RPNHead',
+        in_channels=1024,
+        feat_channels=1024,
+        anchor_generator=dict(
+            type='AnchorGenerator',
+            scales=[2, 4, 8, 16, 32],
+            ratios=[0.5, 1.0, 2.0],
+            strides=[16]),
+        bbox_coder=dict(
+            type='DeltaXYWHBBoxCoder',
+            target_means=[.0, .0, .0, .0],
+            target_stds=[1.0, 1.0, 1.0, 1.0]),
+        loss_cls=dict(
+            type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0),
+        loss_bbox=dict(type='L1Loss', loss_weight=1.0)),
+    roi_head=dict(
+        type='StandardRoIHead',
+        shared_head=dict(
+            type='ResLayer',
+            depth=50,
+            stage=3,
+            stride=2,
+            dilation=1,
+            style='caffe',
+            norm_cfg=norm_cfg,
+            norm_eval=True),
+        bbox_roi_extractor=dict(
+            type='SingleRoIExtractor',
+            roi_layer=dict(type='RoIAlign', output_size=14, sampling_ratio=0),
+            out_channels=1024,
+            featmap_strides=[16]),
+        bbox_head=dict(
+            type='BBoxHead',
+            with_avg_pool=True,
+            roi_feat_size=7,
+            in_channels=2048,
+            num_classes=80,
+            bbox_coder=dict(
+                type='DeltaXYWHBBoxCoder',
+                target_means=[0., 0., 0., 0.],
+                target_stds=[0.1, 0.1, 0.2, 0.2]),
+            reg_class_agnostic=False,
+            loss_cls=dict(
+                type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0),
+            loss_bbox=dict(type='L1Loss', loss_weight=1.0))),
+    # model training and testing settings
+    train_cfg=dict(
+        rpn=dict(
+            assigner=dict(
+                type='MaxIoUAssigner',
+                pos_iou_thr=0.7,
+                neg_iou_thr=0.3,
+                min_pos_iou=0.3,
+                match_low_quality=True,
+                ignore_iof_thr=-1),
+            sampler=dict(
+                type='RandomSampler',
+                num=256,
+                pos_fraction=0.5,
+                neg_pos_ub=-1,
+                add_gt_as_proposals=False),
+            allowed_border=0,
+            pos_weight=-1,
+            debug=False),
+        rpn_proposal=dict(
+            nms_pre=12000,
+            max_per_img=2000,
+            nms=dict(type='nms', iou_threshold=0.7),
+            min_bbox_size=0),
+        rcnn=dict(
+            assigner=dict(
+                type='MaxIoUAssigner',
+                pos_iou_thr=0.5,
+                neg_iou_thr=0.5,
+                min_pos_iou=0.5,
+                match_low_quality=False,
+                ignore_iof_thr=-1),
+            sampler=dict(
+                type='RandomSampler',
+                num=512,
+                pos_fraction=0.25,
+                neg_pos_ub=-1,
+                add_gt_as_proposals=True),
+            pos_weight=-1,
+            debug=False)),
+    test_cfg=dict(
+        rpn=dict(
+            nms_pre=6000,
+            max_per_img=1000,
+            nms=dict(type='nms', iou_threshold=0.7),
+            min_bbox_size=0),
+        rcnn=dict(
+            score_thr=0.05,
+            nms=dict(type='nms', iou_threshold=0.5),
+            max_per_img=100)))

detection/configs/_base_/models/faster_rcnn_r50_caffe_dc5.py

@@ -0,0 +1,105 @@
+# model settings
+norm_cfg = dict(type='BN', requires_grad=False)
+model = dict(
+    type='FasterRCNN',
+    backbone=dict(
+        type='ResNet',
+        depth=50,
+        num_stages=4,
+        strides=(1, 2, 2, 1),
+        dilations=(1, 1, 1, 2),
+        out_indices=(3, ),
+        frozen_stages=1,
+        norm_cfg=norm_cfg,
+        norm_eval=True,
+        style='caffe',
+        init_cfg=dict(
+            type='Pretrained',
+            checkpoint='open-mmlab://detectron2/resnet50_caffe')),
+    rpn_head=dict(
+        type='RPNHead',
+        in_channels=2048,
+        feat_channels=2048,
+        anchor_generator=dict(
+            type='AnchorGenerator',
+            scales=[2, 4, 8, 16, 32],
+            ratios=[0.5, 1.0, 2.0],
+            strides=[16]),
+        bbox_coder=dict(
+            type='DeltaXYWHBBoxCoder',
+            target_means=[.0, .0, .0, .0],
+            target_stds=[1.0, 1.0, 1.0, 1.0]),
+        loss_cls=dict(
+            type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0),
+        loss_bbox=dict(type='L1Loss', loss_weight=1.0)),
+    roi_head=dict(
+        type='StandardRoIHead',
+        bbox_roi_extractor=dict(
+            type='SingleRoIExtractor',
+            roi_layer=dict(type='RoIAlign', output_size=7, sampling_ratio=0),
+            out_channels=2048,
+            featmap_strides=[16]),
+        bbox_head=dict(
+            type='Shared2FCBBoxHead',
+            in_channels=2048,
+            fc_out_channels=1024,
+            roi_feat_size=7,
+            num_classes=80,
+            bbox_coder=dict(
+                type='DeltaXYWHBBoxCoder',
+                target_means=[0., 0., 0., 0.],
+                target_stds=[0.1, 0.1, 0.2, 0.2]),
+            reg_class_agnostic=False,
+            loss_cls=dict(
+                type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0),
+            loss_bbox=dict(type='L1Loss', loss_weight=1.0))),
+    # model training and testing settings
+    train_cfg=dict(
+        rpn=dict(
+            assigner=dict(
+                type='MaxIoUAssigner',
+                pos_iou_thr=0.7,
+                neg_iou_thr=0.3,
+                min_pos_iou=0.3,
+                match_low_quality=True,
+                ignore_iof_thr=-1),
+            sampler=dict(
+                type='RandomSampler',
+                num=256,
+                pos_fraction=0.5,
+                neg_pos_ub=-1,
+                add_gt_as_proposals=False),
+            allowed_border=0,
+            pos_weight=-1,
+            debug=False),
+        rpn_proposal=dict(
+            nms_pre=12000,
+            max_per_img=2000,
+            nms=dict(type='nms', iou_threshold=0.7),
+            min_bbox_size=0),
+        rcnn=dict(
+            assigner=dict(
+                type='MaxIoUAssigner',
+                pos_iou_thr=0.5,
+                neg_iou_thr=0.5,
+                min_pos_iou=0.5,
+                match_low_quality=False,
+                ignore_iof_thr=-1),
+            sampler=dict(
+                type='RandomSampler',
+                num=512,
+                pos_fraction=0.25,
+                neg_pos_ub=-1,
+                add_gt_as_proposals=True),
+            pos_weight=-1,
+            debug=False)),
+    test_cfg=dict(
+        rpn=dict(
+            nms=dict(type='nms', iou_threshold=0.7),
+            nms_pre=6000,
+            max_per_img=1000,
+            min_bbox_size=0),
+        rcnn=dict(
+            score_thr=0.05,
+            nms=dict(type='nms', iou_threshold=0.5),
+            max_per_img=100)))

detection/configs/_base_/models/faster_rcnn_r50_fpn.py

@@ -0,0 +1,108 @@
+# model settings
+model = dict(
+    type='FasterRCNN',
+    backbone=dict(
+        type='ResNet',
+        depth=50,
+        num_stages=4,
+        out_indices=(0, 1, 2, 3),
+        frozen_stages=1,
+        norm_cfg=dict(type='BN', requires_grad=True),
+        norm_eval=True,
+        style='pytorch',
+        init_cfg=dict(type='Pretrained', checkpoint='torchvision://resnet50')),
+    neck=dict(
+        type='FPN',
+        in_channels=[256, 512, 1024, 2048],
+        out_channels=256,
+        num_outs=5),
+    rpn_head=dict(
+        type='RPNHead',
+        in_channels=256,
+        feat_channels=256,
+        anchor_generator=dict(
+            type='AnchorGenerator',
+            scales=[8],
+            ratios=[0.5, 1.0, 2.0],
+            strides=[4, 8, 16, 32, 64]),
+        bbox_coder=dict(
+            type='DeltaXYWHBBoxCoder',
+            target_means=[.0, .0, .0, .0],
+            target_stds=[1.0, 1.0, 1.0, 1.0]),
+        loss_cls=dict(
+            type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0),
+        loss_bbox=dict(type='L1Loss', loss_weight=1.0)),
+    roi_head=dict(
+        type='StandardRoIHead',
+        bbox_roi_extractor=dict(
+            type='SingleRoIExtractor',
+            roi_layer=dict(type='RoIAlign', output_size=7, sampling_ratio=0),
+            out_channels=256,
+            featmap_strides=[4, 8, 16, 32]),
+        bbox_head=dict(
+            type='Shared2FCBBoxHead',
+            in_channels=256,
+            fc_out_channels=1024,
+            roi_feat_size=7,
+            num_classes=80,
+            bbox_coder=dict(
+                type='DeltaXYWHBBoxCoder',
+                target_means=[0., 0., 0., 0.],
+                target_stds=[0.1, 0.1, 0.2, 0.2]),
+            reg_class_agnostic=False,
+            loss_cls=dict(
+                type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0),
+            loss_bbox=dict(type='L1Loss', loss_weight=1.0))),
+    # model training and testing settings
+    train_cfg=dict(
+        rpn=dict(
+            assigner=dict(
+                type='MaxIoUAssigner',
+                pos_iou_thr=0.7,
+                neg_iou_thr=0.3,
+                min_pos_iou=0.3,
+                match_low_quality=True,
+                ignore_iof_thr=-1),
+            sampler=dict(
+                type='RandomSampler',
+                num=256,
+                pos_fraction=0.5,
+                neg_pos_ub=-1,
+                add_gt_as_proposals=False),
+            allowed_border=-1,
+            pos_weight=-1,
+            debug=False),
+        rpn_proposal=dict(
+            nms_pre=2000,
+            max_per_img=1000,
+            nms=dict(type='nms', iou_threshold=0.7),
+            min_bbox_size=0),
+        rcnn=dict(
+            assigner=dict(
+                type='MaxIoUAssigner',
+                pos_iou_thr=0.5,
+                neg_iou_thr=0.5,
+                min_pos_iou=0.5,
+                match_low_quality=False,
+                ignore_iof_thr=-1),
+            sampler=dict(
+                type='RandomSampler',
+                num=512,
+                pos_fraction=0.25,
+                neg_pos_ub=-1,
+                add_gt_as_proposals=True),
+            pos_weight=-1,
+            debug=False)),
+    test_cfg=dict(
+        rpn=dict(
+            nms_pre=1000,
+            max_per_img=1000,
+            nms=dict(type='nms', iou_threshold=0.7),
+            min_bbox_size=0),
+        rcnn=dict(
+            score_thr=0.05,
+            nms=dict(type='nms', iou_threshold=0.5),
+            max_per_img=100)
+        # soft-nms is also supported for rcnn testing
+        # e.g., nms=dict(type='soft_nms', iou_threshold=0.5, min_score=0.05)
+    ))

detection/configs/_base_/models/mask_rcnn_convnext_fpn.py

@@ -0,0 +1,128 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+# model settings
+model = dict(
+    type='MaskRCNN',
+    pretrained=None,
+    backbone=dict(
+        type='ConvNeXt',
+        in_chans=3,
+        depths=[3, 3, 9, 3],
+        dims=[96, 192, 384, 768],
+        drop_path_rate=0.2,
+        layer_scale_init_value=1e-6,
+        out_indices=[0, 1, 2, 3],
+    ),
+    neck=dict(
+        type='FPN',
+        in_channels=[128, 256, 512, 1024],
+        out_channels=256,
+        num_outs=5),
+    rpn_head=dict(
+        type='RPNHead',
+        in_channels=256,
+        feat_channels=256,
+        anchor_generator=dict(
+            type='AnchorGenerator',
+            scales=[8],
+            ratios=[0.5, 1.0, 2.0],
+            strides=[4, 8, 16, 32, 64]),
+        bbox_coder=dict(
+            type='DeltaXYWHBBoxCoder',
+            target_means=[.0, .0, .0, .0],
+            target_stds=[1.0, 1.0, 1.0, 1.0]),
+        loss_cls=dict(
+            type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0),
+        loss_bbox=dict(type='L1Loss', loss_weight=1.0)),
+    roi_head=dict(
+        type='StandardRoIHead',
+        bbox_roi_extractor=dict(
+            type='SingleRoIExtractor',
+            roi_layer=dict(type='RoIAlign', output_size=7, sampling_ratio=0),
+            out_channels=256,
+            featmap_strides=[4, 8, 16, 32]),
+        bbox_head=dict(
+            type='Shared2FCBBoxHead',
+            in_channels=256,
+            fc_out_channels=1024,
+            roi_feat_size=7,
+            num_classes=80,
+            bbox_coder=dict(
+                type='DeltaXYWHBBoxCoder',
+                target_means=[0., 0., 0., 0.],
+                target_stds=[0.1, 0.1, 0.2, 0.2]),
+            reg_class_agnostic=False,
+            loss_cls=dict(
+                type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0),
+            loss_bbox=dict(type='L1Loss', loss_weight=1.0)),
+        mask_roi_extractor=dict(
+            type='SingleRoIExtractor',
+            roi_layer=dict(type='RoIAlign', output_size=14, sampling_ratio=0),
+            out_channels=256,
+            featmap_strides=[4, 8, 16, 32]),
+        mask_head=dict(
+            type='FCNMaskHead',
+            num_convs=4,
+            in_channels=256,
+            conv_out_channels=256,
+            num_classes=80,
+            loss_mask=dict(
+                type='CrossEntropyLoss', use_mask=True, loss_weight=1.0))),
+    # model training and testing settings
+    train_cfg=dict(
+        rpn=dict(
+            assigner=dict(
+                type='MaxIoUAssigner',
+                pos_iou_thr=0.7,
+                neg_iou_thr=0.3,
+                min_pos_iou=0.3,
+                match_low_quality=True,
+                ignore_iof_thr=-1),
+            sampler=dict(
+                type='RandomSampler',
+                num=256,
+                pos_fraction=0.5,
+                neg_pos_ub=-1,
+                add_gt_as_proposals=False),
+            allowed_border=-1,
+            pos_weight=-1,
+            debug=False),
+        rpn_proposal=dict(
+            nms_pre=2000,
+            max_per_img=1000,
+            nms=dict(type='nms', iou_threshold=0.7),
+            min_bbox_size=0),
+        rcnn=dict(
+            assigner=dict(
+                type='MaxIoUAssigner',
+                pos_iou_thr=0.5,
+                neg_iou_thr=0.5,
+                min_pos_iou=0.5,
+                match_low_quality=True,
+                ignore_iof_thr=-1),
+            sampler=dict(
+                type='RandomSampler',
+                num=512,
+                pos_fraction=0.25,
+                neg_pos_ub=-1,
+                add_gt_as_proposals=True),
+            mask_size=28,
+            pos_weight=-1,
+            debug=False)),
+    test_cfg=dict(
+        rpn=dict(
+            nms_pre=1000,
+            max_per_img=1000,
+            nms=dict(type='nms', iou_threshold=0.7),
+            min_bbox_size=0),
+        rcnn=dict(
+            score_thr=0.05,
+            nms=dict(type='nms', iou_threshold=0.5),
+            max_per_img=100,
+            mask_thr_binary=0.5)))

detection/configs/_base_/models/mask_rcnn_r50_caffe_c4.py

@@ -0,0 +1,125 @@
+# model settings
+norm_cfg = dict(type='BN', requires_grad=False)
+model = dict(
+    type='MaskRCNN',
+    backbone=dict(
+        type='ResNet',
+        depth=50,
+        num_stages=3,
+        strides=(1, 2, 2),
+        dilations=(1, 1, 1),
+        out_indices=(2, ),
+        frozen_stages=1,
+        norm_cfg=norm_cfg,
+        norm_eval=True,
+        style='caffe',
+        init_cfg=dict(
+            type='Pretrained',
+            checkpoint='open-mmlab://detectron2/resnet50_caffe')),
+    rpn_head=dict(
+        type='RPNHead',
+        in_channels=1024,
+        feat_channels=1024,
+        anchor_generator=dict(
+            type='AnchorGenerator',
+            scales=[2, 4, 8, 16, 32],
+            ratios=[0.5, 1.0, 2.0],
+            strides=[16]),
+        bbox_coder=dict(
+            type='DeltaXYWHBBoxCoder',
+            target_means=[.0, .0, .0, .0],
+            target_stds=[1.0, 1.0, 1.0, 1.0]),
+        loss_cls=dict(
+            type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0),
+        loss_bbox=dict(type='L1Loss', loss_weight=1.0)),
+    roi_head=dict(
+        type='StandardRoIHead',
+        shared_head=dict(
+            type='ResLayer',
+            depth=50,
+            stage=3,
+            stride=2,
+            dilation=1,
+            style='caffe',
+            norm_cfg=norm_cfg,
+            norm_eval=True),
+        bbox_roi_extractor=dict(
+            type='SingleRoIExtractor',
+            roi_layer=dict(type='RoIAlign', output_size=14, sampling_ratio=0),
+            out_channels=1024,
+            featmap_strides=[16]),
+        bbox_head=dict(
+            type='BBoxHead',
+            with_avg_pool=True,
+            roi_feat_size=7,
+            in_channels=2048,
+            num_classes=80,
+            bbox_coder=dict(
+                type='DeltaXYWHBBoxCoder',
+                target_means=[0., 0., 0., 0.],
+                target_stds=[0.1, 0.1, 0.2, 0.2]),
+            reg_class_agnostic=False,
+            loss_cls=dict(
+                type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0),
+            loss_bbox=dict(type='L1Loss', loss_weight=1.0)),
+        mask_roi_extractor=None,
+        mask_head=dict(
+            type='FCNMaskHead',
+            num_convs=0,
+            in_channels=2048,
+            conv_out_channels=256,
+            num_classes=80,
+            loss_mask=dict(
+                type='CrossEntropyLoss', use_mask=True, loss_weight=1.0))),
+    # model training and testing settings
+    train_cfg=dict(
+        rpn=dict(
+            assigner=dict(
+                type='MaxIoUAssigner',
+                pos_iou_thr=0.7,
+                neg_iou_thr=0.3,
+                min_pos_iou=0.3,
+                match_low_quality=True,
+                ignore_iof_thr=-1),
+            sampler=dict(
+                type='RandomSampler',
+                num=256,
+                pos_fraction=0.5,
+                neg_pos_ub=-1,
+                add_gt_as_proposals=False),
+            allowed_border=0,
+            pos_weight=-1,
+            debug=False),
+        rpn_proposal=dict(
+            nms_pre=12000,
+            max_per_img=2000,
+            nms=dict(type='nms', iou_threshold=0.7),
+            min_bbox_size=0),
+        rcnn=dict(
+            assigner=dict(
+                type='MaxIoUAssigner',
+                pos_iou_thr=0.5,
+                neg_iou_thr=0.5,
+                min_pos_iou=0.5,
+                match_low_quality=False,
+                ignore_iof_thr=-1),
+            sampler=dict(
+                type='RandomSampler',
+                num=512,
+                pos_fraction=0.25,
+                neg_pos_ub=-1,
+                add_gt_as_proposals=True),
+            mask_size=14,
+            pos_weight=-1,
+            debug=False)),
+    test_cfg=dict(
+        rpn=dict(
+            nms_pre=6000,
+            nms=dict(type='nms', iou_threshold=0.7),
+            max_per_img=1000,
+            min_bbox_size=0),
+        rcnn=dict(
+            score_thr=0.05,
+            nms=dict(type='nms', iou_threshold=0.5),
+            max_per_img=100,
+            mask_thr_binary=0.5)))

detection/configs/_base_/models/mask_rcnn_r50_fpn.py

@@ -0,0 +1,120 @@
+# model settings
+model = dict(
+    type='MaskRCNN',
+    backbone=dict(
+        type='ResNet',
+        depth=50,
+        num_stages=4,
+        out_indices=(0, 1, 2, 3),
+        frozen_stages=1,
+        norm_cfg=dict(type='BN', requires_grad=True),
+        norm_eval=True,
+        style='pytorch',
+        init_cfg=dict(type='Pretrained', checkpoint='torchvision://resnet50')),
+    neck=dict(
+        type='FPN',
+        in_channels=[256, 512, 1024, 2048],
+        out_channels=256,
+        num_outs=5),
+    rpn_head=dict(
+        type='RPNHead',
+        in_channels=256,
+        feat_channels=256,
+        anchor_generator=dict(
+            type='AnchorGenerator',
+            scales=[8],
+            ratios=[0.5, 1.0, 2.0],
+            strides=[4, 8, 16, 32, 64]),
+        bbox_coder=dict(
+            type='DeltaXYWHBBoxCoder',
+            target_means=[.0, .0, .0, .0],
+            target_stds=[1.0, 1.0, 1.0, 1.0]),
+        loss_cls=dict(
+            type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0),
+        loss_bbox=dict(type='L1Loss', loss_weight=1.0)),
+    roi_head=dict(
+        type='StandardRoIHead',
+        bbox_roi_extractor=dict(
+            type='SingleRoIExtractor',
+            roi_layer=dict(type='RoIAlign', output_size=7, sampling_ratio=0),
+            out_channels=256,
+            featmap_strides=[4, 8, 16, 32]),
+        bbox_head=dict(
+            type='Shared2FCBBoxHead',
+            in_channels=256,
+            fc_out_channels=1024,
+            roi_feat_size=7,
+            num_classes=80,
+            bbox_coder=dict(
+                type='DeltaXYWHBBoxCoder',
+                target_means=[0., 0., 0., 0.],
+                target_stds=[0.1, 0.1, 0.2, 0.2]),
+            reg_class_agnostic=False,
+            loss_cls=dict(
+                type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0),
+            loss_bbox=dict(type='L1Loss', loss_weight=1.0)),
+        mask_roi_extractor=dict(
+            type='SingleRoIExtractor',
+            roi_layer=dict(type='RoIAlign', output_size=14, sampling_ratio=0),
+            out_channels=256,
+            featmap_strides=[4, 8, 16, 32]),
+        mask_head=dict(
+            type='FCNMaskHead',
+            num_convs=4,
+            in_channels=256,
+            conv_out_channels=256,
+            num_classes=80,
+            loss_mask=dict(
+                type='CrossEntropyLoss', use_mask=True, loss_weight=1.0))),
+    # model training and testing settings
+    train_cfg=dict(
+        rpn=dict(
+            assigner=dict(
+                type='MaxIoUAssigner',
+                pos_iou_thr=0.7,
+                neg_iou_thr=0.3,
+                min_pos_iou=0.3,
+                match_low_quality=True,
+                ignore_iof_thr=-1),
+            sampler=dict(
+                type='RandomSampler',
+                num=256,
+                pos_fraction=0.5,
+                neg_pos_ub=-1,
+                add_gt_as_proposals=False),
+            allowed_border=-1,
+            pos_weight=-1,
+            debug=False),
+        rpn_proposal=dict(
+            nms_pre=2000,
+            max_per_img=1000,
+            nms=dict(type='nms', iou_threshold=0.7),
+            min_bbox_size=0),
+        rcnn=dict(
+            assigner=dict(
+                type='MaxIoUAssigner',
+                pos_iou_thr=0.5,
+                neg_iou_thr=0.5,
+                min_pos_iou=0.5,
+                match_low_quality=True,
+                ignore_iof_thr=-1),
+            sampler=dict(
+                type='RandomSampler',
+                num=512,
+                pos_fraction=0.25,
+                neg_pos_ub=-1,
+                add_gt_as_proposals=True),
+            mask_size=28,
+            pos_weight=-1,
+            debug=False)),
+    test_cfg=dict(
+        rpn=dict(
+            nms_pre=1000,
+            max_per_img=1000,
+            nms=dict(type='nms', iou_threshold=0.7),
+            min_bbox_size=0),
+        rcnn=dict(
+            score_thr=0.05,
+            nms=dict(type='nms', iou_threshold=0.5),
+            max_per_img=100,
+            mask_thr_binary=0.5)))

detection/configs/_base_/models/retinanet_r50_fpn.py

@@ -0,0 +1,60 @@
+# model settings
+model = dict(
+    type='RetinaNet',
+    backbone=dict(
+        type='ResNet',
+        depth=50,
+        num_stages=4,
+        out_indices=(0, 1, 2, 3),
+        frozen_stages=1,
+        norm_cfg=dict(type='BN', requires_grad=True),
+        norm_eval=True,
+        style='pytorch',
+        init_cfg=dict(type='Pretrained', checkpoint='torchvision://resnet50')),
+    neck=dict(
+        type='FPN',
+        in_channels=[256, 512, 1024, 2048],
+        out_channels=256,
+        start_level=1,
+        add_extra_convs='on_input',
+        num_outs=5),
+    bbox_head=dict(
+        type='RetinaHead',
+        num_classes=80,
+        in_channels=256,
+        stacked_convs=4,
+        feat_channels=256,
+        anchor_generator=dict(
+            type='AnchorGenerator',
+            octave_base_scale=4,
+            scales_per_octave=3,
+            ratios=[0.5, 1.0, 2.0],
+            strides=[8, 16, 32, 64, 128]),
+        bbox_coder=dict(
+            type='DeltaXYWHBBoxCoder',
+            target_means=[.0, .0, .0, .0],
+            target_stds=[1.0, 1.0, 1.0, 1.0]),
+        loss_cls=dict(
+            type='FocalLoss',
+            use_sigmoid=True,
+            gamma=2.0,
+            alpha=0.25,
+            loss_weight=1.0),
+        loss_bbox=dict(type='L1Loss', loss_weight=1.0)),
+    # model training and testing settings
+    train_cfg=dict(
+        assigner=dict(
+            type='MaxIoUAssigner',
+            pos_iou_thr=0.5,
+            neg_iou_thr=0.4,
+            min_pos_iou=0,
+            ignore_iof_thr=-1),
+        allowed_border=-1,
+        pos_weight=-1,
+        debug=False),
+    test_cfg=dict(
+        nms_pre=1000,
+        min_bbox_size=0,
+        score_thr=0.05,
+        nms=dict(type='nms', iou_threshold=0.5),
+        max_per_img=100))

detection/configs/_base_/models/rpn_r50_caffe_c4.py

@@ -0,0 +1,58 @@
+# model settings
+model = dict(
+    type='RPN',
+    backbone=dict(
+        type='ResNet',
+        depth=50,
+        num_stages=3,
+        strides=(1, 2, 2),
+        dilations=(1, 1, 1),
+        out_indices=(2, ),
+        frozen_stages=1,
+        norm_cfg=dict(type='BN', requires_grad=False),
+        norm_eval=True,
+        style='caffe',
+        init_cfg=dict(
+            type='Pretrained',
+            checkpoint='open-mmlab://detectron2/resnet50_caffe')),
+    neck=None,
+    rpn_head=dict(
+        type='RPNHead',
+        in_channels=1024,
+        feat_channels=1024,
+        anchor_generator=dict(
+            type='AnchorGenerator',
+            scales=[2, 4, 8, 16, 32],
+            ratios=[0.5, 1.0, 2.0],
+            strides=[16]),
+        bbox_coder=dict(
+            type='DeltaXYWHBBoxCoder',
+            target_means=[.0, .0, .0, .0],
+            target_stds=[1.0, 1.0, 1.0, 1.0]),
+        loss_cls=dict(
+            type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0),
+        loss_bbox=dict(type='L1Loss', loss_weight=1.0)),
+    # model training and testing settings
+    train_cfg=dict(
+        rpn=dict(
+            assigner=dict(
+                type='MaxIoUAssigner',
+                pos_iou_thr=0.7,
+                neg_iou_thr=0.3,
+                min_pos_iou=0.3,
+                ignore_iof_thr=-1),
+            sampler=dict(
+                type='RandomSampler',
+                num=256,
+                pos_fraction=0.5,
+                neg_pos_ub=-1,
+                add_gt_as_proposals=False),
+            allowed_border=0,
+            pos_weight=-1,
+            debug=False)),
+    test_cfg=dict(
+        rpn=dict(
+            nms_pre=12000,
+            max_per_img=2000,
+            nms=dict(type='nms', iou_threshold=0.7),
+            min_bbox_size=0)))