From 6b3e1d6fe736ea59f19f5fcd774ca6b03787a61f Mon Sep 17 00:00:00 2001 From: Yuwen Xiong Date: Fri, 19 Jan 2024 22:18:59 +0800 Subject: [PATCH] update 1.0.0.post2 --- .gitignore | 1 + DCNv4_op/DCNv4/__init__.py | 2 ++ DCNv4_op/DCNv4/functions/flash_deform_attn_func.py | 4 ++-- DCNv4_op/DCNv4/modules/dcnv4.py | 4 +++- DCNv4_op/DCNv4/modules/flash_deform_attn.py | 1 + DCNv4_op/__init__.py | 2 -- DCNv4_op/setup.py | 2 +- DCNv4_op/src/cuda/dcnv4_cuda.cu | 13 +++++++------ DCNv4_op/src/cuda/flash_deform_attn_cuda.cu | 7 ++++--- 9 files changed, 21 insertions(+), 15 deletions(-) create mode 100644 DCNv4_op/DCNv4/__init__.py diff --git a/.gitignore b/.gitignore index 9c57ef3..0ad7194 100644 --- a/.gitignore +++ b/.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/ diff --git a/DCNv4_op/DCNv4/__init__.py b/DCNv4_op/DCNv4/__init__.py new file mode 100644 index 0000000..ab94d0f --- /dev/null +++ b/DCNv4_op/DCNv4/__init__.py @@ -0,0 +1,2 @@ +from .functions import DCNv4Function, FlashDeformAttnFunction +from .modules import DCNv4, FlashDeformAttn \ No newline at end of file diff --git a/DCNv4_op/DCNv4/functions/flash_deform_attn_func.py b/DCNv4_op/DCNv4/functions/flash_deform_attn_func.py index 08c44cc..974a121 100644 --- a/DCNv4_op/DCNv4/functions/flash_deform_attn_func.py +++ b/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 ): diff --git a/DCNv4_op/DCNv4/modules/dcnv4.py b/DCNv4_op/DCNv4/modules/dcnv4.py index 560e9d7..92d5305 100644 --- a/DCNv4_op/DCNv4/modules/dcnv4.py +++ b/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 diff --git a/DCNv4_op/DCNv4/modules/flash_deform_attn.py b/DCNv4_op/DCNv4/modules/flash_deform_attn.py index 4472fa6..0c4922b 100644 --- a/DCNv4_op/DCNv4/modules/flash_deform_attn.py +++ b/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 diff --git a/DCNv4_op/__init__.py b/DCNv4_op/__init__.py index ab94d0f..e69de29 100644 --- a/DCNv4_op/__init__.py +++ b/DCNv4_op/__init__.py @@ -1,2 +0,0 @@ -from .functions import DCNv4Function, FlashDeformAttnFunction -from .modules import DCNv4, FlashDeformAttn \ No newline at end of file diff --git a/DCNv4_op/setup.py b/DCNv4_op/setup.py index 1fa8531..082dbeb 100644 --- a/DCNv4_op/setup.py +++ b/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", diff --git a/DCNv4_op/src/cuda/dcnv4_cuda.cu b/DCNv4_op/src/cuda/dcnv4_cuda.cu index fd158d1..3e0224c 100644 --- a/DCNv4_op/src/cuda/dcnv4_cuda.cu +++ b/DCNv4_op/src/cuda/dcnv4_cuda.cu @@ -91,17 +91,17 @@ at::Tensor dcnv4_cuda_forward( std::vector 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(), diff --git a/DCNv4_op/src/cuda/flash_deform_attn_cuda.cu b/DCNv4_op/src/cuda/flash_deform_attn_cuda.cu index 3038ea5..7496a5f 100644 --- a/DCNv4_op/src/cuda/flash_deform_attn_cuda.cu +++ b/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() + 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() + n * im2col_step_ * per_value_size, grad_offset.data() + n * im2col_step_ * per_offset_size, d_stride, block_thread