Fail to transfer .pt into .onnx
ziqi-jin opened this issue · comments
ziqi-jin commented
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
intomodel.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!