Creating A Network Definition From Scratch Using The C++ API

  • 0 replies
  • 2582 views
*

sisiy

  • *****
  • 246
    • 查看个人资料
Creating A Network Definition From Scratch Using The C++ API
« 于: 七月 30, 2019, 10:36:51 am »
您还可以通过网络定义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));