nn.SiLU() export support (#1713)

models/export.py

@@ -15,7 +15,7 @@ import torch.nn as nn
 
 import models
 from models.experimental import attempt_load
-from utils.activations import Hardswish
+from utils.activations import Hardswish, SiLU
 from utils.general import set_logging, check_img_size
 
 if __name__ == '__main__':
@@ -43,9 +43,12 @@ if __name__ == '__main__':
     # Update model
     for k, m in model.named_modules():
         m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatibility
-        if isinstance(m, models.common.Conv) and isinstance(m.act, nn.Hardswish):
-            m.act = Hardswish()  # assign activation
-        # if isinstance(m, models.yolo.Detect):
+        if isinstance(m, models.common.Conv):  # assign export-friendly activations
+            if isinstance(m.act, nn.Hardswish):
+                m.act = Hardswish()
+            elif isinstance(m.act, nn.SiLU):
+                m.act = SiLU()
+        # elif isinstance(m, models.yolo.Detect):
         #     m.forward = m.forward_export  # assign forward (optional)
     model.model[-1].export = True  # set Detect() layer export=True
     y = model(img)  # dry run

utils/activations.py

@@ -5,8 +5,8 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 
-# Swish https://arxiv.org/pdf/1905.02244.pdf ---------------------------------------------------------------------------
-class Swish(nn.Module):  #
+# SiLU https://arxiv.org/pdf/1905.02244.pdf ----------------------------------------------------------------------------
+class SiLU(nn.Module):  # export-friendly version of nn.SiLU()
     @staticmethod
     def forward(x):
         return x * torch.sigmoid(x)