From f70507168fb128ad4bb28febe111d6d79697e0c5 Mon Sep 17 00:00:00 2001 From: Lei Yang Date: Wed, 30 Sep 2020 10:11:34 +0800 Subject: [PATCH] Remove redundancies in pytorch2onnx (#160) * rm redundancies * re-add some packages --- tools/pytorch2onnx.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/tools/pytorch2onnx.py b/tools/pytorch2onnx.py index b245366..f22d4d3 100644 --- a/tools/pytorch2onnx.py +++ b/tools/pytorch2onnx.py @@ -5,11 +5,11 @@ import mmcv import numpy as np import onnxruntime as rt import torch -from torch import nn import torch._C import torch.serialization from mmcv.onnx import register_extra_symbolics from mmcv.runner import load_checkpoint +from torch import nn from mmseg.models import build_segmentor @@ -186,11 +186,6 @@ if __name__ == '__main__': # convert SyncBN to BN segmentor = _convert_batchnorm(segmentor) - if isinstance(segmentor.decode_head, nn.ModuleList): - num_classes = segmentor.decode_head[-1].num_classes - else: - num_classes = segmentor.decode_head.num_classes - if args.checkpoint: load_checkpoint(segmentor, args.checkpoint, map_location='cpu')