模仿学习模型ACT部署

首先下载act代码：https://github.com/tonyzhaozh/act

修改代码imitate_episodes.py，在第250行处增加：

torch.onnx.export(policy, (qpos, curr_image), "model.onnx", opset_version=13)

导出onnx模型结构如下：
onnxruntime推理脚本：

import numpy as np
import onnxruntimeonnx_session = onnxruntime.InferenceSession("model.onnx", providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])input_name = []
for node in onnx_session.get_inputs():input_name.append(node.name)output_name = []
for node in onnx_session.get_outputs():output_name.append(node.name)qpos = np.random.randn(1, 14).astype(np.float32)
curr_image = np.random.randn(1, 1, 3, 480, 640).astype(np.float32)inputs = {}
inputs['onnx::Gemm_0'] = qpos
inputs['tensor'] = curr_imageoutputs = onnx_session.run(None, inputs)
print(outputs)

tensorrt推理脚本：

import cv2
import numpy as np
import tensorrt as trt
import commonif __name__ == '__main__':logger = trt.Logger(trt.Logger.WARNING)with open("model.engine", "rb") as f, trt.Runtime(logger) as runtime:engine = runtime.deserialize_cuda_engine(f.read())context = engine.create_execution_context()inputs, outputs, bindings, stream = common.allocate_buffers(engine)qpos = np.random.randn(1, 14).astype(np.float32)curr_image = np.random.randn(1, 1, 3, 480, 640).astype(np.float32)np.copyto(inputs[0].host, qpos.ravel())np.copyto(inputs[1].host, curr_image.ravel())output = common.do_inference(context,engine=engine, bindings=bindings,inputs=inputs, outputs=outputs, stream=stream,)print(output)