peizhilin
7 years ago
commit
5670e9ea1a
@ -1,4 +1,5 @@
|
||||
nv_library(tensorrt_engine SRCS engine.cc DEPS framework_proto device_context)
|
||||
nv_test(test_tensorrt SRCS test_tensorrt.cc DEPS dynload_cuda device_context dynamic_loader)
|
||||
nv_test(test_tensorrt_engine SRCS test_engine.cc DEPS dynload_cuda tensorrt_engine)
|
||||
add_subdirectory(plugin)
|
||||
add_subdirectory(convert)
|
||||
|
||||
@ -0,0 +1,75 @@
|
||||
/* 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();
|
||||
|
||||
// Get Attrs
|
||||
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);
|
||||
@ -0,0 +1 @@
|
||||
nv_library(tensorrt_plugin SRCS trt_plugin.cc split_op_plugin.cu DEPS enforce)
|
||||
@ -0,0 +1,111 @@
|
||||
// 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 <cassert>
|
||||
#include <cstring>
|
||||
#include <type_traits>
|
||||
#include <vector>
|
||||
|
||||
template <typename T>
|
||||
inline void SerializeValue(void** buffer, T const& value);
|
||||
|
||||
template <typename T>
|
||||
inline void DeserializeValue(void const** buffer, size_t* buffer_size,
|
||||
T* value);
|
||||
|
||||
namespace {
|
||||
|
||||
template <typename T, class Enable = void>
|
||||
struct Serializer {};
|
||||
|
||||
template <typename T>
|
||||
struct Serializer<T, typename std::enable_if<std::is_arithmetic<T>::value ||
|
||||
std::is_enum<T>::value ||
|
||||
std::is_pod<T>::value>::type> {
|
||||
static size_t SerializedSize(T const& value) { return sizeof(T); }
|
||||
static void Serialize(void** buffer, T const& value) {
|
||||
std::memcpy(*buffer, &value, sizeof(T));
|
||||
reinterpret_cast<char*&>(*buffer) += sizeof(T);
|
||||
}
|
||||
static void Deserialize(void const** buffer, size_t* buffer_size, T* value) {
|
||||
assert(*buffer_size >= sizeof(T));
|
||||
std::memcpy(value, *buffer, sizeof(T));
|
||||
reinterpret_cast<char const*&>(*buffer) += sizeof(T);
|
||||
*buffer_size -= sizeof(T);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct Serializer<const char*> {
|
||||
static size_t SerializedSize(const char* value) { return strlen(value) + 1; }
|
||||
static void Serialize(void** buffer, const char* value) {
|
||||
std::strcpy(static_cast<char*>(*buffer), value);
|
||||
reinterpret_cast<char*&>(*buffer) += strlen(value) + 1;
|
||||
}
|
||||
static void Deserialize(void const** buffer, size_t* buffer_size,
|
||||
const char** value) {
|
||||
*value = static_cast<char const*>(*buffer);
|
||||
size_t data_size = strnlen(*value, *buffer_size) + 1;
|
||||
assert(*buffer_size >= data_size);
|
||||
reinterpret_cast<char const*&>(*buffer) += data_size;
|
||||
*buffer_size -= data_size;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct Serializer<std::vector<T>,
|
||||
typename std::enable_if<std::is_arithmetic<T>::value ||
|
||||
std::is_enum<T>::value ||
|
||||
std::is_pod<T>::value>::type> {
|
||||
static size_t SerializedSize(std::vector<T> const& value) {
|
||||
return sizeof(value.size()) + value.size() * sizeof(T);
|
||||
}
|
||||
static void Serialize(void** buffer, std::vector<T> const& value) {
|
||||
SerializeValue(buffer, value.size());
|
||||
size_t nbyte = value.size() * sizeof(T);
|
||||
std::memcpy(*buffer, value.data(), nbyte);
|
||||
reinterpret_cast<char*&>(*buffer) += nbyte;
|
||||
}
|
||||
static void Deserialize(void const** buffer, size_t* buffer_size,
|
||||
std::vector<T>* value) {
|
||||
size_t size;
|
||||
DeserializeValue(buffer, buffer_size, &size);
|
||||
value->resize(size);
|
||||
size_t nbyte = value->size() * sizeof(T);
|
||||
assert(*buffer_size >= nbyte);
|
||||
std::memcpy(value->data(), *buffer, nbyte);
|
||||
reinterpret_cast<char const*&>(*buffer) += nbyte;
|
||||
*buffer_size -= nbyte;
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
template <typename T>
|
||||
inline size_t SerializedSize(T const& value) {
|
||||
return Serializer<T>::SerializedSize(value);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
inline void SerializeValue(void** buffer, T const& value) {
|
||||
return Serializer<T>::Serialize(buffer, value);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
inline void DeserializeValue(void const** buffer, size_t* buffer_size,
|
||||
T* value) {
|
||||
return Serializer<T>::Deserialize(buffer, buffer_size, value);
|
||||
}
|
||||
@ -0,0 +1,81 @@
|
||||
// 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 <stdio.h>
|
||||
#include <cassert>
|
||||
#include "paddle/fluid/inference/tensorrt/plugin/split_op_plugin.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace inference {
|
||||
namespace tensorrt {
|
||||
|
||||
nvinfer1::Dims SplitPlugin::getOutputDimensions(int index,
|
||||
const nvinfer1::Dims* inputDims,
|
||||
int nbInputs) {
|
||||
assert(nbInputs == 1);
|
||||
assert(index < this->getNbOutputs());
|
||||
nvinfer1::Dims const& input_dims = inputDims[0];
|
||||
nvinfer1::Dims output_dims = input_dims;
|
||||
output_dims.d[axis_] = output_length_.at(index);
|
||||
return output_dims;
|
||||
}
|
||||
|
||||
int SplitPlugin::initialize() {
|
||||
std::vector<int> segment_offsets(1, 0);
|
||||
for (int i = 0; i < this->getNbOutputs(); ++i) {
|
||||
segment_offsets.push_back(segment_offsets.back() + output_length_[i]);
|
||||
}
|
||||
segment_offsets_ = segment_offsets;
|
||||
nvinfer1::Dims dims = this->getInputDims(0);
|
||||
nx_ = 1;
|
||||
for (int i = dims.nbDims - 1; i > axis_; --i) {
|
||||
nx_ *= dims.d[i];
|
||||
}
|
||||
ny_ = dims.d[axis_];
|
||||
nz_ = 1;
|
||||
for (int i = axis_ - 1; i >= 0; --i) {
|
||||
nz_ *= dims.d[i];
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
int SplitPlugin::enqueue(int batchSize, const void* const* inputs,
|
||||
void** outputs, void* workspace, cudaStream_t stream) {
|
||||
auto const& input_dims = this->getInputDims(0);
|
||||
int input_size = 0;
|
||||
float const* idata = reinterpret_cast<float const*>(inputs[0]);
|
||||
float** odatas = reinterpret_cast<float**>(outputs);
|
||||
|
||||
// kernel impl here.
|
||||
int inputBatchOffset = nx_ * ny_ * nz_;
|
||||
for (size_t i = 0; i < this->getNbOutputs(); i++) {
|
||||
for (size_t j = 0; j < batchSize; j++) {
|
||||
cudaMemcpyAsync(
|
||||
odatas[i] +
|
||||
j * (segment_offsets_[i + 1] - segment_offsets_[i]) * nx_ *
|
||||
sizeof(float),
|
||||
inputs[0] +
|
||||
(inputBatchOffset * j + segment_offsets_[i] * nx_) *
|
||||
sizeof(float),
|
||||
(segment_offsets_[i + 1] - segment_offsets_[i]) * nx_ * sizeof(float),
|
||||
cudaMemcpyDeviceToDevice, stream);
|
||||
}
|
||||
}
|
||||
|
||||
return cudaGetLastError() != cudaSuccess;
|
||||
}
|
||||
|
||||
} // tensorrt
|
||||
} // inference
|
||||
} // paddle
|
||||
@ -0,0 +1,74 @@
|
||||
// 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 "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace inference {
|
||||
namespace tensorrt {
|
||||
|
||||
class SplitPlugin : public PluginTensorRT {
|
||||
int axis_;
|
||||
std::vector<int> output_length_;
|
||||
int nx_, ny_, nz_;
|
||||
std::vector<int> segment_offsets_;
|
||||
|
||||
protected:
|
||||
virtual size_t getSerializationSize() override {
|
||||
return SerializedSize(axis_) + SerializedSize(output_length_) +
|
||||
getBaseSerializationSize();
|
||||
}
|
||||
|
||||
// TRT will call this func when we need to serialize the configuration of
|
||||
// tensorrt.
|
||||
// It should not be called by users.
|
||||
virtual void serialize(void *buffer) override {
|
||||
serializeBase(buffer);
|
||||
SerializeValue(&buffer, axis_);
|
||||
SerializeValue(&buffer, output_length_);
|
||||
}
|
||||
|
||||
public:
|
||||
SplitPlugin(int axis, std::vector<int> const &output_lengths)
|
||||
: axis_(axis), output_length_(output_lengths) {
|
||||
assert(axis <= nvinfer1::Dims::MAX_DIMS);
|
||||
}
|
||||
|
||||
// It was used for tensorrt deserialization.
|
||||
// It should not be called by users.
|
||||
SplitPlugin(void const *serialData, size_t serialLength) {
|
||||
deserializeBase(serialData, serialLength);
|
||||
DeserializeValue(&serialData, &serialLength, &axis_);
|
||||
DeserializeValue(&serialData, &serialLength, &output_length_);
|
||||
}
|
||||
|
||||
SplitPlugin *clone() const override {
|
||||
return new SplitPlugin(axis_, output_length_);
|
||||
}
|
||||
|
||||
virtual const char *getPluginType() const override { return "split"; }
|
||||
virtual int getNbOutputs() const override { return output_length_.size(); }
|
||||
virtual nvinfer1::Dims getOutputDimensions(int index,
|
||||
const nvinfer1::Dims *inputs,
|
||||
int nbInputDims) override;
|
||||
virtual int initialize() override;
|
||||
virtual int enqueue(int batchSize, const void *const *inputs, void **outputs,
|
||||
void *workspace, cudaStream_t stream) override;
|
||||
};
|
||||
|
||||
} // tensorrt
|
||||
} // inference
|
||||
} // paddle
|
||||
@ -0,0 +1,61 @@
|
||||
// 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/trt_plugin.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace inference {
|
||||
namespace tensorrt {
|
||||
|
||||
void PluginTensorRT::serializeBase(void*& buffer) {
|
||||
SerializeValue(&buffer, input_dims_);
|
||||
SerializeValue(&buffer, max_batch_size_);
|
||||
SerializeValue(&buffer, data_type_);
|
||||
SerializeValue(&buffer, data_format_);
|
||||
}
|
||||
|
||||
void PluginTensorRT::deserializeBase(void const*& serialData,
|
||||
size_t& serialLength) {
|
||||
DeserializeValue(&serialData, &serialLength, &input_dims_);
|
||||
DeserializeValue(&serialData, &serialLength, &max_batch_size_);
|
||||
DeserializeValue(&serialData, &serialLength, &data_type_);
|
||||
DeserializeValue(&serialData, &serialLength, &data_format_);
|
||||
}
|
||||
|
||||
size_t PluginTensorRT::getBaseSerializationSize() {
|
||||
return (SerializedSize(input_dims_) + SerializedSize(max_batch_size_) +
|
||||
SerializedSize(data_type_) + SerializedSize(data_format_));
|
||||
}
|
||||
|
||||
bool PluginTensorRT::supportsFormat(nvinfer1::DataType type,
|
||||
nvinfer1::PluginFormat format) const {
|
||||
return ((type == nvinfer1::DataType::kFLOAT) &&
|
||||
(format == nvinfer1::PluginFormat::kNCHW));
|
||||
}
|
||||
|
||||
void PluginTensorRT::configureWithFormat(const nvinfer1::Dims* inputDims,
|
||||
int nbInputs,
|
||||
const nvinfer1::Dims* outputDims,
|
||||
int nbOutputs, nvinfer1::DataType type,
|
||||
nvinfer1::PluginFormat format,
|
||||
int maxBatchSize) {
|
||||
data_type_ = type;
|
||||
data_format_ = format;
|
||||
input_dims_.assign(inputDims, inputDims + nbInputs);
|
||||
max_batch_size_ = maxBatchSize;
|
||||
}
|
||||
|
||||
} // namespace tensorrt
|
||||
} // namespace inference
|
||||
} // namespace paddle
|
||||
@ -0,0 +1,80 @@
|
||||
// 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 <cassert>
|
||||
#include <cstring>
|
||||
#include <iostream>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
#include "NvInfer.h"
|
||||
|
||||
#include "paddle/fluid/inference/tensorrt/plugin/serialize.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace inference {
|
||||
namespace tensorrt {
|
||||
|
||||
class PluginTensorRT : public nvinfer1::IPluginExt {
|
||||
public:
|
||||
PluginTensorRT() {}
|
||||
PluginTensorRT(const void* serialized_data, size_t length) {}
|
||||
nvinfer1::Dims const& getInputDims(int index) const {
|
||||
return input_dims_.at(index);
|
||||
}
|
||||
size_t getMaxBatchSize() const { return max_batch_size_; }
|
||||
nvinfer1::DataType getDataType() const { return data_type_; }
|
||||
nvinfer1::PluginFormat getDataFormat() const { return data_format_; }
|
||||
virtual const char* getPluginVersion() const { return "1"; }
|
||||
size_t getWorkspaceSize(int) const override { return 0; }
|
||||
void terminate() override {}
|
||||
virtual ~PluginTensorRT() {}
|
||||
// Check format support. The default is FLOAT32 and NCHW.
|
||||
bool supportsFormat(nvinfer1::DataType type,
|
||||
nvinfer1::PluginFormat format) const override;
|
||||
void configureWithFormat(const nvinfer1::Dims* inputDims, int nbInputs,
|
||||
const nvinfer1::Dims* outputDims, int nbOutputs,
|
||||
nvinfer1::DataType type,
|
||||
nvinfer1::PluginFormat format,
|
||||
int maxBatchSize) override;
|
||||
|
||||
// *NOTE* The following functions need to be overrided in the subclass.
|
||||
virtual nvinfer1::IPluginExt* clone() const = 0;
|
||||
virtual const char* getPluginType() const = 0;
|
||||
// Initialize the layer for execution. This is called when the engine is
|
||||
// created.
|
||||
int initialize() override { return 0; }
|
||||
// Serialize the layer config to buffer.
|
||||
virtual void serialize(void* buffer) = 0;
|
||||
virtual size_t getSerializationSize() = 0;
|
||||
virtual int enqueue(int batchSize, const void* const* inputs, void** outputs,
|
||||
void* workspace, cudaStream_t stream) = 0;
|
||||
|
||||
protected:
|
||||
// Deserialize input_dims, max_batch_size, data_type, data_format
|
||||
void deserializeBase(void const*& serialData, size_t& serialLength);
|
||||
size_t getBaseSerializationSize();
|
||||
// Serialize input_dims, max_batch_size, data_type, data_format
|
||||
void serializeBase(void*& buffer);
|
||||
|
||||
std::vector<nvinfer1::Dims> input_dims_;
|
||||
size_t max_batch_size_;
|
||||
nvinfer1::DataType data_type_;
|
||||
nvinfer1::PluginFormat data_format_;
|
||||
};
|
||||
|
||||
} // namespace tensorrt
|
||||
} // namespace inference
|
||||
} // namespace paddle
|
||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in new issue