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

update 1.0.0.post2

Browse Source
This commit is contained in:
Yuwen Xiong
2024-01-19 22:18:59 +08:00
parent 7d59305b5f
commit 6b3e1d6fe7
9 changed files with 21 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

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

2
DCNv4_op/DCNv4/__init__.py Normal file
View File

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

4
DCNv4_op/DCNv4/modules/dcnv4.py
View File

@@ -139,13 +139,15 @@ class DCNv4(nn.Module):
self.remove_center self.remove_center
) )
x = x.view(N, L, -1)
if self.center_feature_scale: if self.center_feature_scale:
center_feature_scale = self.center_feature_scale_module( center_feature_scale = self.center_feature_scale_module(
x, self.center_feature_scale_proj_weight, self.center_feature_scale_proj_bias) x, self.center_feature_scale_proj_weight, self.center_feature_scale_proj_bias)
center_feature_scale = center_feature_scale[..., None].repeat( center_feature_scale = center_feature_scale[..., None].repeat(
1, 1, 1, 1, self.channels // self.group).flatten(-2) 1, 1, 1, 1, self.channels // self.group).flatten(-2)
x = x * (1 - center_feature_scale) + x_proj * center_feature_scale x = x * (1 - center_feature_scale) + x_proj * center_feature_scale
x = x.view(N, L, -1)
if not self.without_pointwise: if not self.without_pointwise:
x = self.output_proj(x) x = self.output_proj(x)
return x return x

1
DCNv4_op/DCNv4/modules/flash_deform_attn.py
View File

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

2
DCNv4_op/__init__.py
View File

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

2
DCNv4_op/setup.py
View File

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

3
DCNv4_op/src/cuda/dcnv4_cuda.cu
View File

@@ -146,7 +146,8 @@ dcnv4_cuda_backward(
for (int n = 0; n < batch / im2col_step_; ++n) { for (int n = 0; n < batch / im2col_step_; ++n) {
auto columns = grad_output_n.select(0, 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_backward_cuda", ([&] {
dcnv4_col2im_cuda( dcnv4_col2im_cuda(
at::cuda::getCurrentCUDAStream(), at::cuda::getCurrentCUDAStream(),

3
DCNv4_op/src/cuda/flash_deform_attn_cuda.cu
View File

@@ -135,7 +135,8 @@ flash_deform_attn_cuda_backward(
auto per_out_size = num_query * num_heads * num_channels; auto per_out_size = num_query * num_heads * num_channels;
for (int n = 0; n < batch / im2col_step_; ++n) { 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_deform_attn_backward_cuda", ([&] {
flash_deformable_col2im_cuda( flash_deformable_col2im_cuda(
at::cuda::getCurrentCUDAStream(), at::cuda::getCurrentCUDAStream(),