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)
onnx替换pt进行推理代码
猜你喜欢
转载自blog.csdn.net/qq_40962125/article/details/134239176
今日推荐
周排行