WongKinYiu / ScaledYOLOv4

Scaled-YOLOv4: Scaling Cross Stage Partial Network

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Fail to transfer .pt into .onnx

ziqi-jin opened this issue · 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!