[TensorRT] How to save engine to disk
iyLester opened this issue · 2 comments
iyLester commented
I tried to build the engine using the code, and tried to save the engine, but the size of the saved engine file is only 1 bytes, which is obviously a problem.
I have tried to deserialize IHostMemory to ICudaEngine to make sure it can be inference fine.
String onnxPath = "./model.onnx";
ILogger logger = new Logger();
IBuilder builder = createInferBuilder(logger);
int flag = 1 << NetworkDefinitionCreationFlag.kEXPLICIT_BATCH.value;
flag = flag | ( 1 << NetworkDefinitionCreationFlag.kEXPLICIT_PRECISION.value);
INetworkDefinition network = builder.createNetworkV2(flag);
IParser parser = createParser(network, logger);
parser.parseFromFile(onnxPath, ILogger.Severity.kVERBOSE.value);
IBuilderConfig config = builder.createBuilderConfig();
config.setBuilderOptimizationLevel(0);
IOptimizationProfile profile = builder.createOptimizationProfile();
// Set Model MIN/OPT/MAX Shape....
config.addOptimizationProfile(profile);
IHostMemory engine = builder.buildSerializedNetwork(network, config);
FileChannel fc = null;
try {
fc = new FileOutputStream("model.engine").getChannel();
fc.write(engine.data().asByteBuffer());
fc.close();
} catch (Exception e) {
System.out.println(e.getMessage());
}
engine.destroy();
parser.destroy();
network.destroy();
config.destroy();
builder.destroy();
System.out.println("END");
saudet commented
Try to call this instead:
fc.write(engine.data().capacity(engine.size()).asByteBuffer());
iyLester commented
It work, thanks