FDTensor support GPU device (#190)

* fdtensor support GPU

* TRT backend support GPU FDTensor

* FDHostAllocator add FASTDEPLOY_DECL

* fix FDTensor Data

* fix FDTensor dtype

Co-authored-by: Jason <jiangjiajun@baidu.com>
This commit is contained in:
heliqi
2022-09-08 03:53:08 -05:00
committed by GitHub
parent bc8e9e4dae
commit 4d1f264d01
17 changed files with 432 additions and 153 deletions

View File

@@ -14,6 +14,8 @@
#pragma once
#include <cuda_runtime_api.h>
#include <iostream>
#include <map>
#include <string>
@@ -23,7 +25,6 @@
#include "NvOnnxParser.h"
#include "fastdeploy/backends/backend.h"
#include "fastdeploy/backends/tensorrt/utils.h"
#include <cuda_runtime_api.h>
namespace fastdeploy {
@@ -109,12 +110,12 @@ class TrtBackend : public BaseBackend {
std::map<std::string, ShapeRangeInfo> shape_range_info_;
void GetInputOutputInfo();
void AllocateBufferInDynamicShape(const std::vector<FDTensor>& inputs,
std::vector<FDTensor>* outputs);
bool CreateTrtEngineFromOnnx(const std::string& onnx_model_buffer);
bool BuildTrtEngine();
bool LoadTrtCache(const std::string& trt_engine_file);
int ShapeRangeInfoUpdated(const std::vector<FDTensor>& inputs);
void SetInputs(const std::vector<FDTensor>& inputs);
void AllocateOutputsBuffer(std::vector<FDTensor>* outputs);
};
} // namespace fastdeploy
} // namespace fastdeploy