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

Remove a redundant .softmax() in FlashDeformAttn.

Browse Source
This commit is contained in:
HELLORPG
2024-05-17 11:12:15 +08:00
parent b2b6028d3c
commit 398607aeb2
1 changed files with 0 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -111,7 +111,6 @@ class FlashDeformAttn(nn.Module):
value = value.view(N, Len_in, self.n_heads, self.d_model // self.n_heads) 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) 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 = 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 # N, Len_q, n_heads, n_levels, n_points, 2
if reference_points.shape[-1] == 2: if reference_points.shape[-1] == 2:
offset_normalizer = torch.stack([input_spatial_shapes[..., 1], input_spatial_shapes[..., 0]], -1) 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: # Cat sampling_offsets and attention_weights, generate sampling_loc_attn:
sampling_locations = sampling_locations.flatten(-3).half() sampling_locations = sampling_locations.flatten(-3).half()
attention_weights = attention_weights.flatten(-2)
sampling_loc_attn = torch.cat([sampling_locations, attention_weights], dim=-1) sampling_loc_attn = torch.cat([sampling_locations, attention_weights], dim=-1)
output = FlashDeformAttnFunction.apply( output = FlashDeformAttnFunction.apply(