onnx替换pt进行推理代码

最新推荐文章于 2024-07-12 16:16:27 发布

SetMaker

最新推荐文章于 2024-07-12 16:16:27 发布

阅读量343

点赞数

文章标签： python pytorch 深度学习算法

本文链接：https://blog.csdn.net/qq_40962125/article/details/134239176

版权

import onnxruntime as rt
def onnx_infer(im, onnx_model):
    # InferenceSession获取onnxruntime解释器
    sess = rt.InferenceSession(onnx_model)
    # 模型的输入输出名，必须和onnx的输入输出名相同，可以通过netron查看，如何查看参考下文
    input_name = "input1"
    
    # run方法用于模型推理，run(输出张量名列表，输入值字典)
    output = sess.run(None, {input_name: im})
    return output

var_L = data['LQ'].to("cuda")
#这里要注意的是 传进去的input类型要是np的，然后inputname要去onnx可视化里面看一下。

val_np = var_L.cpu().numpy()
out = onnx_infer(val_np,"../pretrained_models/100w_crop.onnx")
#这里可以打印一下out，不同的任务out不一样，也是和onnx对应
#下面是我任务的专门的后处理，仅供参考。
o1 = out[0]
o1 = np.squeeze(o1,0)
o1[o1 < 0] = 0
o1 = np.transpose(o1[[2, 1, 0], :, :], (1, 2, 0))  # HWC, BGR
o1.astype(np.float32)
sr_img = o1
util.save_img_with_ratio(image_path, sr_img, alignratio_path)