英文:
Python set ONNX runtime to return tensor instead of numpy array
问题
在Python中，我加载了我的预定义模型（super-gradients，yolox-s）：
onnx_session = onnxrt.InferenceSession("yolox_s_640_640.onnx")
然后我加载一些数据并运行它：
dataset = MyCostumeDataset(args.path, 'val')
val_dataloader = DataLoader(dataset, batch_size=args.bsize)
for inputs in val_dataloader:
onnx_inputs = {onnx_session.get_inputs()[0].name: inputs}
# inputs.shape: torch.Size([4, 3, 640, 640]), 即这是一个Tensor
raw_predictions = onnx_session.run(None, onnx_inputs)
# 这返回一个numpy数组的列表:
# type(raw_predictions[0])
# <class 'numpy.ndarray'>
# raw_predictions[0].shape
# (4, 8400, 85)
到目前为止，它正在按预期工作，除了我希望它默认返回torch.Tensor的列表，而不是numpy数组。我对ONNX和PyTorch都不熟悉，感觉我可能错过了一些基本的东西。
如何让onnx_session返回torch.Tensor的列表，而不是numpy数组？这将减少一些转换的开销。谢谢！
英文:
In python I'm loading my predefined model (super-gradients, yolox-s):
onnx_session = onnxrt.InferenceSession("yolox_s_640_640.onnx")
Then I load some data and run it:
dataset = MyCostumeDataset(args.path, 'val')
val_dataloader = DataLoader(dataset, batch_size=args.bsize)
for inputs in val_dataloader:
onnx_inputs = {onnx_session.get_inputs()[0].name: inputs}
# inputs.shape: torch.Size([4, 3, 640, 640]), i.e., this is a Tensor
raw_predictions = onnx_session.run(None, onnx_inputs)
# this returns a list of numpy arrays:
# type(raw_predictions[0])
# <class 'numpy.ndarray'>
# raw_predictions[0].shape
# (4, 8400, 85)
So far it is working as it should, except I'd like it to return, by default, a list of Tensors (torch.Tensor) instead of numpy array. I'm new to both ONNX and PyTorch, and I'm feeling like this is something basic that I'm missing here.
How can I get onnx_session to return a list of torch.Tensor, instead of numpy arrays? This will same some overhead in the conversion. Thanks!
答案1
得分: 1
根据文档，您可以获取以下内容 - 返回结果：结果列表，每个结果可以是NumPy数组、稀疏张量、列表或字典。
英文:
According to the documentation you can get either - Returns: list of results, every result is either a numpy array, a sparse tensor, a list or a dictionary.
通过集体智慧和协作来改善编程学习和解决问题的方式。致力于成为全球开发者共同参与的知识库，让每个人都能够通过互相帮助和分享经验来进步。
评论