1、创建模型
class TRTLogger : public nvinfer1::ILogger{
public:
virtual void log(Severity severity, nvinfer1::AsciiChar const* msg) noexcept override{
if(severity <= Severity::kINFO){
printf("%d: %s\n", severity, msg);
}
}
} logger;
nvinfer1::Weights make_weights(float* ptr, int n){
nvinfer1::Weights w;
w.count = n;
w.type = nvinfer1::DataType::kFLOAT;
w.values = ptr;
return w;
}
bool build_model(){
TRTLogger logger;
// 这是基本需要的组件
nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger);
nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
nvinfer1::INetworkDefinition* network = builder->createNetworkV2(1);
// 构建一个模型
/*
Network definition:
image
|
linear (fully connected) input = 3, output = 2, bias = True w=[[1.0, 2.0, 0.5], [0.1, 0.2, 0.5]], b=[0.3, 0.8]
|
sigmoid
|
prob
*/
const int num_input = 3;
const int num_output = 2;
float layer1_weight_values[] = {1.0, 2.0, 0.5, 0.1, 0.2, 0.5};
float layer1_bias_values[] = {0.3, 0.8};
nvinfer1::ITensor* input = network->addInput("image", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4(1, num_input, 1, 1));
nvinfer1::Weights layer1_weight = make_weights(layer1_weight_values, 6);
nvinfer1::Weights layer1_bias = make_weights(layer1_bias_values, 2);
auto layer1 = network->addFullyConnected(*input, num_output, layer1_weight, layer1_bias);
auto prob = network->addActivation(*layer1->getOutput(0), nvinfer1::ActivationType::kSIGMOID);
// 将我们需要的prob标记为输出
network->markOutput(*prob->getOutput(0));
printf("Workspace Size = %.2f MB\n", (1 << 28) / 1024.0f / 1024.0f);
config->setMaxWorkspaceSize(1 << 28);
builder->setMaxBatchSize(1);
nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
if(engine == nullptr){
printf("Build engine failed.\n");
return false;
}
// 将模型序列化,并储存为文件
nvinfer1::IHostMemory* model_data = engine->serialize();
FILE* f = fopen("engine.trtmodel", "wb");
fwrite(model_data->data(), 1, model_data->size(), f);
fclose(f);
// 卸载顺序按照构建顺序倒序
model_data->destroy();
engine->destroy();
network->destroy();
config->destroy();
builder->destroy();
printf("Done.\n");
return true;
}
这些都是上一节的代码,就不多说了。
2、main函数
int main(){
if(!build_model()){
return -1;
}
inference();
return 0;
}
这个也很简单,就是模型如果build成功了就推理,反之则不推理。
3、inference
3.1 inference前文-load_file
vector<unsigned char> load_file(const string& file){
ifstream in(file, ios::in | ios::binary);
if (!in.is_open())
return {};
in.seekg(0, ios::end);
size_t length = in.tellg();
std::vector<uint8_t> data;
if (length > 0){
in.seekg(0, ios::beg);
data.resize(length);
in.read((char*)&data[0], length);
}
in.close();
return data;
}
这段代码实现了一个函数 load_file,用于从文件中读取二进制数据并返回一个包含数据的无符号字符向量。
以下是对代码的解释:
函数声明:vector<unsigned char> load_file(const string& file)
函数名:load_file
参数:const string& file,表示要加载的文件的路径和名称
返回类型:vector<unsigned char>,表示读取的二进制数据
打开文件:ifstream in(file, ios::in | ios::binary);
使用 ifstream 类创建输入流对象 in ,并指定打开方式为以二进制形式打开进行读取。
检查文件是否成功打开: if (!in.is_open()) return {};
如果文件打开失败,则返回一个空向量({}),表示读取失败。
获取文件长度:in.seekg(0, ios::end); size_t length = in.tellg();
将文件位置移至文件末尾,并使用 tellg 函数获取当前位置的偏移量,即文件的大小。
创建存储数据的向量:std::vector<uint8_t> data;
声明一个名为 data 的无符号字符型向量,用于存储从文件中读取的数据。
判断文件长度是否大于0:if (length > 0){...}
如果文件长度大于0,说明文件非空,则执行下面的代码块。
设置文件位置为开头并调整向量大小:in.seekg(0, ios::beg); data.resize(length);
将文件位置移至文件开头,并使用 resize 函数调整向量 data 的大小,使其能够容纳读取的数据。
读取文件内容到向量:in.read((char*)&data[0], length);
使用 read 函数将文件中长度为 length 的二进制数据读取到向量 data 中。(char*) 是进行类型转换操作。
关闭文件流:in.close();
使用 close 函数关闭文件流,释放资源。
返回读取的数据向量: return data;
将存储数据的向量 data 返回给调用该函数的部分,表示成功读取了文件的二进制数据。
3.2 inference正文-1模型加载
void inference(){
// ------------------------------ 1. 准备模型并加载 ----------------------------
TRTLogger logger;
auto engine_data = load_file("engine.trtmodel");
// 执行推理前,需要创建一个推理的runtime接口实例。与builer一样,runtime需要logger:
nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(logger);
// 将模型从读取到engine_data中,则可以对其进行反序列化以获得engine
nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(engine_data.data(), engine_data.size());
if(engine == nullptr){
printf("Deserialize cuda engine failed.\n");
runtime->destroy();
return;
}
nvinfer1::IExecutionContext* execution_context = engine->createExecutionContext();
cudaStream_t stream = nullptr;
// 因为是异步的所以创建CUDA流,以确定这个batch的推理是独立的
cudaStreamCreate(&stream);
先将我们上一篇所设计的模型用load file读取出来
nvinfer1::createInferRuntime(logger)
函数通过传入日志记录器 logger
来创建一个 TensorRT 推理运行时对象。
创建后的 runtime
变量将用于加载和执行 TensorRT 模型,并进行推理任务。
3.2 inference正文-2 搬运到GPU
// ------------------------------ 2. 准备好要推理的数据并搬运到GPU ----------------------------
float input_data_host[] = {1, 2, 3};
float* input_data_device = nullptr;
float output_data_host[2];
float* output_data_device = nullptr;
cudaMalloc(&input_data_device, sizeof(input_data_host));
cudaMalloc(&output_data_device, sizeof(output_data_host));
cudaMemcpyAsync(input_data_device, input_data_host, sizeof(input_data_host), cudaMemcpyHostToDevice, stream);
// 用一个指针数组指定input和output在gpu中的指针。
float* bindings[] = {input_data_device, output_data_device};
bindings是tensorRT对输入输出张量的描述,bindings = input-tensor + output-tensor。比如input有a,output有b, c, d,那么bindings = [a, b, c, d],bindings[0] = a,bindings[2] = c。此时看到engine->getBindingDimensions(0) 我们得知道获取的是什么
3.2 inference正文-3 拿回到CPU
// ------------------------------ 3. 推理并将结果搬运回CPU ----------------------------
bool success = execution_context->enqueueV2((void**)bindings, stream, nullptr);
cudaMemcpyAsync(output_data_host, output_data_device, sizeof(output_data_host), cudaMemcpyDeviceToHost, stream);
cudaStreamSynchronize(stream);
printf("output_data_host = %f, %f\n", output_data_host[0], output_data_host[1]);
这里一定要提前准备好 内存,比如上面的float output_data_host[2];
执行enqueueV2就可以进行推理,注意是异步推理,加入到stream队列等待执行。输入的bindings则是tensors的指针(注意是device pointer)。其shape对应于编译时指定的输入输出的shape(这里只演示全部shape静态)
3.2 inference正文-4 释放内存
// ------------------------------ 4. 释放内存 ----------------------------
printf("Clean memory\n");
cudaStreamDestroy(stream);
execution_context->destroy();
engine->destroy();
runtime->destroy();
3.2 inference正文-5 推理验证
// ------------------------------ 5. 手动推理进行验证 ----------------------------
const int num_input = 3;
const int num_output = 2;
float layer1_weight_values[] = {1.0, 2.0, 0.5, 0.1, 0.2, 0.5};
float layer1_bias_values[] = {0.3, 0.8};
printf("手动验证计算结果:\n");
for(int io = 0; io < num_output; ++io){
float output_host = layer1_bias_values[io];
for(int ii = 0; ii < num_input; ++ii){
output_host += layer1_weight_values[io * num_input + ii] * input_data_host[ii];
}
// sigmoid
float prob = 1 / (1 + exp(-output_host));
printf("output_prob[%d] = %f\n", io, prob);
}
}