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

Fix the issue of mismatch of params for FlashDeformAttnFunction.

Browse Source
This commit is contained in:
HELLORPG
2024-05-17 11:09:29 +08:00
parent 97dcd914fc
commit b2b6028d3c
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -129,12 +129,16 @@ class FlashDeformAttn(nn.Module):
"Last dim of reference_points must be 2 or 4, but get {} instead.".format(reference_points.shape[-1]) "Last dim of reference_points must be 2 or 4, but get {} instead.".format(reference_points.shape[-1])
) )
# Cat sampling_offsets and attention_weights, generate sampling_loc_attn:
sampling_locations = sampling_locations.flatten(-3).half()
attention_weights = attention_weights.flatten(-2)
sampling_loc_attn = torch.cat([sampling_locations, attention_weights], dim=-1)
output = FlashDeformAttnFunction.apply( output = FlashDeformAttnFunction.apply(
value, value,
input_spatial_shapes, input_spatial_shapes,
input_level_start_index, input_level_start_index,
sampling_locations, sampling_loc_attn,
attention_weights,
self.im2col_step, self.im2col_step,
self.n_points self.n_points
) )