Add DWConvClass() (#4274)

* Add `DWConvClass()`

* Cleanup

* Cleanup2
This commit is contained in:
Glenn Jocher authored and GitHub committed 2021-08-02 15:36:30 +02:00
1 parent 9c6732f61c
commit 587c4b4b81
3 files changed
+12 -5

No files matched your search

+9 -2
View File
@@ -30,7 +30,7 @@ def autopad(k, p=None): # kernel, padding
def DWConv(c1, c2, k=1, s=1, act=True):
# Depth-wise convolution
# Depth-wise convolution function
return Conv(c1, c2, k, s, g=math.gcd(c1, c2), act=act)
@@ -45,10 +45,17 @@ class Conv(nn.Module):
def forward(self, x):
return self.act(self.bn(self.conv(x)))
def fuseforward(self, x):
def forward_fuse(self, x):
return self.act(self.conv(x))
class DWConvClass(Conv):
# Depth-wise convolution class
def __init__(self, c1, c2, k=1, s=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
super().__init__(c1, c2, k, s, act)
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k), groups=math.gcd(c1, c2), bias=False)
class TransformerLayer(nn.Module):
# Transformer layer https://arxiv.org/abs/2010.11929 (LayerNorm layers removed for better performance)
def __init__(self, c, num_heads):