以一个简单的CNN网络为例
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
一、准备模型并加载
1. 加载模型
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;
}
TRTLogger logger;
auto engine_data = load_file("engine.trtmodel");
2. 创建runtime:createInferRuntime(logger)
// 执行推理前,需要创建一个推理的runtime接口实例。与builer一样,runtime需要logger:
nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(logger);
3. 反序列化engine
得到engine提供数据:runtime->deserializeCudaEngine(modelData, modelSize)
其中modelData
包含的是input和output的名字,形状,大小和数据类型
class ModelData(object):
INPUT_NAME = "data"
INPUT_SHAPE = (1, 1, 28, 28) // [B, C, H, W]
OUTPUT_NAME = "prob"
OUTPUT_SIZE = 10
DTYPE = trt.float32
// 将模型从读取到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;
}
4. 从engine创建执行上下文:engine->createExecutionContext()
nvinfer1::IExecutionContext* execution_context = engine->createExecutionContext();
5. 创建cuda流:cudaStreamCreate(&stream)
:
- CUDA编程流是组织异步工作的一种方式,创建流来确定batch推理的独立
- 为每个独立batch使用IExecutionContext,并为每个独立批次使用cudaStreamCreate创建CUDA流
cudaStream_t stream = nullptr;
// 创建CUDA流,以确定这个batch的推理是独立的
cudaStreamCreate(&stream);
二、准备数据并搬运到gpu
1. 在host上声明input
数据和output
数组大小,搬运到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);
2. 用一个指针数组指定input
和output
在gpu中的指针
执行inference的必须操作
float* bindings[] = {input_data_device, output_data_device};
三、推理并将结果搬运回cpu
1. 推理并结果搬运回cpu
bool success = execution_context->enqueueV2((void**)bindings, stream, nullptr);
cudaMemcpyAsync(output_data_host, output_data_device, sizeof(output_data_host), cudaMemcpyDeviceToHost, stream);
2. 启动所有工作后,与所有流同步以等待结果:cudaStreamSynchronize
cudaStreamSynchronize(stream);
四、按照与创建相反顺序释放内存
cudaStreamDestroy(stream);
execution_context->destroy();
engine->destroy();
runtime->destroy();