Skip to content

Commit

Permalink
fix(tensorrt): read onnx model to find topk
Browse files Browse the repository at this point in the history
  • Loading branch information
@Bycob @mergify
Bycob authored and mergify[bot] committed Dec 10, 2021
1 parent a03eb87 commit 5cce134
Show file tree
Hide file tree
Showing 5 changed files with 244 additions and 147 deletions.
14 changes: 14 additions & 0 deletions CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -990,6 +990,20 @@ if (USE_TENSORRT)

# Use our tensorrt-oss header versions first
list(INSERT TENSORRT_INC_DIR 0 ${CMAKE_BINARY_DIR}/tensorrt-oss/src/tensorrt-oss/include)

# onnx.pb.h
set(ONNX_PROTO ${CMAKE_BINARY_DIR}/tensorrt-oss/src/tensorrt-oss/parsers/onnx/third_party/onnx/onnx/onnx.proto)
add_custom_command(
OUTPUT ${CMAKE_BINARY_DIR}/src/onnx.pb.h ${CMAKE_BINARY_DIR}/src/onnx.pb.cc
COMMAND LD_LIBRARY_PATH=${PROTOBUF_LIB_DIR} ${PROTOBUF_PROTOC} --proto_path=${CMAKE_BINARY_DIR}/tensorrt-oss/src/tensorrt-oss/parsers/onnx/third_party/onnx/onnx/ --cpp_out=${CMAKE_BINARY_DIR}/src ${ONNX_PROTO}
COMMENT Generating onnx.pb.h and onnx.pb.cc
WORKING_DIRECTORY ${CMAKE_BINARY_DIR}
DEPENDS protobuf tensorrt-oss
)
add_custom_target(
onnx_pb_h
DEPENDS ${CMAKE_BINARY_DIR}/src/onnx.pb.h ${CMAKE_BINARY_DIR}/src/onnx.pb.cc
)
endif()

set(OATPP_VERSION "1.2.5")
Expand Down
9 changes: 8 additions & 1 deletion src/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,11 @@ if (USE_TORCH)
endif()

if (USE_TENSORRT)
set_source_files_properties(
${CMAKE_BINARY_DIR}/src/onnx.pb.cc
${CMAKE_BINARY_DIR}/src/onnx.pb.h
PROPERTIES GENERATED TRUE)

list(APPEND ddetect_SOURCES
backends/tensorrt/tensorrtlib.h
backends/tensorrt/tensorrtlib.cc
Expand All @@ -121,7 +126,8 @@ if (USE_TENSORRT)
backends/tensorrt/tensorrtinputconns.cc
backends/tensorrt/tensorrtmodel.h
backends/tensorrt/tensorrtmodel.cc
)
${CMAKE_BINARY_DIR}/src/onnx.pb.cc
)
endif()


Expand All @@ -140,6 +146,7 @@ if (USE_CAFFE)
add_dependencies(ddetect caffe_dd)
endif()
if (USE_TENSORRT)
add_dependencies(ddetect onnx_pb_h)
if (USE_CAFFE2)
add_dependencies(ddetect pytorch)
endif()
Expand Down
310 changes: 182 additions & 128 deletions src/backends/tensorrt/protoUtils.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,11 +27,14 @@
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <iostream>
#include <fstream>

#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-parameter"
#pragma GCC diagnostic ignored "-Wmaybe-uninitialized"
#include "src/caffe.pb.h"
#include "src/onnx.pb.h"
#pragma GCC diagnostic pop

using google::protobuf::io::CodedInputStream;
Expand All @@ -43,144 +46,186 @@ using google::protobuf::io::ZeroCopyOutputStream;

namespace dd
{

bool findInputDimensions(const std::string &source, int &width, int &height)
namespace caffe_proto
{
caffe::NetParameter net;
if (!TRTReadProtoFromTextFile(source.c_str(), &net))
return false;
bool findInputDimensions(const std::string &source, int &width,
int &height)
{
caffe::NetParameter net;
if (!TRTReadProtoFromTextFile(source.c_str(), &net))
return false;

for (int i = 0; i < net.layer_size(); ++i)
{
caffe::LayerParameter lparam = net.layer(i);
if (lparam.type() == "MemoryData")
{
width = lparam.memory_data_param().width();
height = lparam.memory_data_param().height();
}
break; // don't go further than first layer.
}

return true;
}
for (int i = 0; i < net.layer_size(); ++i)
{
caffe::LayerParameter lparam = net.layer(i);
if (lparam.type() == "MemoryData")
{
width = lparam.memory_data_param().width();
height = lparam.memory_data_param().height();
}
break; // don't go further than first layer.
}

int findNClasses(const std::string source, bool bbox)
{
caffe::NetParameter net;
if (!TRTReadProtoFromTextFile(source.c_str(), &net))
return -1;
int nlayers = net.layer_size();
if (bbox)
{
for (int i = nlayers - 1; i >= 0; --i)
{
caffe::LayerParameter lparam = net.layer(i);
if (lparam.type() == "DetectionOutput")
return lparam.detection_output_param().num_classes();
}
}
for (int i = nlayers - 1; i >= 0; --i)
{
caffe::LayerParameter lparam = net.layer(i);
if (lparam.type() == "InnerProduct")
return lparam.inner_product_param().num_output();

if (net.layer(0).name() == "squeezenet"
&& lparam.type() == "Convolution")
return lparam.convolution_param().num_output();
}
return -1;
}
return true;
}

int findTopK(const std::string source)
{
caffe::NetParameter net;
if (!TRTReadProtoFromTextFile(source.c_str(), &net))
int findNClasses(const std::string source, bool bbox)
{
caffe::NetParameter net;
if (!TRTReadProtoFromTextFile(source.c_str(), &net))
return -1;
int nlayers = net.layer_size();
if (bbox)
{
for (int i = nlayers - 1; i >= 0; --i)
{
caffe::LayerParameter lparam = net.layer(i);
if (lparam.type() == "DetectionOutput")
return lparam.detection_output_param().num_classes();
}
}
for (int i = nlayers - 1; i >= 0; --i)
{
caffe::LayerParameter lparam = net.layer(i);
if (lparam.type() == "InnerProduct")
return lparam.inner_product_param().num_output();

if (net.layer(0).name() == "squeezenet"
&& lparam.type() == "Convolution")
return lparam.convolution_param().num_output();
}
return -1;
int nlayers = net.layer_size();
for (int i = nlayers - 1; i >= 0; --i)
{
caffe::LayerParameter lparam = net.layer(i);
if (lparam.type() == "DetectionOutput")
return lparam.detection_output_param().nms_param().top_k();
}
return -1;
}
}

bool isRefinedet(const std::string source)
{
caffe::NetParameter net;
if (!TRTReadProtoFromTextFile(source.c_str(), &net))
int findTopK(const std::string source)
{
caffe::NetParameter net;
if (!TRTReadProtoFromTextFile(source.c_str(), &net))
return -1;
int nlayers = net.layer_size();
for (int i = nlayers - 1; i >= 0; --i)
{
caffe::LayerParameter lparam = net.layer(i);
if (lparam.type() == "DetectionOutput")
return lparam.detection_output_param().nms_param().top_k();
}
return -1;
int nlayers = net.layer_size();

for (int i = nlayers - 1; i >= 0; --i)
{
caffe::LayerParameter lparam = net.layer(i);
if (lparam.type() == "DetectionOutput" && lparam.bottom_size() > 3)
return true;
}
return false;
}

bool isRefinedet(const std::string source)
{
caffe::NetParameter net;
if (!TRTReadProtoFromTextFile(source.c_str(), &net))
return -1;
int nlayers = net.layer_size();

for (int i = nlayers - 1; i >= 0; --i)
{
caffe::LayerParameter lparam = net.layer(i);
if (lparam.type() == "DetectionOutput" && lparam.bottom_size() > 3)
return true;
}
return false;
}

int fixProto(const std::string dest, const std::string source)
{
caffe::NetParameter source_net;
caffe::NetParameter dest_net;
if (!TRTReadProtoFromTextFile(source.c_str(), &source_net))
return 1;

dest_net.set_name(source_net.name());
int nlayers = source_net.layer_size();

for (int i = 0; i < nlayers; ++i)
{
caffe::LayerParameter lparam = source_net.layer(i);
if (lparam.type() == "MemoryData")
{
dest_net.add_input(lparam.top(0));
caffe::BlobShape *is = dest_net.add_input_shape();
is->add_dim(lparam.memory_data_param().batch_size());
is->add_dim(lparam.memory_data_param().channels());
is->add_dim(lparam.memory_data_param().height());
is->add_dim(lparam.memory_data_param().width());
}
else if (lparam.type() == "Flatten")
{
caffe::LayerParameter *rparam = dest_net.add_layer();
rparam->set_name(lparam.name());
rparam->set_type("Reshape");
rparam->add_bottom(lparam.bottom(0));
rparam->add_top(lparam.top(0));
int faxis = lparam.flatten_param().axis();
caffe::ReshapeParameter *rp = rparam->mutable_reshape_param();
caffe::BlobShape *bs = rp->mutable_shape();
for (int i = 0; i < faxis; ++i)
bs->add_dim(0);
bs->add_dim(-1);
for (int i = faxis + 1; i < 4; ++i)
bs->add_dim(1);
}
else if (lparam.type() == "DetectionOutput")
{
caffe::LayerParameter *dlparam = dest_net.add_layer();
caffe::NonMaximumSuppressionParameter *nmsp
= lparam.mutable_detection_output_param()
->mutable_nms_param();
nmsp->clear_soft_nms();
nmsp->clear_theta();
*dlparam = lparam;
dlparam->add_top("keep_count");
}
else
{
caffe::LayerParameter *dlparam = dest_net.add_layer();
*dlparam = lparam;
}
}

if (!TRTWriteProtoToTextFile(dest_net, dest.c_str()))
return 2;
return 0;
}
}

int fixProto(const std::string dest, const std::string source)
namespace onnx_proto
{
caffe::NetParameter source_net;
caffe::NetParameter dest_net;
if (!TRTReadProtoFromTextFile(source.c_str(), &source_net))
return 1;

dest_net.set_name(source_net.name());
int nlayers = source_net.layer_size();

for (int i = 0; i < nlayers; ++i)
{
caffe::LayerParameter lparam = source_net.layer(i);
if (lparam.type() == "MemoryData")
{
dest_net.add_input(lparam.top(0));
caffe::BlobShape *is = dest_net.add_input_shape();
is->add_dim(lparam.memory_data_param().batch_size());
is->add_dim(lparam.memory_data_param().channels());
is->add_dim(lparam.memory_data_param().height());
is->add_dim(lparam.memory_data_param().width());
}
else if (lparam.type() == "Flatten")
{
caffe::LayerParameter *rparam = dest_net.add_layer();
rparam->set_name(lparam.name());
rparam->set_type("Reshape");
rparam->add_bottom(lparam.bottom(0));
rparam->add_top(lparam.top(0));
int faxis = lparam.flatten_param().axis();
caffe::ReshapeParameter *rp = rparam->mutable_reshape_param();
caffe::BlobShape *bs = rp->mutable_shape();
for (int i = 0; i < faxis; ++i)
bs->add_dim(0);
bs->add_dim(-1);
for (int i = faxis + 1; i < 4; ++i)
bs->add_dim(1);
}
else if (lparam.type() == "DetectionOutput")
{
caffe::LayerParameter *dlparam = dest_net.add_layer();
caffe::NonMaximumSuppressionParameter *nmsp
= lparam.mutable_detection_output_param()->mutable_nms_param();
nmsp->clear_soft_nms();
nmsp->clear_theta();
*dlparam = lparam;
dlparam->add_top("keep_count");
}
else
{
caffe::LayerParameter *dlparam = dest_net.add_layer();
*dlparam = lparam;
}
}

if (!TRTWriteProtoToTextFile(dest_net, dest.c_str()))
return 2;
return 0;
// https://github.com/onnx/onnx/blob/master/onnx/onnx.proto
int findTopK(const std::string &source, const std::string &out_name)
{
onnx::ModelProto net;
if (!TRTReadProtoFromBinaryFile(source.c_str(), &net))
{
return -1;
}

int node_count = net.graph().output_size();

for (int i = 0; i < node_count; ++i)
{
const auto &out = net.graph().output(i);
std::string name = out.has_name() ? out.name() : "";

if (name == out_name)
{
if (out.type().has_tensor_type())
{
auto &shape = out.type().tensor_type().shape();
switch (shape.dim_size())
{
case 2:
return shape.dim(0).dim_value();
case 3:
case 4:
return shape.dim(1).dim_value();
}
}
}
}
return -1;
}
}

bool TRTReadProtoFromTextFile(const char *filename,
Expand Down Expand Up @@ -209,4 +254,13 @@ namespace dd
return success;
}

bool TRTReadProtoFromBinaryFile(const char *filename,
google::protobuf::Message *proto)
{
std::ifstream input(filename, std::ios::in | std::ios::binary);
if (!input.good())
return false;
bool success = proto->ParseFromIstream(&input);
return success;
}
}
Loading

0 comments on commit 5cce134

Please sign in to comment.