您还可以通过网络定义API直接向TensorRT定义网络,而不是使用解析器。此场景假设主机内存中的每层权重已经准备好在网络创建期间传递给TensorRT。
在下面的示例中,我们将创建一个简单的网络,其中包含输入、卷积、池、完全连接、激活和SoftMax层。要全面了解代码,请参考在/usr/src/tensorrt/samples/sampleMNISTAPI目录中
逐层构建简单MNIST网络https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html#mnistapi_sample示例。
1.创建构建器和网络:
IBuilder* builder = createInferBuilder(gLogger);
INetworkDefinition* network = builder->createNetwork();
2.将输入层添加到具有输入维度的网络中。一个网络可以有多个输入,尽管在这个例子中只有一个:
auto data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{1, INPUT_H, INPUT_W});
3.添加含有隐含层输入节点的卷积层、滤波和偏置的步长和权重。为了从层中检索张量引用,我们可以使用:
auto conv1 = network->addConvolution(*data->getOutput(0), 20, DimsHW{5, 5}, weightMap["conv1filter"], weightMap["conv1bias"]);
conv1->setStride(DimsHW{1, 1});
注意:传递给TensorRT层的权重位于主机内存中。
4.添加池层:
auto pool1 = network->addPooling(*conv1->getOutput(0), PoolingType::kMAX, DimsHW{2, 2});
pool1->setStride(DimsHW{2, 2});
5.添加全连接激活层:
auto ip1 = network->addFullyConnected(*pool1->getOutput(0), 500, weightMap["ip1filter"], weightMap["ip1bias"]);
auto relu1 = network->addActivation(*ip1->getOutput(0), ActivationType::kRELU);
6.添加SoftMax层,计算最终概率,设置为输出:
auto prob = network->addSoftMax(*relu1->getOutput(0));
prob->getOutput(0)->setName(OUTPUT_BLOB_NAME);
7.Mark输出:
network->markOutput(*prob->getOutput(0));