From b2b6028d3c790d3221bca51a106b3b85f2a876a1 Mon Sep 17 00:00:00 2001 From: HELLORPG Date: Fri, 17 May 2024 11:09:29 +0800 Subject: [PATCH 1/2] Fix the issue of mismatch of params for FlashDeformAttnFunction. --- DCNv4_op/DCNv4/modules/flash_deform_attn.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/DCNv4_op/DCNv4/modules/flash_deform_attn.py b/DCNv4_op/DCNv4/modules/flash_deform_attn.py index 0c4922b..d10f72d 100644 --- a/DCNv4_op/DCNv4/modules/flash_deform_attn.py +++ b/DCNv4_op/DCNv4/modules/flash_deform_attn.py @@ -128,13 +128,17 @@ class FlashDeformAttn(nn.Module): raise ValueError( "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( value, input_spatial_shapes, input_level_start_index, - sampling_locations, - attention_weights, + sampling_loc_attn, self.im2col_step, self.n_points ) From 398607aeb2c8eecf76a82485a3d96f0d06892577 Mon Sep 17 00:00:00 2001 From: HELLORPG Date: Fri, 17 May 2024 11:12:15 +0800 Subject: [PATCH 2/2] Remove a redundant .softmax() in FlashDeformAttn. --- DCNv4_op/DCNv4/modules/flash_deform_attn.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/DCNv4_op/DCNv4/modules/flash_deform_attn.py b/DCNv4_op/DCNv4/modules/flash_deform_attn.py index d10f72d..d7340de 100644 --- a/DCNv4_op/DCNv4/modules/flash_deform_attn.py +++ b/DCNv4_op/DCNv4/modules/flash_deform_attn.py @@ -111,7 +111,6 @@ class FlashDeformAttn(nn.Module): value = value.view(N, Len_in, self.n_heads, self.d_model // self.n_heads) sampling_offsets = self.sampling_offsets(query).view(N, Len_q, self.n_heads, self.n_levels, self.n_points, 2) attention_weights = self.attention_weights(query).view(N, Len_q, self.n_heads, self.n_levels * self.n_points) - attention_weights = F.softmax(attention_weights, -1).view(N, Len_q, self.n_heads, self.n_levels, self.n_points) # N, Len_q, n_heads, n_levels, n_points, 2 if reference_points.shape[-1] == 2: offset_normalizer = torch.stack([input_spatial_shapes[..., 1], input_spatial_shapes[..., 0]], -1) @@ -131,7 +130,6 @@ class FlashDeformAttn(nn.Module): # 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(