deep-learning - 函数 IBuilder::buildEngineWithConfig() 返回 null
问题描述
我正在使用 tensorRT 构建一个小模型,如下所示:
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#include <fstream>
#include <map>
#include <chrono>
#include <iostream>
#include "include/Utils.h"
#include <memory>
#include <vector>
#include <cassert>
#include "src/InferDeleter.cpp"
using namespace std;
using namespace nvinfer1;
class MyLogger : public ILogger {
void log(Severity severity, const char *msg) override {
if (severity != Severity::kINFO) {
cout << msg << endl;
}
}
} gLogger;
int main() {
//load weights
map<string, Weights> mWeightMap = Utils::getInstance().loadWeights("Weights/mnistapi.wts");
//a few configuration parameters
const char *INPUT_BLOB_NAME = "input";
const char *OUTPUT_BLOB_NAME = "output";
DataType dataType = nvinfer1::DataType::kFLOAT;
int INPUT_H = 28, INPUT_W = 28;
int batchSize = 1;
//define the network
IBuilder *builder = createInferBuilder(gLogger);
INetworkDefinition *network = builder->createNetworkV2(0U);
// Create input tensor of shape { 1, 1, 28, 28 }
ITensor *data = network->addInput(
INPUT_BLOB_NAME, DataType::kFLOAT, Dims3{1, INPUT_H, INPUT_W});
// Create scale layer with default power/shift and specified scale parameter.
const float scaleParam = 0.0125f;
const Weights power{DataType::kFLOAT, nullptr, 0};
const Weights shift{DataType::kFLOAT, nullptr, 0};
const Weights scale{DataType::kFLOAT, &scaleParam, 1};
IScaleLayer *scale_1 = network->addScale(*data, ScaleMode::kUNIFORM, shift, scale, power);
// Add convolution layer with 20 outputs and a 5x5 filter.
IConvolutionLayer *conv1 = network->addConvolutionNd(
*scale_1->getOutput(0), 20, Dims{2, {5, 5}, {}}, mWeightMap["conv1filter"], mWeightMap["conv1bias"]);
conv1->setStride(DimsHW{1, 1});
// Add max pooling layer with stride of 2x2 and kernel size of 2x2.
IPoolingLayer *pool1 = network->addPoolingNd(*conv1->getOutput(0), PoolingType::kMAX, Dims{2, {2, 2}, {}});
pool1->setStride(DimsHW{2, 2});
// Add second convolution layer with 50 outputs and a 5x5 filter.
IConvolutionLayer *conv2 = network->addConvolutionNd(
*pool1->getOutput(0), 50, Dims{2, {5, 5}, {}}, mWeightMap["conv2filter"], mWeightMap["conv2bias"]);
conv2->setStride(DimsHW{1, 1});
// Add second max pooling layer with stride of 2x2 and kernel size of 2x3>
IPoolingLayer *pool2 = network->addPoolingNd(*conv2->getOutput(0), PoolingType::kMAX, Dims{2, {2, 2}, {}});
pool2->setStride(DimsHW{2, 2});
// Add fully connected layer with 500 outputs.
IFullyConnectedLayer *ip1
= network->addFullyConnected(*pool2->getOutput(0), 500, mWeightMap["ip1filter"], mWeightMap["ip1bias"]);
// Add activation layer using the ReLU algorithm.
IActivationLayer *relu1 = network->addActivation(*ip1->getOutput(0), ActivationType::kRELU);
// Add second fully connected layer with 20 outputs.
IFullyConnectedLayer *ip2 = network->addFullyConnected(
*relu1->getOutput(0), 10, mWeightMap["ip2filter"], mWeightMap["ip2bias"]);
// Add softmax layer to determine the probability.
ISoftMaxLayer *prob = network->addSoftMax(*ip2->getOutput(0));
prob->getOutput(0)->setName(OUTPUT_BLOB_NAME);
network->markOutput(*prob->getOutput(0));
//build engine
IBuilderConfig *builderConfig = builder->createBuilderConfig();
builder->setMaxBatchSize(batchSize);
builderConfig->setMaxWorkspaceSize(1<<24);
//engine null
ICudaEngine *engine = builder->buildEngineWithConfig(*network, *builderConfig);
//later uses of engine.
return 0;
}
但是,该函数builder->buildEngineWithConfig(*network, *builderConfig)
返回nullptr
. 我试图更改maxWorkSpace
为其他值,但它仍然不起作用。我也访问了这篇文章,但没有任何帮助。任何人都指出了问题的原因。谢谢!
解决方案
经过几天的滚动处理这个问题。我发现如果模型中的层与传入的权重不匹配,则不会出现错误,但您无法创建 TensorRT 引擎来执行后续任务。因此,在这种情况下最好的方法是仔细检查逐层和.wts
文件。
推荐阅读
- r - 从数据框中提取特定单元格并将数据展平为单行
- linux - 如何返回创建的菜单的不同部分
- reactjs - 在 REACT 中从 cyrptoicons 导入 SVG 时遇到问题
- python - Python:将 csv 转换为 xlsx 时出现问题。存储为文本的数字
- json - 是否可以在满足条件时使用 JSON 将用户重定向到不同的视图
- c++ - 如果我在代码块中运行这个程序,我没有得到正确的 ans,但如果我调试它,我就会得到正确的 ans。它正在发生一圈又一圈
- r - R - 数据框之间的匹配列
- c - 尝试切换字符串中的字母时出现总线错误
- latex - 使用 Science.bst 模板文件,无法正确显示期刊文章的参考文献
- r - 将数据框中的列文本转换为一行内的多行