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/dist
ops_dcnv3/DCNv3.egg-info
DCNv4_op/DCNv4.egg-info
build/
dist/
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
)
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

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

@@ -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

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(
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",

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) {
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(),

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;
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(),