mindspore-Ecosystem
/
mindspore
21
0
Fork 27
Code Issues Pull Requests 6 Projects Releases Wiki Activity

[MSLITE] fix new tensor api using for tensorrt

This commit is contained in:
Liu_Xuu 2021-07-20 16:55:33 +08:00
parent 4d403f5a39
commit 63313b0f32
9 changed files with 40 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mindspore/lite/src/delegate/tensorrt/op/gather_tensorrt.cc
View File

@ -53,7 +53,7 @@ int GatherTensorRT::AddInnerOp(nvinfer1::INetworkDefinition *network) {
return RET_ERROR;
}
nvinfer1::IGatherLayer *gather_layer =
network->addGather(*tensorrt_in_tensors_[0], *add_tensor /*indices*/, axis_ /*axis*/);
network->addGather(*tensorrt_in_tensors_[0], *add_tensor /* indices */, axis_ /* axis */);
if (gather_layer == nullptr) {
MS_LOG(ERROR) << "addGather failed for TensorRT.";
return RET_ERROR;

2
mindspore/lite/src/delegate/tensorrt/tensorrt_allocator.cc
View File

@ -21,7 +21,7 @@
#include "src/delegate/tensorrt/tensorrt_utils.h"
namespace mindspore::lite {
void *TensorRTAllocator::MallocDeviceMem(mindspore::MSTensor host_tensor, size_t size) {
void *TensorRTAllocator::MallocDeviceMem(const mindspore::MSTensor &host_tensor, size_t size) {
if (host_tensor == NULL) {
return nullptr;
}

2
mindspore/lite/src/delegate/tensorrt/tensorrt_allocator.h
View File

@ -33,7 +33,7 @@ class TensorRTAllocator {
~TensorRTAllocator() = default;
void *MallocDeviceMem(mindspore::MSTensor host_tensor, size_t size);
void *MallocDeviceMem(const mindspore::MSTensor &host_tensor, size_t size);
void *GetDevicePtr(const std::string &tensor_name);

2
mindspore/lite/src/delegate/tensorrt/tensorrt_delegate.cc
View File

@ -93,7 +93,6 @@ int TensorRTDelegate::Build(DelegateModel *model) {
for (KernelIter iter = model->BeginKernelIterator(); iter != model->EndKernelIterator(); iter++) {
kernel::Kernel *kernel = *iter;
auto tensorrt_op = FindTensorRTOp(kernel, model->GetPrimitive(kernel));
if (tensorrt_op != nullptr) {
// If tensorrt_ops does not equal nullptr, this kernel can be supported by delegate
if (tensorrt_ops.size() == 0) {
@ -132,7 +131,6 @@ TensorRTOp *TensorRTDelegate::FindTensorRTOp(kernel::Kernel *kernel, const schem
auto out_tensors = kernel->outputs();
auto name = kernel->name();
auto node_type = primitive->value_type();
if (op_func_lists_.find(node_type) != op_func_lists_.end()) {
return op_func_lists_[node_type](primitive, in_tensors, out_tensors, name);
} else {

2
mindspore/lite/src/delegate/tensorrt/tensorrt_runtime.h
View File

@ -26,7 +26,7 @@ using mindspore::lite::RET_OK;
namespace mindspore::lite {
class TensorRTLogger : public nvinfer1::ILogger {
void log(Severity severity, const char *msg) override {
void log(Severity severity, const char *msg) noexcept override {
if (severity == Severity::kINTERNAL_ERROR || severity == Severity::kERROR) {
MS_LOG(ERROR) << msg;
} else if (severity == Severity::kWARNING) {

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

@ -99,7 +99,7 @@ int TensorRTSubGraph::SetDeviceConfig() {
config_->setFlag(nvinfer1::BuilderFlag::kFP16);
}
// config setMaxWorkspaceSize to x MB
// config setMaxWorkspaceSize to 256 MB for max limit
config_->setMaxWorkspaceSize(256 * (1 << 20));
return RET_OK;
}
@ -248,7 +248,7 @@ int TensorRTSubGraph::Execute() {
return RET_OK;
}
nvinfer1::ITensor *TensorRTSubGraph::FindTensorRTInputs(TensorRTOp *cur_op, mindspore::MSTensor in_tensor) {
nvinfer1::ITensor *TensorRTSubGraph::FindTensorRTInputs(TensorRTOp *cur_op, const mindspore::MSTensor &in_tensor) {
for (auto input_op : cur_op->in_ops()) {
for (size_t i = 0; i < input_op->outputs().size(); i++) {
auto out_tensor = input_op->outputs().at(i);

2
mindspore/lite/src/delegate/tensorrt/tensorrt_subgraph.h
View File

@ -62,7 +62,7 @@ class TensorRTSubGraph : public kernel::Kernel {
bool SupportFP16();
static nvinfer1::ITensor *FindTensorRTInputs(TensorRTOp *cur_op, mindspore::MSTensor in_tensor);
static nvinfer1::ITensor *FindTensorRTInputs(TensorRTOp *cur_op, const mindspore::MSTensor &in_tensor);
TensorRTRuntime *runtime_{nullptr};

40
mindspore/lite/src/delegate/tensorrt/tensorrt_utils.cc
View File

@ -86,15 +86,18 @@ nvinfer1::IShuffleLayer *NCHW2NHWC(nvinfer1::INetworkDefinition *network, const
return SetTranspose(network, input, perm);
}
nvinfer1::ITensor *ConvertConstantTensor(nvinfer1::INetworkDefinition *network, mindspore::MSTensor ms_tensor) {
nvinfer1::ITensor *ConvertConstantTensor(nvinfer1::INetworkDefinition *network, const mindspore::MSTensor &ms_tensor) {
if (network == nullptr) {
MS_LOG(ERROR) << "network is null for ConvertConstantTensor";
return nullptr;
}
nvinfer1::Dims dims = ConvertCudaDims(ms_tensor.Shape());
nvinfer1::DataType data_type = ConvertDataType(ms_tensor.DataType());
nvinfer1::Weights weights{data_type, ms_tensor.MutableData(), ms_tensor.ElementNum()};
if (ms_tensor.Data() == nullptr) {
MS_LOG(ERROR) << "ConvertConstantTensor from a MSTensor with nullptr data";
return nullptr;
}
nvinfer1::Weights weights{data_type, ms_tensor.Data().get(), ms_tensor.ElementNum()};
nvinfer1::IConstantLayer *constant_tensor = network->addConstant(dims, weights);
if (constant_tensor == nullptr) {
MS_LOG(ERROR) << "create constant_tensor failed.";
@ -140,8 +143,8 @@ ActivationParams ConvertActivationType(schema::ActivationType activation_type) {
return action_param;
}
nvinfer1::ITensor *ConvertTensorWithExpandDims(nvinfer1::INetworkDefinition *network, mindspore::MSTensor ms_tensor,
size_t expand_shape_size) {
nvinfer1::ITensor *ConvertTensorWithExpandDims(nvinfer1::INetworkDefinition *network,
const mindspore::MSTensor &ms_tensor, size_t expand_shape_size) {
if (network == nullptr) {
MS_LOG(ERROR) << "network is null for ConvertConstantTensor";
return nullptr;
@ -158,8 +161,11 @@ nvinfer1::ITensor *ConvertTensorWithExpandDims(nvinfer1::INetworkDefinition *net
}
nvinfer1::Dims dims = ConvertCudaDims(shape);
nvinfer1::DataType data_type = ConvertDataType(ms_tensor.DataType());
nvinfer1::Weights weights{data_type, ms_tensor.MutableData(), ms_tensor.ElementNum()};
if (ms_tensor.Data() == nullptr) {
MS_LOG(ERROR) << "ConvertTensorWithExpandDims from a MSTensor with nullptr data";
return nullptr;
}
nvinfer1::Weights weights{data_type, ms_tensor.Data().get(), ms_tensor.ElementNum()};
nvinfer1::IConstantLayer *constant_tensor = network->addConstant(dims, weights);
if (constant_tensor == nullptr) {
MS_LOG(ERROR) << "create constant_tensor failed.";
@ -169,7 +175,7 @@ nvinfer1::ITensor *ConvertTensorWithExpandDims(nvinfer1::INetworkDefinition *net
constant_tensor->setName(name.c_str());
return constant_tensor->getOutput(0);
}
nvinfer1::Weights TransposeWeight(mindspore::MSTensor ms_tensor, float **pack_weight) {
nvinfer1::Weights TransposeWeight(const mindspore::MSTensor &ms_tensor, float **pack_weight) {
// usage notice: malloc addr saved to pack_weight, save pack_weight ptr and free it when deconstruct
nvinfer1::Weights weights{};
weights.count = ms_tensor.ElementNum();
@ -178,7 +184,15 @@ nvinfer1::Weights TransposeWeight(mindspore::MSTensor ms_tensor, float **pack_we
}
weights.type = nvinfer1::DataType::kFLOAT;
auto weight_shape = ms_tensor.Shape();
float *src_val = reinterpret_cast<float *>(ms_tensor.MutableData());
const void *src_ptr = ms_tensor.Data().get();
const float *src_val;
if (src_ptr == nullptr) {
src_val = nullptr;
MS_LOG(ERROR) << "TransposeWeight from a MSTensor with nullptr data";
return weights;
}
src_val = reinterpret_cast<const float *>(src_ptr);
float *pack_weight_tmp = reinterpret_cast<float *>(malloc(ms_tensor.ElementNum() * sizeof(float)));
if (pack_weight_tmp == nullptr) {
MS_LOG(ERROR) << "Malloc buffer failed.";
@ -191,12 +205,14 @@ nvinfer1::Weights TransposeWeight(mindspore::MSTensor ms_tensor, float **pack_we
return weights;
}
nvinfer1::Weights ConvertWeight(mindspore::MSTensor ms_tensor) {
nvinfer1::Weights ConvertWeight(const mindspore::MSTensor &ms_tensor) {
nvinfer1::Weights weights{};
weights.type = ConvertDataType(ms_tensor.DataType());
weights.values = ms_tensor.MutableData();
weights.values = ms_tensor.Data().get();
weights.count = ms_tensor.ElementNum();
if (weights.values == nullptr) {
MS_LOG(ERROR) << "ConvertWeight from a MSTensor with nullptr data";
}
return weights;
}
} // namespace mindspore::lite

10
mindspore/lite/src/delegate/tensorrt/tensorrt_utils.h
View File

@ -46,16 +46,16 @@ nvinfer1::IShuffleLayer *NCHW2NHWC(nvinfer1::INetworkDefinition *network, const
ActivationParams ConvertActivationType(schema::ActivationType activation_type);
nvinfer1::ITensor *ConvertConstantTensor(nvinfer1::INetworkDefinition *network, mindspore::MSTensor ms_tensor);
nvinfer1::ITensor *ConvertConstantTensor(nvinfer1::INetworkDefinition *network, const mindspore::MSTensor &ms_tensor);
nvinfer1::ITensor *ConvertTensorWithExpandDims(nvinfer1::INetworkDefinition *network, mindspore::MSTensor ms_tensor,
size_t expand_shape_size);
nvinfer1::ITensor *ConvertTensorWithExpandDims(nvinfer1::INetworkDefinition *network,
const mindspore::MSTensor &ms_tensor, size_t expand_shape_size);
nvinfer1::ITensor *ConvertScalarToITensor(nvinfer1::INetworkDefinition *network, size_t shape_size, void *value);
nvinfer1::Weights TransposeWeight(mindspore::MSTensor ms_tensor, float **pack_weight);
nvinfer1::Weights TransposeWeight(const mindspore::MSTensor &ms_tensor, float **pack_weight);
nvinfer1::Weights ConvertWeight(mindspore::MSTensor ms_tensor);
nvinfer1::Weights ConvertWeight(const mindspore::MSTensor &ms_tensor);
} // namespace mindspore::lite
#endif // MINDSPORE_LITE_SRC_RUNTIME_DELEGATE_TENSORRT_UTILS_H_