bytedeco/javacpp-presets

[TensorRT] How to save engine to disk

iyLester opened this issue · 2 comments

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");

Try to call this instead:

fc.write(engine.data().capacity(engine.size()).asByteBuffer());

It work, thanks