p19647583
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!27725 [MSLITE] optimize input profile in tensorrt 

Merge pull request !27725 from Liu_Xuu/trt_1215_opt
This commit is contained in:
i-robot 2021-12-17 02:02:24 +00:00 committed by Gitee
parent 4daacb6f7f 6a8533afa1
commit 2510e9f3d3
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
mindspore/lite/src/delegate/tensorrt/tensorrt_subgraph.cc
View File

@ -108,8 +108,8 @@ int TensorRTSubGraph::SetDeviceConfig(cudaStream_t stream) {
config_->setProfileStream(stream);
// config setMaxWorkspaceSize to 1024 MB for max limit
config_->setMaxWorkspaceSize(1024 * (1 << 20));
// config setMaxWorkspaceSize to 1152 MB for max limit
config_->setMaxWorkspaceSize(1152 * (1 << 20));
return RET_OK;
}
@ -194,6 +194,9 @@ nvinfer1::ITensor *TensorRTSubGraph::SetTensorRTNetworkInput(const mindspore::MS
return nullptr;
}
nvinfer1::Dims input_dims_opt = ConvertCudaDims(in_tensor.Shape());
if (input_batchsize_index_ != -1) {
input_dims_opt.d[input_batchsize_index_] = 1;
}
if (!profile_->setDimensions(in_tensor.Name().c_str(), nvinfer1::OptProfileSelector::kOPT, input_dims_opt)) {
MS_LOG(ERROR) << "setDimensions of kOPT failed for " << in_tensor.Name();
return nullptr;