the result in onnx model is all zero.

[Bryce1998]

I convert the .pth model to .onnx model by this code:

import torch
import torch.onnx

batch_size = 1
input_shape = (3, 300, 300)


def convert_onnx(model):
    model.eval()
    dummy_input = torch.randn(batch_size, *input_shape, requires_grad=True)
    torch.onnx.export(
        model=model,
        args=dummy_input,
        f="ssd.onnx",
        export_params=True,
        opset_version=10,
        do_constant_folding=True,
        input_names=['input'],
        output_names=['output'],
        dynamic_axes={'input': {0: 'batch_size'},  # variable length axes
                      'output': {0: 'batch_size'}}
    )

    print('Model has been converted to ONNX.')

and run the model by this code:

import cv2
import torch
import onnxruntime
import matplotlib.pyplot as plt
import numpy as np
from dataset_loader import VOC_CLASSES as LABELS

from config import voc
from SSD_module import build_ssd
from convert import convert_onnx

ssd_onnx_model = './ssd.onnx'
img_file = '/home/bryce/datasets/VOCdevkit/VOC2007/JPEGImages/000015.jpg'

VOC_MEAN = (123.0, 117.0, 104.0)
GENERAL_MEAN = (0.485, 0.456, 0.406)
GENERAL_STD = (0.229, 0.224, 0.225)

def imageProcess(imagePath, isVOC=True):
    image = cv2.imread(imagePath, cv2.IMREAD_COLOR)
    rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    x = cv2.resize(image, (300, 300)).astype(np.float32)  # x shape [300, 300, 3]
    x = cv2.cvtColor(x, cv2.COLOR_BGR2RGB)  # BGR -> RGB

    if isVOC:
        x -= VOC_MEAN
    else:
        for i in range(x.shape[2]):
            x[:, :, i] = (x[:, :, i] / 255.0 - GENERAL_MEAN[i]) / GENERAL_STD[i]

    x = np.transpose(x, [2, 0, 1])  # [300, 300, 3] -> [3, 300, 300]
    # x = torch.from_numpy(x).permute(2, 0, 1)
    x = np.expand_dims(x, 0)  # [3, 300, 300] -> [1, 3, 300, 300]

    print(x.shape, rgb_image.shape)
    return x, rgb_image

def ssd_infer(onnxPath, image):
    session = onnxruntime.InferenceSession(onnxPath)

    # get the name of the first input of the model
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name

    print('Input Name:', input_name)
    print('Output Name:', output_name)
    print('image shape:', image.shape)
    res = session.run(None, {input_name: image})
    res = torch.tensor(res[0])
    return res


def show_result(rgb_image, result):
    plt.figure(figsize=(10, 10))
    colors = plt.cm.hsv(np.linspace(0, 1, 21)).tolist()
    plt.imshow(rgb_image)  # plot the image for matplotlib
    currentAxis = plt.gca()

    detections = result.data
    # scale each detection back up to the image
    scale = torch.Tensor(rgb_image.shape[1::-1]).repeat(2)
    print(detections.shape)

    for i in range(detections.size(1)):
        j = 0
        print(detections[0, i, j, 0])
        while detections[0, i, j, 0] >= 0.6:
            score = detections[0, i, j, 0]
            label_name = LABELS[i - 1]
            display_txt = '%s: %.2f' % (label_name, score)
            pt = (detections[0, i, j, 1:] * scale).cpu().numpy()
            coords = (pt[0], pt[1]), pt[2] - pt[0] + 1, pt[3] - pt[1] + 1
            color = colors[i]
            currentAxis.add_patch(plt.Rectangle(*coords, fill=False, edgecolor=color, linewidth=2))
            currentAxis.text(pt[0], pt[1], display_txt, bbox={'facecolor': color, 'alpha': 0.5})
            j += 1
    # plt.show()

if __name__ == '__main__':
    net = build_ssd('test', 300, 21)
    net.load_state_dict(
        {k.replace('module.', ''): v for k, v in torch.load("../weights/VOC.pth").items()})

    convert_onnx(net)
    input_image, rgb_image = imageProcess(img_file)
    res = ssd_infer(ssd_onnx_model, input_image)
    show_result(rgb_image, res)

Everything works fine but predictions are incorrect.

The model output is a tensor with shape [1, 21, 200, 5] but all the element is zero.

I have no idea for this issue.

Am I missing someting important?

By the way, load the .pth model I trained to run eval.py and demo.py all going well and correct.