mirrors/yolov5
1
0
Fork 0
mirror of https://github.com/ultralytics/yolov5.git synced 2025-06-03 14:49:29 +08:00
Code Issues Projects Releases Wiki Activity

Replace 2 transpose() with 1 permute in TransformerBlock()` (#5645)

Browse Source
This commit is contained in:
Ding Yiwei 2021-11-15 17:06:18 +08:00 committed by GitHub
parent 540ef0dd30
commit c2523be634
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
models/common.py
View File

@ -86,8 +86,8 @@ class TransformerBlock(nn.Module):
if self.conv is not None: if self.conv is not None:
x = self.conv(x) x = self.conv(x)
b, _, w, h = x.shape b, _, w, h = x.shape
p = x.flatten(2).unsqueeze(0).transpose(0, 3).squeeze(3) p = x.flatten(2).permute(2, 0, 1)
return self.tr(p + self.linear(p)).unsqueeze(3).transpose(0, 3).reshape(b, self.c2, w, h) return self.tr(p + self.linear(p)).permute(1, 2, 0).reshape(b, self.c2, w, h)
class Bottleneck(nn.Module): class Bottleneck(nn.Module):