mirrors/mmsegmentation
1
0
Fork 0
mirror of https://github.com/open-mmlab/mmsegmentation.git synced 2025-06-03 22:03:48 +08:00
Code Issues Projects Releases Wiki Activity

fix verify bugs (#547)

Browse Source 
* fix verify bugs

* rename args
This commit is contained in:
q.yao 2021-05-13 12:14:41 +08:00 committed by GitHub
parent 805eda81ea
commit 8558163753
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 9 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
tools/pytorch2onnx.py
View File

@ -96,13 +96,16 @@ def _prepare_input_img(img_path,
return mm_inputs return mm_inputs
def _update_input_img(img_list, img_meta_list): def _update_input_img(img_list, img_meta_list, update_ori_shape=False):
# update img and its meta list # update img and its meta list
N = img_list[0].size(0) N, C, H, W = img_list[0].shape
img_meta = img_meta_list[0][0] img_meta = img_meta_list[0][0]
img_shape = img_meta['img_shape'] img_shape = (H, W, C)
ori_shape = img_meta['ori_shape'] if update_ori_shape:
pad_shape = img_meta['pad_shape'] ori_shape = img_shape
else:
ori_shape = img_meta['ori_shape']
pad_shape = img_shape
new_img_meta_list = [[{ new_img_meta_list = [[{
'img_shape': 'img_shape':
img_shape, img_shape,
@ -220,7 +223,7 @@ def pytorch2onnx(model,
# update img_meta # update img_meta
img_list, img_meta_list = _update_input_img( img_list, img_meta_list = _update_input_img(
img_list, img_meta_list) img_list, img_meta_list, test_mode == 'whole')
# check the numerical value # check the numerical value
# get pytorch output # get pytorch output