How to export a trained lane-det model from PyTorch to ONNX

[cumtchw]

How to export a trained lane-det model from PyTorch to ONNX

[cfzd]

@cumtchw
I haven't try ONNX for now. But I think it would be pretty easy since the model only contains conv, bn, relu, and mlp, which are all standard components in PyTorch. Please refer to here.

[cumtchw]

line 46 to line 63 is the codes about export pth model to onne model:
`import cv2
from model.model import parsingNet
from utils.common import merge_config
from utils.dist_utils import dist_print
import torch
import scipy.special
import numpy as np
import torchvision.transforms as transforms
from data.dataset import LaneTestDataset
from data.constant import culane_row_anchor, tusimple_row_anchor
from PIL import Image

import onnx
import onnxruntime as ort

test_img = 'tmp/1.jpg'

if name == "main":
torch.backends.cudnn.benchmark = True
args, cfg = merge_config()
dist_print('start testing...')
assert cfg.backbone in ['18', '34', '50', '101', '152', '50next', '101next', '50wide', '101wide']

if cfg.dataset == 'CULane':
    cls_num_per_lane = 18
elif cfg.dataset == 'Tusimple':
    cls_num_per_lane = 56
else:
    raise NotImplementedError

net = parsingNet(pretrained=False, backbone=cfg.backbone, cls_dim=(cfg.griding_num + 1, cls_num_per_lane, 4),
                 use_aux=False).cuda()  # we dont need auxiliary segmentation in testing
state_dict = torch.load(cfg.test_model, map_location='cpu')['model']
compatible_state_dict = {}
for k, v in state_dict.items():
    if 'module.' in k:
        compatible_state_dict[k[7:]] = v
    else:
        compatible_state_dict[k] = v

net.load_state_dict(compatible_state_dict, strict=False)
print("net====================", net)


#在这里增加转换为onnx模型的代码。
#img = torch.zeros(1, 3, 720, 1280).to('cuda')
img = torch.zeros(1, 3, 288, 800).to('cuda')
onne_name = "lane.onnx"
torch.onnx.export(net, img, onne_name, verbose=True)


model = onnx.load(onne_name)
# Check that the IR is well formed
onnx.checker.check_model(model)
# Print a human readable representation of the graph
print("==============================================================================================")
onnx.helper.printable_graph(model.graph)


ort_session = ort.InferenceSession(onne_name)
outputs = ort_session.run(None,  {'input.1':np.random.randn(1, 3, 288, 800).astype(np.float32)})
print("outputs[0]:::::", outputs)






net.eval()

img_transforms = transforms.Compose([
    # transforms.CenterCrop((590,1640)),
    transforms.Resize((288, 800)),
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])

frame = cv2.imread(test_img)
print('输入图片大小为:',frame.shape)
# oir_shape = (frame_ori.shape[1],frame_ori.shape[0])
# frame = cv2.resize(frame_ori, (1280, 720), interpolation=cv2.INTER_LINEAR)
img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
img_ = Image.fromarray(img)
imgs = img_transforms(img_)
imgs = imgs.unsqueeze(0)
imgs = imgs.cuda()
with torch.no_grad():
    out = net(imgs)

col_sample = np.linspace(0, 800 - 1, cfg.griding_num)
col_sample_w = col_sample[1] - col_sample[0]

out_j = out[0].data.cpu().numpy()
out_j = out_j[:, ::-1, :]
prob = scipy.special.softmax(out_j[:-1, :, :], axis=0)
idx = np.arange(cfg.griding_num) + 1
idx = idx.reshape(-1, 1, 1)
loc = np.sum(prob * idx, axis=0)
out_j = np.argmax(out_j, axis=0)
loc[out_j == cfg.griding_num] = 0
out_j = loc

for i in range(out_j.shape[1]):
    if np.sum(out_j[:, i] != 0) > 2:
        for k in range(out_j.shape[0]):
            if out_j[k, i] > 0:
                ppp = (int(out_j[k, i] * col_sample_w * frame.shape[1] / 800) - 1, int(frame.shape[0] * (tusimple_row_anchor[56-1-k]/288)) - 1 )
                cv2.circle(frame, ppp, 5, (0, 255, 0), -1)
# frame = cv2.resize(frame,oir_shape,interpolation=cv2.INTER_LINEAR)
print('输出图片大小为:',frame.shape)
#cv2.imshow('result',frame)
#cv2.waitKey(0)

`