update 1.0.0.post2

.gitignore

@@ -11,6 +11,7 @@ detection/work_dirs
 ops_dcnv3/build
 ops_dcnv3/dist
 ops_dcnv3/DCNv3.egg-info
+DCNv4_op/DCNv4.egg-info
 build/
 dist/
 ckpts/

DCNv4_op/DCNv4/__init__.py

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

DCNv4_op/DCNv4/functions/flash_deform_attn_func.py

@@ -57,7 +57,7 @@ def findspec_bwd(B, Q, G, C):
         d_stride = 2
     else:
         d_stride = 1
-    
+
     ms = factors(B*Q)
     multiplier = 1
     for m in ms:
@@ -70,7 +70,7 @@ class FlashDeformAttnFunction(Function):
     @staticmethod
     @torch.autocast("cuda", enabled=True, dtype=torch.float16)
     def forward(
-        ctx, value, value_spatial_shapes, value_level_start_index, 
+        ctx, value, value_spatial_shapes, value_level_start_index,
         sampling_loc_attn, im2col_step, K=8
     ):
 

DCNv4_op/DCNv4/modules/dcnv4.py

@@ -139,13 +139,15 @@ class DCNv4(nn.Module):
             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
+
+        x = x.view(N, L, -1)
+
         if not self.without_pointwise:
             x = self.output_proj(x)
         return x

DCNv4_op/DCNv4/modules/flash_deform_attn.py

@@ -136,6 +136,7 @@ class FlashDeformAttn(nn.Module):
             sampling_locations,
             attention_weights,
             self.im2col_step,
+            self.n_points           
         )
         output = self.output_proj(output)
         return output

DCNv4_op/__init__.py

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

DCNv4_op/setup.py

@@ -62,7 +62,7 @@ def get_extensions():
 
 setup(
     name="DCNv4",
-    version="1.0.0",
+    version="1.0.0.post2",
     author="Yuwen Xiong, Feng Wang",
     url="",
     description="PyTorch Wrapper for CUDA Functions of DCNv4",

DCNv4_op/src/cuda/dcnv4_cuda.cu

@@ -91,17 +91,17 @@ at::Tensor dcnv4_cuda_forward(
 
 std::vector<at::Tensor>
 dcnv4_cuda_backward(
-    const at::Tensor &value, 
-    const at::Tensor &p_offset, 
+    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 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(), 
+  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");
@@ -129,7 +129,7 @@ dcnv4_cuda_backward(
       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;
@@ -146,7 +146,8 @@ dcnv4_cuda_backward(
 
   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(),
+    AT_DISPATCH_FLOATING_TYPES_AND2(
+        at::ScalarType::Half, at::ScalarType::BFloat16, value.scalar_type(),
         "dcnv4_backward_cuda", ([&] {
           dcnv4_col2im_cuda(
               at::cuda::getCurrentCUDAStream(),

DCNv4_op/src/cuda/flash_deform_attn_cuda.cu

@@ -26,7 +26,7 @@
 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 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(),
@@ -135,7 +135,8 @@ flash_deform_attn_cuda_backward(
   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(),
+    AT_DISPATCH_FLOATING_TYPES_AND2(
+        at::ScalarType::Half, at::ScalarType::BFloat16, value.scalar_type(),
         "flash_deform_attn_backward_cuda", ([&] {
           flash_deformable_col2im_cuda(
               at::cuda::getCurrentCUDAStream(),
@@ -145,7 +146,7 @@ flash_deform_attn_cuda_backward(
                   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, 
+              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