Update export.py with v3.0 Hardswish() support

This commit is contained in:
Glenn Jocher committed 2020-08-24 21:47:49 -07:00
1 parent 8666bc507e
commit 4d7f222f73
2 files changed
+15 -10

No files matched your search

+3 -3
View File
@@ -10,11 +10,11 @@ class Swish(nn.Module): #
return x * torch.sigmoid(x)
class Hardswish(nn.Module): # alternative to nn.Hardswish() for export
class Hardswish(nn.Module): # export-friendly version of nn.Hardswish()
@staticmethod
def forward(x):
# return x * F.hardsigmoid(x)
return x * F.hardtanh(x + 3, 0., 6.) / 6.
# return x * F.hardsigmoid(x) # for torchscript and CoreML
return x * F.hardtanh(x + 3, 0., 6.) / 6. # for torchscript, CoreML and ONNX
class MemoryEfficientSwish(nn.Module):