ScaledYOLOv4 icon indicating copy to clipboard operation
ScaledYOLOv4 copied to clipboard

Fail to transfer .pt into .onnx

Open ziqi-jin opened this issue 1 year ago • 0 comments

if you can not generate .onnx file when you are using the models/export.py, maybe you could try my method.

1. using the docker container supplied by the author in the README.md

2. modify the models/export.py in 2 steps:

  • step1: change the model.model[-1].export = True into model.model[-1].export = False
  • step2: add code
    for k, m in model.named_modules():
        m._non_persistent_buffers_set = set()  # pytorch 1.6.0 compatability
        if isinstance(m, models.common.Conv) and isinstance(m.act, models.common.Mish):
            m.act = Mish()  # assign activation
        if isinstance(m, models.common.BottleneckCSP) or isinstance(m, models.common.BottleneckCSP2) \
                or isinstance(m, models.common.SPPCSP):
            if isinstance(m.bn, nn.SyncBatchNorm):
                bn = nn.BatchNorm2d(m.bn.num_features, eps=m.bn.eps, momentum=m.bn.momentum)
                bn.training = False
                bn._buffers = m.bn._buffers
                bn._non_persistent_buffers_set = set()
                m.bn = bn
            if isinstance(m.act, models.common.Mish):
                m.act = Mish()  # assign activation

above the line model.model[-1].export = False

  • you may also need to import some necessary package like torch.nn enjoy it!

ziqi-jin avatar Jul 22 '22 03:07 ziqi-jin