TensorRT 4でONNX

@Vengineerの戯言 : Twitter
SystemVerilogの世界へようこそ、すべては、SystemC v0.9公開から始まった

昨日の続き、

TensorRT 4 Accelerates Neural Machine Translation, Recommenders, and Speechの下の方にいくと、

Support for ONNX format

とあります。Python API と C++ API があるようです。
下記は、Python API での ONNX フォーマットのファイルを読み込んで、TensorRT に変換しているところまで。

引用
import tensorrt as trt
// Import NvOnnxParser, use config object to pass user args to the parser object 
   from tensorrt.parsers import onnxparser
   apex = onnxparser.create_onnxconfig()
// Parse the trained model and generate TensorRT engine
   apex.set_model_file_name("model_file_path")
   apex.set_model_dtype(trt.infer.DataType.FLOAT)
//create the ONNX parser
   trt_parser = onnxparser.create_onnxparser(apex)
   data_type = apex.get_model_dtype()
   onnx_filename = apex.get_model_file_name()
// Generate the TensorRT network after parsing the model file
   trt_parser.parse(onnx_filename, data_type)
// retrieve the network from the parser
   trt_parser.convert_to_trtnetwork()
   trt_network = trt_parser.get_trtnetwork()

C++ API では、下記のような感じ

引用
//Create the ONNX parser using SampleConfig object to pass the input arguments
   nvonnxparser::IOnnxConfig* config = nvonnxparser::createONNXConfig();
//Create Parser
   nvonnxparser::IONNXParser* parser = nvonnxparser::createONNXParser(*config);
//Ingest the model
   parser->parse(onnx_filename, DataType::kFLOAT);
//Convert the model to a TensorRT network:
   parser->convertToTRTNetwork();
//Obtain the network from the model:
   nvinfer1::INetworkDefinition* trtNetwork = parser->getTRTNetwork();

そんでもって、この ONNX サポートは、githubにてオープンソースとして公開していますよ。
まだ、TensorRT は、3 +となっていて、一部のファイルではTRT4とコメントがあります。