3.2. onnx export multi_batch

前言

将onnx bs=1 修改为多batch操作

参考链接:

https://www.cnblogs.com/tangjunjun/p/16500116.html

https://blog.csdn.net/weixin_43863869/article/details/128638397?spm=1001.2101.3001.6650.3&utm_medium=distribute.pc_relevant.none-task-blog-2%7Edefault%7ECTRLIST%7ERate-3-128638397-blog-115749006.235%5Ev38%5Epc_relevant_anti_vip&depth_1-utm_source=distribute.pc_relevant.none-task-blog-2%7Edefault%7ECTRLIST%7ERate-3-128638397-blog-115749006.235%5Ev38%5Epc_relevant_anti_vip&utm_relevant_index=6

Code

# 方式一: 固定shape 
input_data = torch.rand((64, 3, 224, 224))
with torch.no_grad():
     torch.onnx.export(
            model, input_data, # (img_list, ),
            output_file,
            input_names=['input'],
            output_names=['probs'],
            # export_params=True,
            # keep_initializers_as_inputs=True,
            # 

你可能感兴趣的:(模型推理,batch,深度学习,开发语言)