refactor tensorrt node teller (#15181)
Yan Chunwei
6 years ago
committed by
GitHub
parent
c8f101e5da
commit
6ccf8685f7
@ -1,9 +1,13 @@
|
||||
cc_library(subgraph_detector SRCS subgraph_detector.cc DEPS proto_desc)
|
||||
cc_library(tensorrt_subgraph_pass SRCS tensorrt_subgraph_pass.cc DEPS subgraph_detector)
|
||||
set(analysis_deps ${analysis_deps}
|
||||
subgraph_detector tensorrt_subgraph_pass
|
||||
CACHE INTERNAL "")
|
||||
|
||||
set(pass_file ${PADDLE_BINARY_DIR}/paddle/fluid/inference/api/paddle_inference_pass.h)
|
||||
file(APPEND ${pass_file} "USE_PASS(tensorrt_subgraph_pass);\n")
|
||||
set(INFER_IR_PASSES ${INFER_IR_PASSES} tensorrt_subgraph_pass CACHE INTERNAL "")
|
||||
if (TENSORRT_FOUND)
|
||||
cc_library(tensorrt_subgraph_pass SRCS tensorrt_subgraph_pass.cc DEPS subgraph_detector tensorrt_op_teller)
|
||||
|
||||
set(analysis_deps ${analysis_deps}
|
||||
subgraph_detector tensorrt_subgraph_pass
|
||||
CACHE INTERNAL "")
|
||||
|
||||
set(pass_file ${PADDLE_BINARY_DIR}/paddle/fluid/inference/api/paddle_inference_pass.h)
|
||||
file(APPEND ${pass_file} "USE_PASS(tensorrt_subgraph_pass);\n")
|
||||
set(INFER_IR_PASSES ${INFER_IR_PASSES} tensorrt_subgraph_pass CACHE INTERNAL "")
|
||||
endif()
|
||||
|
||||
@ -0,0 +1,49 @@
|
||||
// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include "paddle/fluid/inference/tensorrt/op_teller.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace inference {
|
||||
namespace tensorrt {
|
||||
|
||||
// Just tell by the op_types.
|
||||
struct SimpleOpTypeSetTeller : public Teller {
|
||||
SimpleOpTypeSetTeller() {}
|
||||
|
||||
bool operator()(const std::string& op_type,
|
||||
const framework::OpDesc& desc) override {
|
||||
return teller_set.count(op_type);
|
||||
}
|
||||
|
||||
private:
|
||||
std::unordered_set<std::string> teller_set{
|
||||
{"mul", "conv2d", "pool2d", "relu", "softmax", "sigmoid",
|
||||
"depthwise_conv2d", "batch_norm", "concat", "tanh", "pad",
|
||||
"elementwise_add", "elementwise_mul", "dropout", "split", "prelu",
|
||||
"conv2d_transpose", "leaky_relu"}};
|
||||
};
|
||||
|
||||
bool OpTeller::Tell(const std::string& op_type, const framework::OpDesc& desc) {
|
||||
for (auto& teller : tellers_) {
|
||||
if ((*teller)(op_type, desc)) return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
OpTeller::OpTeller() { tellers_.emplace_back(new SimpleOpTypeSetTeller); }
|
||||
|
||||
} // namespace tensorrt
|
||||
} // namespace inference
|
||||
} // namespace paddle
|
||||
@ -0,0 +1,68 @@
|
||||
// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#pragma once
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include "paddle/fluid/framework/op_desc.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace inference {
|
||||
namespace tensorrt {
|
||||
|
||||
/*
|
||||
* Single Op teller definition.
|
||||
* One can override this and define a more complex tell logic, considerring more
|
||||
* issues such as op_desc.
|
||||
*/
|
||||
struct Teller {
|
||||
virtual bool operator()(const std::string& op_type,
|
||||
const framework::OpDesc& desc) = 0;
|
||||
|
||||
virtual ~Teller() = default;
|
||||
};
|
||||
/*
|
||||
* A real example:
|
||||
*
|
||||
* struct SomeTeller : public Teller {
|
||||
* bool operator()(const std::string& op_type,
|
||||
* const framework::OpDesc& desc) override {
|
||||
* return op_type == "fc" && desc.Inputs().size() == 2;
|
||||
* }
|
||||
*};
|
||||
*/
|
||||
|
||||
/*
|
||||
* class OpTeller helps to tell whether a fluid
|
||||
* operator can be transformed to a TensorRT layer.
|
||||
*/
|
||||
class OpTeller {
|
||||
public:
|
||||
static OpTeller& Global() {
|
||||
static std::unique_ptr<OpTeller> x(new OpTeller);
|
||||
return *x;
|
||||
}
|
||||
|
||||
bool Tell(const std::string& op_type, const framework::OpDesc& desc);
|
||||
|
||||
private:
|
||||
OpTeller();
|
||||
|
||||
private:
|
||||
std::vector<std::unique_ptr<Teller>> tellers_;
|
||||
};
|
||||
|
||||
} // namespace tensorrt
|
||||
} // namespace inference
|
||||
} // namespace paddle
|
||||
Loading…
Reference in new issue