zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Concise `TransformerBlock()` (#3821)

This commit is contained in:
Glenn Jocher 2021-06-29 16:03:10 +02:00 committed by GitHub
parent 5ea771d93d
commit 57c5d02bbe
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
models/common.py
View File

@ -77,18 +77,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) p = x.flatten(2).unsqueeze(0).transpose(0, 3).squeeze(3)
p = p.unsqueeze(0) return self.tr(p + self.linear(p)).unsqueeze(3).transpose(0, 3).reshape(b, self.c2, w, h)
p = p.transpose(0, 3)
p = p.squeeze(3)
e = self.linear(p)
x = p + e
x = self.tr(x)
x = x.unsqueeze(3)
x = x.transpose(0, 3)
x = x.reshape(b, self.c2, w, h)
return x
class Bottleneck(nn.Module): class Bottleneck(nn.Module):