diff --git a/models/export.py b/models/export.py index 168d24d50..4715596d7 100644 --- a/models/export.py +++ b/models/export.py @@ -9,7 +9,7 @@ import argparse import torch import torch.nn as nn -from models.common import Conv +import models from models.experimental import attempt_load from utils.activations import Hardswish from utils.general import set_logging @@ -33,10 +33,10 @@ if __name__ == '__main__': # Update model for k, m in model.named_modules(): m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatability - if isinstance(m, Conv) and isinstance(m.act, nn.Hardswish): + if isinstance(m, models.common.Conv) and isinstance(m.act, nn.Hardswish): m.act = Hardswish() # assign activation - # if isinstance(m, Detect): - # m.forward = m.forward_export # assign forward (optional) + # if 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