diff --git a/models/export.py b/models/export.py index c5e96f1..057658a 100644 --- a/models/export.py +++ b/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 diff --git a/utils/activations.py b/utils/activations.py index 24f5a30..954d2e1 100644 --- a/utils/activations.py +++ b/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)