pytorch转换onnx,测试onnx,pytorch 模型结果是否一致

def to_onnx():
    dummy_input = torch.randn(1, 3, 112, 112, dtype=torch.float)
    # model = model_res()
    model = model_osnet()

    input_names = ["data"]
    output_names = ["fc"]
    torch.onnx.export(
        model,
        dummy_input,
        "./osnet.onnx",
        verbose=True,
        input_names=input_names,
        output_names=output_names,
    )
    print("转换模型成功^^")


def pytorch_out(input):
    model = model_res() #model.eval
    # input = input.cuda()
    # model.cuda()
    torch.no_grad()
    output = model(input)
    # print output[0].flatten()[70:80]
    return output

def pytorch_onnx_test():
    import onnxruntime
    from onnxruntime.datasets import get_example

    def to_numpy(tensor):
        return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

    # 测试数据
    torch.manual_seed(66)
    dummy_input = torch.randn(1, 3, 112, 112, device='cpu')

    example_model = get_example("/home/shiyy/nas/all_workspace/pytorch_code/reid/InsightFace-v2/res50.onnx")
    # netron.start(example_model) 使用 netron python 包可视化网络
    sess = onnxruntime.InferenceSession(example_model)

    # onnx 网络输出
    onnx_out = np.array(sess.run(None, { "data": to_numpy(dummy_input)}))  #fc 输出是三维列表
    print("==============>")
    print(onnx_out)
    print(onnx_out.shape)
    print("==============>")
    torch_out_res = pytorch_out(dummy_input).detach().numpy()   #fc输出是二维 列表
    print(torch_out_res)
    print(torch_out_res.shape)

    print("===================================>")
    print("输出结果验证小数点后五位是否正确,都变成一维np")

    torch_out_res = torch_out_res.flatten()
    onnx_out = onnx_out.flatten()

    pytor = np.array(torch_out_res,dtype="float32") #need to float32
    onn=np.array(onnx_out,dtype="float32")  ##need to float32
    np.testing.assert_almost_equal(pytor,onn, decimal=5)  #精确到小数点后5位，验证是否正确，不正确会自动打印信息
    print("恭喜你 ^^ ，onnx 和 pytorch 结果一致 ， Exported model has been executed decimal=5 and the result looks good!")

	 [[...
	  -5.49954772e-02 -7.31383190e-02  2.37192452e-01  1.48571879e-02
	  -2.97113061e-02 -5.73663861e-02 -6.34638742e-02 -5.24816178e-02]]
	(1, 512)
	===================================>
	输出结果验证小数点后五位是否正确,都变成一维np
	恭喜你 ^^ ，onnx 和 pytorch 结果一致 ， Exported model has been executed decimal=5 and the result looks good!