|
|
|
@ -220,6 +220,29 @@ class TensorRTEngine {
|
|
|
|
|
void Deserialize(const std::string& engine_serialized_data) {
|
|
|
|
|
freshDeviceId();
|
|
|
|
|
infer_ptr<nvinfer1::IRuntime> runtime(createInferRuntime(&logger_));
|
|
|
|
|
|
|
|
|
|
if (use_dla_) {
|
|
|
|
|
if (precision_ != AnalysisConfig::Precision::kInt8 &&
|
|
|
|
|
precision_ != AnalysisConfig::Precision::kHalf) {
|
|
|
|
|
LOG(WARNING) << "TensorRT DLA must be used with int8 or fp16, but you "
|
|
|
|
|
"set float32, so DLA is not used.";
|
|
|
|
|
} else if (runtime->getNbDLACores() == 0) {
|
|
|
|
|
LOG(WARNING)
|
|
|
|
|
<< "TensorRT DLA is set by config, but your device does not have "
|
|
|
|
|
"DLA, so DLA is not used.";
|
|
|
|
|
} else {
|
|
|
|
|
if (dla_core_ < 0 || dla_core_ >= runtime->getNbDLACores()) {
|
|
|
|
|
dla_core_ = 0;
|
|
|
|
|
LOG(WARNING) << "Invalid DLACore, must be 0 < DLACore < "
|
|
|
|
|
<< runtime->getNbDLACores() << ", but got " << dla_core_
|
|
|
|
|
<< ", so use use 0 as default.";
|
|
|
|
|
}
|
|
|
|
|
runtime->setDLACore(dla_core_);
|
|
|
|
|
LOG(INFO) << "TensorRT DLA enabled in Deserialize(), DLACore "
|
|
|
|
|
<< dla_core_;
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
if (with_dynamic_shape_) {
|
|
|
|
|
#if IS_TRT_VERSION_GE(6000)
|
|
|
|
|
infer_engine_.reset(runtime->deserializeCudaEngine(
|
|
|
|
@ -287,6 +310,8 @@ class TensorRTEngine {
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
void SetUseOSS(bool use_oss) { use_oss_ = use_oss; }
|
|
|
|
|
void SetUseDLA(bool use_dla) { use_dla_ = use_dla; }
|
|
|
|
|
void SetDLACore(int dla_core) { dla_core_ = dla_core; }
|
|
|
|
|
void SetWithErnie(bool with_ernie) { with_ernie_ = with_ernie; }
|
|
|
|
|
|
|
|
|
|
void ClearWeights() {
|
|
|
|
@ -316,8 +341,8 @@ class TensorRTEngine {
|
|
|
|
|
ShapeMapType min_input_shape() { return min_input_shape_; }
|
|
|
|
|
ShapeMapType max_input_shape() { return max_input_shape_; }
|
|
|
|
|
ShapeMapType optim_input_shape() { return optim_input_shape_; }
|
|
|
|
|
bool use_oss() { return use_oss_; };
|
|
|
|
|
bool with_ernie() { return with_ernie_; };
|
|
|
|
|
bool use_oss() { return use_oss_; }
|
|
|
|
|
bool with_ernie() { return with_ernie_; }
|
|
|
|
|
bool disable_trt_plugin_fp16() { return disable_trt_plugin_fp16_; }
|
|
|
|
|
bool with_dynamic_shape() { return with_dynamic_shape_; }
|
|
|
|
|
|
|
|
|
@ -354,6 +379,8 @@ class TensorRTEngine {
|
|
|
|
|
ShapeMapType optim_input_shape_;
|
|
|
|
|
bool disable_trt_plugin_fp16_{false};
|
|
|
|
|
bool use_oss_{false};
|
|
|
|
|
bool use_dla_{false};
|
|
|
|
|
int dla_core_{0};
|
|
|
|
|
bool with_ernie_{false};
|
|
|
|
|
nvinfer1::ILogger& logger_;
|
|
|
|
|
|
|
|
|
|