chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,22 @@
|
||||
#
|
||||
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
|
||||
add_plugin_source(
|
||||
flattenConcat.cpp
|
||||
flattenConcat.h
|
||||
)
|
||||
|
||||
@@ -0,0 +1,41 @@
|
||||
#
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
---
|
||||
name: FlattenConcat_TRT
|
||||
interface: "IPluginV2Ext"
|
||||
versions:
|
||||
"1":
|
||||
attributes:
|
||||
- axis
|
||||
- ignoreBatch
|
||||
attribute_types:
|
||||
axis: int32
|
||||
ignoreBatch: int32
|
||||
attribute_length:
|
||||
axis: 1
|
||||
ignoreBatch: 1
|
||||
attribute_options:
|
||||
axis:
|
||||
min: "=1"
|
||||
max: "=3"
|
||||
ignoreBatch:
|
||||
- 0
|
||||
- 1
|
||||
attributes_required:
|
||||
- axis
|
||||
- ignoreBatch
|
||||
...
|
||||
@@ -0,0 +1,182 @@
|
||||
# flattenConcat Plugin [DEPRECATED]
|
||||
|
||||
**This plugin is deprecated since TensorRT 10.12 and will be removed in a future release. `IConcatenationLayer` can be used as an alternative as appropriate to replace its functionality.**
|
||||
|
||||
**Table Of Contents**
|
||||
- [Description](#description)
|
||||
* [Structure](#structure)
|
||||
- [Parameters](#parameters)
|
||||
- [Additional resources](#additional-resources)
|
||||
- [License](#license)
|
||||
- [Changelog](#changelog)
|
||||
- [Known issues](#known-issues)
|
||||
|
||||
## Description
|
||||
|
||||
The `flattenConcat` plugin performs input tensor flattening followed by concatenation in a single step. The SSD object detection model has bounding box predictions, including bounding box location predictions and object classifications, from different feature maps in the neural network. The bounding box location prediction tensors from different feature maps are flattened and concatenated into a single tensor, as well as the object classification tensors. Merging tensor flattening and concatenation into a single step accelerates the inference speed in TensorRT.
|
||||
|
||||
|
||||
### Structure
|
||||
|
||||
This plugin supports the NCHW format. It takes an arbitrary number of input tensors of shape `[N, C_1, H, W], [N, C_2, H, W], ..., [N, C_k, H, W]`, flattens and concatenates these input tensors, and generates an output tensor of shape `[N, C, 1, 1]` where `C = (C_1 + C_2 + ... + C_k) * H * W`.
|
||||
|
||||
For example, you have input tensor `A` of shape `[2, 2, 2, 2]`:
|
||||
```
|
||||
[[[[ 0 1]
|
||||
[ 2 3]]
|
||||
|
||||
[[ 4 5]
|
||||
[ 6 7]]]
|
||||
|
||||
|
||||
[[[ 8 9]
|
||||
[10 11]]
|
||||
|
||||
[[12 13]
|
||||
[14 15]]]]
|
||||
```
|
||||
|
||||
and input tensor `B` of shape `[2, 3, 2, 2]`:
|
||||
|
||||
```
|
||||
[[[[16 17]
|
||||
[18 19]]
|
||||
|
||||
[[20 21]
|
||||
[22 23]]
|
||||
|
||||
[[24 25]
|
||||
[26 27]]]
|
||||
|
||||
|
||||
[[[28 29]
|
||||
[30 31]]
|
||||
|
||||
[[32 33]
|
||||
[34 35]]
|
||||
|
||||
[[36 37]
|
||||
[38 39]]]]
|
||||
```
|
||||
|
||||
After `flattenConcat` for the two inputs, the output tensor of shape `[2, 20, 1, 1]` is:
|
||||
```
|
||||
[[[[ 0]]
|
||||
|
||||
[[ 1]]
|
||||
|
||||
[[ 2]]
|
||||
|
||||
[[ 3]]
|
||||
|
||||
[[ 4]]
|
||||
|
||||
[[ 5]]
|
||||
|
||||
[[ 6]]
|
||||
|
||||
[[ 7]]
|
||||
|
||||
[[16]]
|
||||
|
||||
[[17]]
|
||||
|
||||
[[18]]
|
||||
|
||||
[[19]]
|
||||
|
||||
[[20]]
|
||||
|
||||
[[21]]
|
||||
|
||||
[[22]]
|
||||
|
||||
[[23]]
|
||||
|
||||
[[24]]
|
||||
|
||||
[[25]]
|
||||
|
||||
[[26]]
|
||||
|
||||
[[27]]]
|
||||
|
||||
|
||||
[[[ 8]]
|
||||
|
||||
[[ 9]]
|
||||
|
||||
[[10]]
|
||||
|
||||
[[11]]
|
||||
|
||||
[[12]]
|
||||
|
||||
[[13]]
|
||||
|
||||
[[14]]
|
||||
|
||||
[[15]]
|
||||
|
||||
[[28]]
|
||||
|
||||
[[29]]
|
||||
|
||||
[[30]]
|
||||
|
||||
[[31]]
|
||||
|
||||
[[32]]
|
||||
|
||||
[[33]]
|
||||
|
||||
[[34]]
|
||||
|
||||
[[35]]
|
||||
|
||||
[[36]]
|
||||
|
||||
[[37]]
|
||||
|
||||
[[38]]
|
||||
|
||||
[[39]]]]
|
||||
```
|
||||
|
||||
## Parameters
|
||||
|
||||
This plugin has the plugin creator class `FlattenConcatPluginCreator` and the plugin class `FlattenConcat`.
|
||||
|
||||
The following parameters were used to create `FlattenConcat` instance:
|
||||
|
||||
| Type | Parameter | Description
|
||||
|------------------|--------------------------------|--------------------------------------------------------
|
||||
|`int` |`axis` |The dimension along which to concatenate. Currently only `axis = 1` is supported.
|
||||
|`bool` |`ignoreBatch` |Whether to ignore batch or not. Currently only `ignoreBatch = false` is supported.
|
||||
|
||||
|
||||
## Additional resources
|
||||
|
||||
The following resources provide a deeper understanding of the `flattenConcat` plugin:
|
||||
|
||||
- [SSD Caffe Implementation](https://github.com/weiliu89/caffe/tree/ssd)
|
||||
|
||||
|
||||
## License
|
||||
|
||||
For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html)
|
||||
documentation.
|
||||
|
||||
|
||||
## Changelog
|
||||
|
||||
May 2025
|
||||
Add deprecation note.
|
||||
|
||||
May 2019
|
||||
This is the first release of this `README.md` file.
|
||||
|
||||
|
||||
## Known issues
|
||||
|
||||
There are no known issues in this plugin.
|
||||
@@ -0,0 +1,411 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* 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 "flattenConcat.h"
|
||||
#include "common/dimsHelpers.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <numeric>
|
||||
#include <string>
|
||||
#include <string_view>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
using namespace nvinfer1;
|
||||
using namespace nvinfer1::pluginInternal;
|
||||
using nvinfer1::plugin::FlattenConcat;
|
||||
using nvinfer1::plugin::FlattenConcatPluginCreator;
|
||||
|
||||
static char const* const kFLATTENCONCAT_PLUGIN_VERSION{"1"};
|
||||
static char const* const kFLATTENCONCAT_PLUGIN_NAME{"FlattenConcat_TRT"};
|
||||
|
||||
FlattenConcat::FlattenConcat(int32_t concatAxis, bool ignoreBatch)
|
||||
: mIgnoreBatch(ignoreBatch)
|
||||
, mConcatAxisID(concatAxis)
|
||||
{
|
||||
PLUGIN_VALIDATE(mConcatAxisID == 1 || mConcatAxisID == 2 || mConcatAxisID == 3);
|
||||
}
|
||||
|
||||
FlattenConcat::FlattenConcat(int32_t concatAxis, bool ignoreBatch, int32_t numInputs, int32_t outputConcatAxis,
|
||||
int32_t const* inputConcatAxis, size_t const* copySize, nvinfer1::Dims const& chwDims)
|
||||
: mCopySize(numInputs)
|
||||
, mInputConcatAxis(numInputs)
|
||||
, mIgnoreBatch(ignoreBatch)
|
||||
, mConcatAxisID(concatAxis)
|
||||
, mOutputConcatAxis(outputConcatAxis)
|
||||
, mNumInputs(numInputs)
|
||||
, mCHW(chwDims)
|
||||
{
|
||||
PLUGIN_VALIDATE(mConcatAxisID >= 1 && mConcatAxisID <= 3);
|
||||
|
||||
std::copy(copySize, copySize + mNumInputs, mCopySize.begin());
|
||||
std::copy(inputConcatAxis, inputConcatAxis + mNumInputs, mInputConcatAxis.begin());
|
||||
}
|
||||
|
||||
FlattenConcat::FlattenConcat(void const* data, size_t length)
|
||||
{
|
||||
char const* d = static_cast<char const*>(data);
|
||||
char const* const a = d;
|
||||
mIgnoreBatch = read<bool>(d);
|
||||
mConcatAxisID = read<int32_t>(d);
|
||||
PLUGIN_VALIDATE(mConcatAxisID >= 1 && mConcatAxisID <= 3);
|
||||
mOutputConcatAxis = read<int32_t>(d);
|
||||
mNumInputs = read<int32_t>(d);
|
||||
|
||||
mInputConcatAxis.resize(mNumInputs);
|
||||
std::for_each(mInputConcatAxis.begin(), mInputConcatAxis.end(), [&](int32_t& inp) { inp = read<int32_t>(d); });
|
||||
|
||||
mCHW = read<nvinfer1::Dims3>(d);
|
||||
|
||||
mCopySize.resize(mNumInputs);
|
||||
std::for_each(mCopySize.begin(), mCopySize.end(), [&](size_t& inp) { inp = read<size_t>(d); });
|
||||
|
||||
PLUGIN_VALIDATE(d == a + length);
|
||||
}
|
||||
|
||||
FlattenConcat::~FlattenConcat() {}
|
||||
|
||||
int32_t FlattenConcat::getNbOutputs() const noexcept
|
||||
{
|
||||
return 1;
|
||||
}
|
||||
|
||||
Dims FlattenConcat::getOutputDimensions(int32_t index, Dims const* inputs, int32_t nbInputDims) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
PLUGIN_ASSERT(nbInputDims >= 1);
|
||||
PLUGIN_ASSERT(index == 0);
|
||||
|
||||
mNumInputs = nbInputDims;
|
||||
mCopySize.resize(mNumInputs);
|
||||
mInputConcatAxis.resize(mNumInputs);
|
||||
int32_t outputConcatAxis = 0;
|
||||
|
||||
for (int32_t i = 0; i < nbInputDims; ++i)
|
||||
{
|
||||
int32_t flattenInput = 0;
|
||||
PLUGIN_ASSERT(inputs[i].nbDims == 3);
|
||||
if (mConcatAxisID != 1)
|
||||
{
|
||||
PLUGIN_ASSERT(inputs[i].d[0] == inputs[0].d[0]);
|
||||
}
|
||||
if (mConcatAxisID != 2)
|
||||
{
|
||||
PLUGIN_ASSERT(inputs[i].d[1] == inputs[0].d[1]);
|
||||
}
|
||||
if (mConcatAxisID != 3)
|
||||
{
|
||||
PLUGIN_ASSERT(inputs[i].d[2] == inputs[0].d[2]);
|
||||
}
|
||||
flattenInput = inputs[i].d[0] * inputs[i].d[1] * inputs[i].d[2];
|
||||
outputConcatAxis += flattenInput;
|
||||
}
|
||||
|
||||
return Dims3(mConcatAxisID == 1 ? outputConcatAxis : 1, mConcatAxisID == 2 ? outputConcatAxis : 1,
|
||||
mConcatAxisID == 3 ? outputConcatAxis : 1);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return Dims{};
|
||||
}
|
||||
|
||||
int32_t FlattenConcat::initialize() noexcept
|
||||
{
|
||||
return STATUS_SUCCESS;
|
||||
}
|
||||
|
||||
void FlattenConcat::terminate() noexcept {}
|
||||
|
||||
size_t FlattenConcat::getWorkspaceSize(int32_t) const noexcept
|
||||
{
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t FlattenConcat::enqueue(
|
||||
int32_t batchSize, void const* const* inputs, void* const* outputs, void*, cudaStream_t stream) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
PLUGIN_ASSERT(mConcatAxisID != 0);
|
||||
// mCHW is the first input tensor
|
||||
auto numConcats = static_cast<int32_t>(pluginInternal::volume(mCHW, /*start*/ 0, /*stop*/ mConcatAxisID - 1));
|
||||
|
||||
// Num concats will be proportional to number of samples in a batch
|
||||
if (!mIgnoreBatch)
|
||||
{
|
||||
numConcats *= batchSize;
|
||||
}
|
||||
|
||||
auto* output = static_cast<float*>(outputs[0]);
|
||||
int32_t offset = 0;
|
||||
for (int32_t i = 0; i < mNumInputs; ++i)
|
||||
{
|
||||
auto const* input = static_cast<float const*>(inputs[i]);
|
||||
for (int32_t n = 0; n < numConcats; ++n)
|
||||
{
|
||||
auto status = mCublasWrapper->cublasScopy(mCublas, mInputConcatAxis[i], input + n * mInputConcatAxis[i],
|
||||
1, output + (n * mOutputConcatAxis + offset), 1);
|
||||
|
||||
if (status != CUBLAS_STATUS_SUCCESS)
|
||||
{
|
||||
return STATUS_FAILURE;
|
||||
}
|
||||
}
|
||||
offset += mInputConcatAxis[i];
|
||||
}
|
||||
|
||||
return STATUS_SUCCESS;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return -1;
|
||||
}
|
||||
|
||||
size_t FlattenConcat::getSerializationSize() const noexcept
|
||||
{
|
||||
return sizeof(bool) + sizeof(int32_t) * (3 + mNumInputs) + sizeof(nvinfer1::Dims)
|
||||
+ (sizeof(decltype(mCopySize)::value_type) * mNumInputs);
|
||||
}
|
||||
|
||||
void FlattenConcat::serialize(void* buffer) const noexcept
|
||||
{
|
||||
char* d = static_cast<char*>(buffer);
|
||||
char const* const a = d;
|
||||
write(d, mIgnoreBatch);
|
||||
write(d, mConcatAxisID);
|
||||
write(d, mOutputConcatAxis);
|
||||
write(d, mNumInputs);
|
||||
for (int32_t i = 0; i < mNumInputs; ++i)
|
||||
{
|
||||
write(d, mInputConcatAxis[i]);
|
||||
}
|
||||
write(d, mCHW);
|
||||
for (int32_t i = 0; i < mNumInputs; ++i)
|
||||
{
|
||||
write(d, mCopySize[i]);
|
||||
}
|
||||
PLUGIN_ASSERT(d == a + getSerializationSize());
|
||||
}
|
||||
|
||||
// Attach the plugin object to an execution context and grant the plugin the access to some context resource.
|
||||
void FlattenConcat::attachToContext(
|
||||
cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
mCublasWrapper = createPluginCublasWrapper(gpuAllocator);
|
||||
mCublas = mCublasWrapper->getCublasHandle();
|
||||
PLUGIN_VALIDATE(mCublas != nullptr);
|
||||
}
|
||||
catch (const std::exception& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
}
|
||||
|
||||
// Detach the plugin object from its execution context.
|
||||
void FlattenConcat::detachFromContext() noexcept {}
|
||||
|
||||
// Set plugin namespace
|
||||
void FlattenConcat::setPluginNamespace(char const* pluginNamespace) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
}
|
||||
|
||||
char const* FlattenConcat::getPluginNamespace() const noexcept
|
||||
{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
|
||||
// Return the DataType of the plugin output at the requested index
|
||||
DataType FlattenConcat::getOutputDataType(
|
||||
int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept
|
||||
{
|
||||
PLUGIN_ASSERT(index < 3);
|
||||
return DataType::kFLOAT;
|
||||
}
|
||||
|
||||
void FlattenConcat::configurePlugin(Dims const* inputDims, int32_t nbInputs, Dims const* outputDims, int32_t nbOutputs,
|
||||
DataType const* inputTypes, DataType const* outputTypes, bool const* inputIsBroadcast,
|
||||
bool const* outputIsBroadcast, PluginFormat floatFormat, int32_t maxBatchSize) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
PLUGIN_ASSERT(nbOutputs == 1);
|
||||
mCHW = inputDims[0];
|
||||
mNumInputs = nbInputs;
|
||||
PLUGIN_ASSERT(inputDims[0].nbDims == 3);
|
||||
|
||||
mInputConcatAxis.resize(mNumInputs);
|
||||
for (int32_t i = 0; i < nbInputs; ++i)
|
||||
{
|
||||
int32_t flattenInput = 0;
|
||||
PLUGIN_ASSERT(inputDims[i].nbDims == 3);
|
||||
if (mConcatAxisID != 1)
|
||||
{
|
||||
PLUGIN_ASSERT(inputDims[i].d[0] == inputDims[0].d[0]);
|
||||
}
|
||||
if (mConcatAxisID != 2)
|
||||
{
|
||||
PLUGIN_ASSERT(inputDims[i].d[1] == inputDims[0].d[1]);
|
||||
}
|
||||
if (mConcatAxisID != 3)
|
||||
{
|
||||
PLUGIN_ASSERT(inputDims[i].d[2] == inputDims[0].d[2]);
|
||||
}
|
||||
flattenInput = inputDims[i].d[0] * inputDims[i].d[1] * inputDims[i].d[2];
|
||||
mInputConcatAxis[i] = flattenInput;
|
||||
mOutputConcatAxis += mInputConcatAxis[i];
|
||||
}
|
||||
|
||||
mCopySize.resize(mNumInputs);
|
||||
for (int32_t i = 0; i < nbInputs; ++i)
|
||||
{
|
||||
mCopySize[i] = inputDims[i].d[0] * inputDims[i].d[1] * inputDims[i].d[2] * sizeof(float);
|
||||
}
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
}
|
||||
|
||||
bool FlattenConcat::supportsFormat(DataType type, PluginFormat format) const noexcept
|
||||
{
|
||||
return (type == DataType::kFLOAT && format == PluginFormat::kLINEAR);
|
||||
}
|
||||
char const* FlattenConcat::getPluginType() const noexcept
|
||||
{
|
||||
return "FlattenConcat_TRT";
|
||||
}
|
||||
|
||||
char const* FlattenConcat::getPluginVersion() const noexcept
|
||||
{
|
||||
return "1";
|
||||
}
|
||||
|
||||
void FlattenConcat::destroy() noexcept
|
||||
{
|
||||
delete this;
|
||||
}
|
||||
|
||||
IPluginV2Ext* FlattenConcat::clone() const noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
auto plugin = std::make_unique<FlattenConcat>(mConcatAxisID, mIgnoreBatch, mNumInputs, mOutputConcatAxis,
|
||||
mInputConcatAxis.data(), mCopySize.data(), mCHW);
|
||||
plugin->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return plugin.release();
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
FlattenConcatPluginCreator::FlattenConcatPluginCreator()
|
||||
{
|
||||
mPluginAttributes.clear();
|
||||
mPluginAttributes.emplace_back(PluginField("axis", nullptr, PluginFieldType::kINT32, 1));
|
||||
mPluginAttributes.emplace_back(PluginField("ignoreBatch", nullptr, PluginFieldType::kINT32, 1));
|
||||
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
|
||||
char const* FlattenConcatPluginCreator::getPluginName() const noexcept
|
||||
{
|
||||
return kFLATTENCONCAT_PLUGIN_NAME;
|
||||
}
|
||||
|
||||
char const* FlattenConcatPluginCreator::getPluginVersion() const noexcept
|
||||
{
|
||||
return kFLATTENCONCAT_PLUGIN_VERSION;
|
||||
}
|
||||
|
||||
PluginFieldCollection const* FlattenConcatPluginCreator::getFieldNames() noexcept
|
||||
{
|
||||
return &mFC;
|
||||
}
|
||||
|
||||
IPluginV2Ext* FlattenConcatPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
using namespace std::string_view_literals;
|
||||
plugin::validateRequiredAttributesExist({"axis", "ignoreBatch"}, fc);
|
||||
PluginField const* fields = fc->fields;
|
||||
for (int32_t i = 0; i < fc->nbFields; ++i)
|
||||
{
|
||||
std::string_view const attrName = fields[i].name;
|
||||
if (attrName == "axis"sv)
|
||||
{
|
||||
PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
|
||||
mConcatAxisID = *(static_cast<int32_t const*>(fields[i].data));
|
||||
}
|
||||
if (attrName == "ignoreBatch"sv)
|
||||
{
|
||||
PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
|
||||
auto ignoreBatch = *(static_cast<int32_t const*>(fields[i].data));
|
||||
PLUGIN_VALIDATE(ignoreBatch == 0 || ignoreBatch == 1);
|
||||
mIgnoreBatch = static_cast<bool>(ignoreBatch);
|
||||
}
|
||||
}
|
||||
|
||||
auto plugin = std::make_unique<FlattenConcat>(mConcatAxisID, mIgnoreBatch);
|
||||
plugin->setPluginNamespace(mNamespace.c_str());
|
||||
return plugin.release();
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
IPluginV2Ext* FlattenConcatPluginCreator::deserializePlugin(
|
||||
char const* name, void const* serialData, size_t serialLength) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
// This object will be deleted when the network is destroyed, which will
|
||||
// call Concat::destroy()
|
||||
auto plugin = std::make_unique<FlattenConcat>(serialData, serialLength);
|
||||
plugin->setPluginNamespace(mNamespace.c_str());
|
||||
return plugin.release();
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
@@ -0,0 +1,132 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* 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.
|
||||
*/
|
||||
#ifndef TRT_FLATTENCONCAT_PLUGIN_H
|
||||
#define TRT_FLATTENCONCAT_PLUGIN_H
|
||||
|
||||
#include "NvInferPlugin.h"
|
||||
#include "common/plugin.h"
|
||||
#include <cstdlib>
|
||||
#include <iostream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
class FlattenConcat : public IPluginV2Ext
|
||||
{
|
||||
public:
|
||||
FlattenConcat(int32_t concatAxis, bool ignoreBatch);
|
||||
|
||||
FlattenConcat(int32_t concatAxis, bool ignoreBatch, int32_t numInputs, int32_t outputConcatAxis,
|
||||
int32_t const* inputConcatAxis, size_t const* copySize, nvinfer1::Dims const& chwDims);
|
||||
|
||||
FlattenConcat(void const* data, size_t length);
|
||||
|
||||
~FlattenConcat() override;
|
||||
|
||||
FlattenConcat() = delete;
|
||||
|
||||
int32_t getNbOutputs() const noexcept override;
|
||||
|
||||
Dims getOutputDimensions(int32_t index, Dims const* inputs, int32_t nbInputDims) noexcept override;
|
||||
|
||||
int32_t initialize() noexcept override;
|
||||
|
||||
void terminate() noexcept override;
|
||||
|
||||
size_t getWorkspaceSize(int32_t) const noexcept override;
|
||||
|
||||
int32_t enqueue(int32_t batchSize, void const* const* inputs, void* const* outputs, void* workspace,
|
||||
cudaStream_t stream) noexcept override;
|
||||
|
||||
DataType getOutputDataType(
|
||||
int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept override;
|
||||
|
||||
size_t getSerializationSize() const noexcept override;
|
||||
|
||||
void serialize(void* buffer) const noexcept override;
|
||||
|
||||
void configurePlugin(Dims const* inputDims, int32_t nbInputs, Dims const* outputDims, int32_t nbOutputs,
|
||||
DataType const* inputTypes, DataType const* outputTypes, bool const* inputIsBroadcast,
|
||||
bool const* outputIsBroadcast, PluginFormat floatFormat, int32_t maxBatchSize) noexcept override;
|
||||
|
||||
bool supportsFormat(DataType type, PluginFormat format) const noexcept override;
|
||||
|
||||
void detachFromContext() noexcept override;
|
||||
|
||||
char const* getPluginType() const noexcept override;
|
||||
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
|
||||
void destroy() noexcept override;
|
||||
|
||||
void attachToContext(
|
||||
cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) noexcept override;
|
||||
IPluginV2Ext* clone() const noexcept override;
|
||||
|
||||
void setPluginNamespace(char const* pluginNamespace) noexcept override;
|
||||
|
||||
char const* getPluginNamespace() const noexcept override;
|
||||
|
||||
private:
|
||||
Weights copyToDevice(void const* hostData, size_t count) noexcept;
|
||||
|
||||
void serializeFromDevice(char*& hostBuffer, Weights deviceWeights) const noexcept;
|
||||
|
||||
Weights deserializeToDevice(char const*& hostBuffer, size_t count) noexcept;
|
||||
|
||||
std::vector<size_t> mCopySize;
|
||||
std::vector<int32_t> mInputConcatAxis;
|
||||
bool mIgnoreBatch{false};
|
||||
int32_t mConcatAxisID{0}, mOutputConcatAxis{0}, mNumInputs{0};
|
||||
nvinfer1::Dims mCHW;
|
||||
std::string mPluginNamespace;
|
||||
nvinfer1::pluginInternal::cublasHandle_t mCublas{nullptr};
|
||||
// the wrapper pointer is shared among all plugins attached to the same context.
|
||||
std::shared_ptr<nvinfer1::pluginInternal::CublasWrapper> mCublasWrapper;
|
||||
};
|
||||
|
||||
class FlattenConcatPluginCreator : public nvinfer1::pluginInternal::BaseCreator
|
||||
{
|
||||
public:
|
||||
FlattenConcatPluginCreator();
|
||||
|
||||
~FlattenConcatPluginCreator() override = default;
|
||||
|
||||
char const* getPluginName() const noexcept override;
|
||||
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
|
||||
PluginFieldCollection const* getFieldNames() noexcept override;
|
||||
|
||||
IPluginV2Ext* createPlugin(char const* name, PluginFieldCollection const* fc) noexcept override;
|
||||
|
||||
IPluginV2Ext* deserializePlugin(char const* name, void const* serialData, size_t serialLength) noexcept override;
|
||||
|
||||
private:
|
||||
PluginFieldCollection mFC;
|
||||
bool mIgnoreBatch{false};
|
||||
int32_t mConcatAxisID;
|
||||
std::vector<PluginField> mPluginAttributes;
|
||||
};
|
||||
|
||||
} // namespace plugin
|
||||
} // namespace nvinfer1
|
||||
|
||||
#endif // TRT_FLATTENCONCAT_PLUGIN_H
|
||||
Reference in New Issue
Block a user