1
2
3
4
5
6
7
8
9
10
11
12
        import torch.onnx.symbolic

        # Override Upsample's ONNX export until new opset is supported
        @torch.onnx.symbolic.parse_args('v', 'is')
        def upsample_nearest2d(g, input, output_size):
            height_scale = float(output_size[-2]) / input.type().sizes()[-2]
            width_scale = float(output_size[-1]) / input.type().sizes()[-1]
            return g.op("Upsample", input,
                scales_f=(1, 1, height_scale, width_scale),
                mode_s="nearest")
        torch.onnx.symbolic.upsample_nearest2d = upsample_nearest2d

1
2
Attribute not found: height_scale

1
2
onnx2trt my_model.onnx -o my_engine.trt

1
2
Attribute not found: height_scale

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
void onnxToTRTModel(const std::string& modelFile, // name of the onnx model
                    unsigned int maxBatchSize,    // batch size - NB must be at least as large as the batch we want to run with
                    nvinfer1::IHostMemory*& trtModelStream,
                    nvinfer1::DataType dataType,
                    nvinfer1::IInt8Calibrator* calibrator,
                    std::string save_name) // output buffer for the TensorRT model
{
    int verbosity = (int)nvinfer1::ILogger::Severity::kINFO;
    // create the builder
    nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(gLogger);
    nvinfer1::INetworkDefinition* network = builder->createNetwork();

    auto parser = nvonnxparser::createParser(*network, gLogger);

    if (!parser->parseFromFile(modelFile.c_str(), verbosity))
    {
        string msg("failed to parse onnx file");
        gLogger.log(nvinfer1::ILogger::Severity::kERROR, msg.c_str());
        exit(EXIT_FAILURE);
    }
    if ((dataType == nvinfer1::DataType::kINT8 && !builder->platformHasFastInt8()) )
        exit(EXIT_FAILURE);  //如果不支持kint8或不支持khalf就返回false
    // Build the engine

    builder->setMaxBatchSize(maxBatchSize);
    builder->setMaxWorkspaceSize(4_GB); //不能超过你的实际能用的显存的大小，例如我的1060的可用为4.98GB，超过4.98GB会报错
    builder->setInt8Mode(dataType == nvinfer1::DataType ::kINT8);  //
    builder->setInt8Calibrator(calibrator);  //
    samplesCommon::enableDLA(builder, gUseDLACore);
    nvinfer1::ICudaEngine* engine = builder->buildCudaEngine(*network);
    assert(engine);

    // we can destroy the parser
    parser->destroy();

    // serialize the engine, then close everything down  序列化
    trtModelStream = engine->serialize();

    gieModelStream.write((const char*)trtModelStream->data(), trtModelStream->size());
    std::ofstream SaveFile(save_name, std::ios::out | std::ios::binary);
    SaveFile.seekp(0, std::ios::beg);
    SaveFile << gieModelStream.rdbuf();
    gieModelStream.seekg(0, gieModelStream.beg);


    engine->destroy();
    network->destroy();
    builder->destroy();
}