mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Fix bottleneck attn transpose typo, hopefully these train better now..

Browse Source
This commit is contained in:
Ross Wightman 2021-09-28 16:38:41 -07:00
parent 80075b0b8a
commit b81e79aae9
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
timm/models/layers/bottleneck_attn.py
View File

@ -122,7 +122,7 @@ class BottleneckAttn(nn.Module):
attn_logits = attn_logits + self.pos_embed(q) # B, num_heads, H * W, H * W attn_logits = attn_logits + self.pos_embed(q) # B, num_heads, H * W, H * W
attn_out = attn_logits.softmax(dim=-1) attn_out = attn_logits.softmax(dim=-1)
attn_out = (attn_out @ v).transpose(1, 2).reshape(B, self.dim_out, H, W) # B, dim_out, H, W attn_out = (attn_out @ v).transpose(-1, -2).reshape(B, self.dim_out, H, W) # B, dim_out, H, W
attn_out = self.pool(attn_out) attn_out = self.pool(attn_out)
return attn_out return attn_out