nhzlx
7 years ago
parent
2d7134bc37
commit
d38fd6a0fc
@ -0,0 +1,73 @@
|
||||
/* Copyright (c) 2018 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/convert/op_converter.h"
|
||||
#include "paddle/fluid/inference/tensorrt/plugin/split_op_plugin.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace inference {
|
||||
namespace tensorrt {
|
||||
|
||||
/*
|
||||
* SplitOp.
|
||||
*/
|
||||
class SplitOpConverter : public OpConverter {
|
||||
public:
|
||||
void operator()(const framework::proto::OpDesc& op,
|
||||
const framework::Scope& scope, bool test_mode) override {
|
||||
VLOG(40) << "convert a fluid split op to tensorrt split layer";
|
||||
|
||||
framework::OpDesc op_desc(op, nullptr);
|
||||
// Declare inputs
|
||||
auto* input = engine_->GetITensor(op_desc.Input("X")[0]);
|
||||
auto input_dims = input->getDimensions();
|
||||
int input_num = op_desc.Input("X").size();
|
||||
size_t output_num = op_desc.Output("Out").size();
|
||||
|
||||
PADDLE_ENFORCE(input_num == 1);
|
||||
int axis = boost::get<int>(op_desc.GetAttr("axis"));
|
||||
std::vector<int> output_lengths =
|
||||
boost::get<std::vector<int>>(op_desc.GetAttr("sections"));
|
||||
PADDLE_ENFORCE(axis != 0);
|
||||
if (axis < 0) {
|
||||
axis += input_dims.nbDims;
|
||||
} else {
|
||||
axis -= 1;
|
||||
}
|
||||
|
||||
PADDLE_ENFORCE(output_lengths.size() == output_num);
|
||||
|
||||
SplitPlugin* plugin = new SplitPlugin(axis, output_lengths);
|
||||
nvinfer1::IPluginLayer* layer =
|
||||
engine_->addPlugin(&input, input_num, plugin);
|
||||
|
||||
std::string layer_name = "split (Output: ";
|
||||
for (size_t i = 0; i < output_num; i++) {
|
||||
auto output_name = op_desc.Output("Out")[i];
|
||||
layer->getOutput(i)->setName(output_name.c_str());
|
||||
engine_->SetITensor(output_name, layer->getOutput(i));
|
||||
layer_name += output_name;
|
||||
if (test_mode) {
|
||||
engine_->DeclareOutput(output_name);
|
||||
}
|
||||
}
|
||||
layer->setName((layer_name + ")").c_str());
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace tensorrt
|
||||
} // namespace inference
|
||||
} // namespace paddle
|
||||
|
||||
REGISTER_TRT_OP_CONVERTER(split, SplitOpConverter);
|
||||
@ -0,0 +1,53 @@
|
||||
/* Copyright (c) 2018 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 <gtest/gtest.h>
|
||||
#include "paddle/fluid/inference/tensorrt/convert/op_converter.h"
|
||||
#include "paddle/fluid/inference/tensorrt/convert/ut_helper.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace inference {
|
||||
namespace tensorrt {
|
||||
|
||||
TEST(split_op, test) {
|
||||
std::unordered_set<std::string> parameters({""});
|
||||
framework::Scope scope;
|
||||
TRTConvertValidation validator(10, parameters, scope, 1000);
|
||||
validator.DeclInputVar("split_input", nvinfer1::DimsCHW(3, 2, 2));
|
||||
validator.DeclOutputVar("split_out1", nvinfer1::DimsCHW(2, 2, 2));
|
||||
validator.DeclOutputVar("split_out2", nvinfer1::DimsCHW(1, 2, 2));
|
||||
|
||||
// Prepare Op description
|
||||
framework::OpDesc desc;
|
||||
desc.SetType("split");
|
||||
desc.SetInput("X", {"split_input"});
|
||||
desc.SetOutput("Out", {"split_out1", "split_out2"});
|
||||
|
||||
int num = 0;
|
||||
int axis = 1;
|
||||
std::vector<int> output_lengths = {2, 1};
|
||||
desc.SetAttr("axis", axis);
|
||||
desc.SetAttr("num", num);
|
||||
desc.SetAttr("sections", output_lengths);
|
||||
|
||||
validator.SetOp(*desc.Proto());
|
||||
|
||||
validator.Execute(1);
|
||||
}
|
||||
|
||||
} // namespace tensorrt
|
||||
} // namespace inference
|
||||
} // namespace paddle
|
||||
|
||||
USE_OP(split);
|
||||
@ -1,2 +1 @@
|
||||
nv_library(tensorrt_plugin SRCS plugin_factory.cc plugin_utils.cc
|
||||
trt_plugin.cc split_op_plugin.cu DEPS enforce)
|
||||
nv_library(tensorrt_plugin SRCS trt_plugin.cc split_op_plugin.cu DEPS enforce)
|
||||
|
||||
@ -1,64 +0,0 @@
|
||||
/* Copyright (c) 2018 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/plugin/plugin_factory.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace inference {
|
||||
namespace tensorrt {
|
||||
|
||||
PluginTensorRT* PluginFactoryTensorRT::createPlugin(const char* layer_name,
|
||||
const void* serial_data,
|
||||
size_t serial_length) {
|
||||
size_t parsed_byte = 0;
|
||||
std::string encoded_op_name =
|
||||
ExtractOpName(serial_data, serial_length, &parsed_byte);
|
||||
|
||||
if (!IsPlugin(encoded_op_name)) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
auto plugin_ptr =
|
||||
plugin_registry_[encoded_op_name].first(serial_data, serial_length);
|
||||
owned_plugins_.emplace_back(plugin_ptr);
|
||||
|
||||
return plugin_ptr;
|
||||
}
|
||||
|
||||
PluginTensorRT* PluginFactoryTensorRT::CreatePlugin(
|
||||
const std::string& op_name) {
|
||||
if (!IsPlugin(op_name)) return nullptr;
|
||||
|
||||
auto plugin_ptr = plugin_registry_[op_name].second();
|
||||
owned_plugins_.emplace_back(plugin_ptr);
|
||||
|
||||
return plugin_ptr;
|
||||
}
|
||||
|
||||
bool PluginFactoryTensorRT::RegisterPlugin(
|
||||
const std::string& op_name, PluginDeserializeFunc deserialize_func,
|
||||
PluginConstructFunc construct_func) {
|
||||
if (IsPlugin(op_name)) return false;
|
||||
|
||||
auto ret = plugin_registry_.emplace(
|
||||
op_name, std::make_pair(deserialize_func, construct_func));
|
||||
|
||||
return ret.second;
|
||||
}
|
||||
|
||||
void PluginFactoryTensorRT::DestroyPlugins() { owned_plugins_.clear(); }
|
||||
|
||||
} // namespace tensorrt
|
||||
} // namespace inference
|
||||
} // namespace paddle
|
||||
@ -1,91 +0,0 @@
|
||||
/* Copyright (c) 2018 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 <memory>
|
||||
#include <unordered_map>
|
||||
|
||||
#include "NvInfer.h"
|
||||
#include "paddle/fluid/inference/tensorrt/plugin/plugin_utils.h"
|
||||
#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
|
||||
#include "paddle/fluid/platform/enforce.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace inference {
|
||||
namespace tensorrt {
|
||||
|
||||
class PluginFactoryTensorRT : public nvinfer1::IPluginFactory {
|
||||
public:
|
||||
static PluginFactoryTensorRT* GetInstance() {
|
||||
static PluginFactoryTensorRT* factory_instance =
|
||||
new PluginFactoryTensorRT();
|
||||
return factory_instance;
|
||||
}
|
||||
|
||||
// Deserialization method
|
||||
PluginTensorRT* createPlugin(const char* layer_name, const void* serial_data,
|
||||
size_t serial_length) override;
|
||||
|
||||
// Plugin construction, PluginFactoryTensorRT owns the plugin.
|
||||
PluginTensorRT* CreatePlugin(const std::string& op_name);
|
||||
|
||||
bool RegisterPlugin(const std::string& op_name,
|
||||
PluginDeserializeFunc deserialize_func,
|
||||
PluginConstructFunc construct_func);
|
||||
|
||||
bool IsPlugin(const std::string& op_name) {
|
||||
return plugin_registry_.find(op_name) != plugin_registry_.end();
|
||||
}
|
||||
|
||||
size_t CountOwnedPlugins() { return owned_plugins_.size(); }
|
||||
|
||||
void DestroyPlugins();
|
||||
|
||||
protected:
|
||||
std::unordered_map<std::string,
|
||||
std::pair<PluginDeserializeFunc, PluginConstructFunc>>
|
||||
plugin_registry_;
|
||||
std::vector<std::unique_ptr<PluginTensorRT>> owned_plugins_;
|
||||
};
|
||||
|
||||
class TrtPluginRegistrar {
|
||||
public:
|
||||
TrtPluginRegistrar(const std::string& name,
|
||||
PluginDeserializeFunc deserialize_func,
|
||||
PluginConstructFunc construct_func) {
|
||||
auto factory = PluginFactoryTensorRT::GetInstance();
|
||||
// platform::PADDLE_ENFORCE(factory->RegisterPlugin(name, deserialize_func,
|
||||
// construct_func), "Falied to register plugin [%s]", name);
|
||||
// platform::PADDLE_ENFORCE(factory->RegisterPlugin(name, deserialize_func,
|
||||
// construct_func));
|
||||
factory->RegisterPlugin(name, deserialize_func, construct_func);
|
||||
}
|
||||
};
|
||||
|
||||
#define REGISTER_TRT_PLUGIN(name, deserialize_func, construct_func) \
|
||||
REGISTER_TRT_PLUGIN_UNIQ_HELPER(__COUNTER__, name, deserialize_func, \
|
||||
construct_func)
|
||||
#define REGISTER_TRT_PLUGIN_UNIQ_HELPER(ctr, name, deserialize_func, \
|
||||
construct_func) \
|
||||
REGISTER_TRT_PLUGIN_UNIQ(ctr, name, deserialize_func, construct_func)
|
||||
#define REGISTER_TRT_PLUGIN_UNIQ(ctr, name, deserialize_func, construct_func) \
|
||||
static ::paddle::inference::tensorrt::TrtPluginRegistrar \
|
||||
trt_plugin_registrar##ctr __attribute__((unused)) = \
|
||||
::paddle::inference::tensorrt::TrtPluginRegistrar( \
|
||||
name, deserialize_func, construct_func)
|
||||
|
||||
} // namespace tensorrt
|
||||
} // namespace inference
|
||||
} // namespace paddle
|
||||
@ -1,37 +0,0 @@
|
||||
/* Copyright (c) 2018 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/plugin/plugin_utils.h"
|
||||
#include <cassert>
|
||||
|
||||
namespace paddle {
|
||||
namespace inference {
|
||||
namespace tensorrt {
|
||||
|
||||
std::string ExtractOpName(const void* serial_data, size_t serial_length,
|
||||
size_t* incremental) {
|
||||
size_t op_name_char_count = *static_cast<const size_t*>(serial_data);
|
||||
*incremental = sizeof(size_t) + op_name_char_count;
|
||||
|
||||
assert(serial_length >= *incremental);
|
||||
|
||||
const char* buffer = static_cast<const char*>(serial_data) + sizeof(size_t);
|
||||
std::string op_name(buffer, op_name_char_count);
|
||||
|
||||
return op_name;
|
||||
}
|
||||
|
||||
} // namespace tensorrt
|
||||
} // namespace inference
|
||||
} // namespace paddle
|
||||
@ -1,34 +0,0 @@
|
||||
/* Copyright (c) 2018 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 <functional>
|
||||
|
||||
#include "NvInfer.h"
|
||||
#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace inference {
|
||||
namespace tensorrt {
|
||||
|
||||
typedef std::function<PluginTensorRT*(const void*, size_t)>
|
||||
PluginDeserializeFunc;
|
||||
typedef std::function<PluginTensorRT*(void)> PluginConstructFunc;
|
||||
|
||||
std::string ExtractOpName(const void* serial_data, size_t serial_length,
|
||||
size_t* incremental);
|
||||
|
||||
} // namespace tensorrt
|
||||
} // namespace inference
|
||||
} // namespze paddle
|
||||
Loading…
Reference in new issue