Files
nvidia--tensorrt/plugin/gridAnchorPlugin/gridAnchorPlugin.cpp
T
wehub-resource-sync c8a779b1bb
Docker Image CI / build-ubuntu2004 (push) Waiting to run
chore: import upstream snapshot with attribution
2026-07-13 13:36:55 +08:00

529 lines
19 KiB
C++

/*
* 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 "gridAnchorPlugin.h"
#include <cmath>
#include <iostream>
#include <memory>
#include <sstream>
#include <string_view>
#include <vector>
using namespace nvinfer1;
using namespace nvinfer1::pluginInternal;
namespace nvinfer1::plugin
{
namespace
{
std::string const kGRID_ANCHOR_PLUGIN_NAMES[] = {"GridAnchor_TRT", "GridAnchorRect_TRT"};
char const* const kGRID_ANCHOR_PLUGIN_VERSION = "1";
} // namespace
GridAnchorGenerator::GridAnchorGenerator(GridAnchorParameters const* paramIn, int32_t numLayers, char const* name)
: mPluginName(name)
, mNumLayers(numLayers)
{
PLUGIN_CUASSERT(cudaMallocHost((void**) &mNumPriors, mNumLayers * sizeof(int32_t)));
PLUGIN_CUASSERT(cudaMallocHost((void**) &mDeviceWidths, mNumLayers * sizeof(Weights)));
PLUGIN_CUASSERT(cudaMallocHost((void**) &mDeviceHeights, mNumLayers * sizeof(Weights)));
mParam.resize(mNumLayers);
for (int32_t id = 0; id < mNumLayers; id++)
{
mParam[id] = paramIn[id];
PLUGIN_VALIDATE(mParam[id].numAspectRatios >= 0 && mParam[id].aspectRatios != nullptr);
mParam[id].aspectRatios = (float*) malloc(sizeof(float) * mParam[id].numAspectRatios);
for (int32_t i = 0; i < paramIn[id].numAspectRatios; ++i)
{
mParam[id].aspectRatios[i] = paramIn[id].aspectRatios[i];
}
for (int32_t i = 0; i < 4; ++i)
{
mParam[id].variance[i] = paramIn[id].variance[i];
}
std::vector<float> tmpScales(mNumLayers + 1);
// Calculate the scales of SSD model for each layer
for (int32_t i = 0; i < mNumLayers; i++)
{
tmpScales[i] = (mParam[id].minSize + (mParam[id].maxSize - mParam[id].minSize) * id / (mNumLayers - 1));
}
// Add another 1.0f to tmpScales to prevent going out side of the vector in calculating the scale_next.
tmpScales.push_back(1.0F); // has 7 entries
// scale0 are for the first layer specifically
std::vector<float> scale0 = {0.1F, tmpScales[0], tmpScales[0]};
std::vector<float> aspect_ratios;
std::vector<float> scales;
// The first layer is different
if (id == 0)
{
for (int32_t i = 0; i < mParam[id].numAspectRatios; i++)
{
aspect_ratios.push_back(mParam[id].aspectRatios[i]);
scales.push_back(scale0[i]);
}
mNumPriors[id] = mParam[id].numAspectRatios;
}
else
{
for (int32_t i = 0; i < mParam[id].numAspectRatios; i++)
{
aspect_ratios.push_back(mParam[id].aspectRatios[i]);
}
// Additional aspect ratio of 1.0 as described in the paper
aspect_ratios.push_back(1.0);
// scales
for (int32_t i = 0; i < mParam[id].numAspectRatios; i++)
{
scales.push_back(tmpScales[id]);
}
auto scale_next = (id == mNumLayers - 1)
? 1.0
: (mParam[id].minSize + (mParam[id].maxSize - mParam[id].minSize) * (id + 1) / (mNumLayers - 1));
scales.push_back(std::sqrt(tmpScales[id] * scale_next));
mNumPriors[id] = mParam[id].numAspectRatios + 1;
}
std::vector<float> tmpWidths;
std::vector<float> tmpHeights;
// Calculate the width and height of the prior boxes
for (int32_t i = 0; i < mNumPriors[id]; i++)
{
float sqrt_AR = std::sqrt(aspect_ratios[i]);
tmpWidths.push_back(scales[i] * sqrt_AR);
tmpHeights.push_back(scales[i] / sqrt_AR);
}
mDeviceWidths[id] = copyToDevice(tmpWidths.data(), tmpWidths.size());
mDeviceHeights[id] = copyToDevice(tmpHeights.data(), tmpHeights.size());
}
}
GridAnchorGenerator::GridAnchorGenerator(void const* data, size_t length, char const* name)
: mPluginName(name)
{
char const *d = reinterpret_cast<char const*>(data), *a = d;
mNumLayers = read<int32_t>(d);
PLUGIN_CUASSERT(cudaMallocHost((void**) &mNumPriors, mNumLayers * sizeof(int32_t)));
PLUGIN_CUASSERT(cudaMallocHost((void**) &mDeviceWidths, mNumLayers * sizeof(Weights)));
PLUGIN_CUASSERT(cudaMallocHost((void**) &mDeviceHeights, mNumLayers * sizeof(Weights)));
mParam.resize(mNumLayers);
for (int32_t id = 0; id < mNumLayers; id++)
{
// we have to deserialize GridAnchorParameters by hand
mParam[id].minSize = read<float>(d);
mParam[id].maxSize = read<float>(d);
mParam[id].numAspectRatios = read<int32_t>(d);
mParam[id].aspectRatios = (float*) malloc(sizeof(float) * mParam[id].numAspectRatios);
for (int32_t i = 0; i < mParam[id].numAspectRatios; ++i)
{
mParam[id].aspectRatios[i] = read<float>(d);
}
mParam[id].H = read<int32_t>(d);
mParam[id].W = read<int32_t>(d);
for (int32_t i = 0; i < 4; ++i)
{
mParam[id].variance[i] = read<float>(d);
}
mNumPriors[id] = read<int32_t>(d);
mDeviceWidths[id] = deserializeToDevice(d, mNumPriors[id]);
mDeviceHeights[id] = deserializeToDevice(d, mNumPriors[id]);
}
PLUGIN_VALIDATE(d == a + length);
}
GridAnchorGenerator::~GridAnchorGenerator()
{
for (int32_t id = 0; id < mNumLayers; id++)
{
PLUGIN_CUERROR(cudaFree(const_cast<void*>(mDeviceWidths[id].values)));
PLUGIN_CUERROR(cudaFree(const_cast<void*>(mDeviceHeights[id].values)));
free(mParam[id].aspectRatios);
}
PLUGIN_CUERROR(cudaFreeHost(mNumPriors));
PLUGIN_CUERROR(cudaFreeHost(mDeviceWidths));
PLUGIN_CUERROR(cudaFreeHost(mDeviceHeights));
}
int32_t GridAnchorGenerator::getNbOutputs() const noexcept
{
return mNumLayers;
}
Dims GridAnchorGenerator::getOutputDimensions(int32_t index, Dims const* inputs, int32_t nbInputDims) noexcept
{
// Particularity of the PriorBox layer: no batchSize dimension needed
// 2 channels. First channel stores the mean of each prior coordinate.
// Second channel stores the variance of each prior coordinate.
return Dims3(2, mParam[index].H * mParam[index].W * mNumPriors[index] * 4, 1);
}
int32_t GridAnchorGenerator::initialize() noexcept
{
return STATUS_SUCCESS;
}
void GridAnchorGenerator::terminate() noexcept {}
size_t GridAnchorGenerator::getWorkspaceSize(int32_t maxBatchSize) const noexcept
{
return 0;
}
int32_t GridAnchorGenerator::enqueue(
int32_t batchSize, void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept
{
// Generate prior boxes for each layer
for (int32_t id = 0; id < mNumLayers; id++)
{
void* outputData = outputs[id];
pluginStatus_t status = anchorGridInference(
stream, mParam[id], mNumPriors[id], mDeviceWidths[id].values, mDeviceHeights[id].values, outputData);
if (status != STATUS_SUCCESS)
{
return status;
}
}
return STATUS_SUCCESS;
}
size_t GridAnchorGenerator::getSerializationSize() const noexcept
{
size_t sum = sizeof(int32_t); // mNumLayers
for (int32_t i = 0; i < mNumLayers; i++)
{
sum += 4 * sizeof(int32_t); // mNumPriors, mParam[i].{numAspectRatios, H, W}
sum += (6 + mParam[i].numAspectRatios)
* sizeof(float); // mParam[i].{minSize, maxSize, aspectRatios, variance[4]}
sum += mDeviceWidths[i].count * sizeof(float);
sum += mDeviceHeights[i].count * sizeof(float);
}
return sum;
}
void GridAnchorGenerator::serialize(void* buffer) const noexcept
{
char *d = reinterpret_cast<char*>(buffer), *a = d;
write(d, mNumLayers);
for (int32_t id = 0; id < mNumLayers; id++)
{
// we have to serialize GridAnchorParameters by hand
write(d, mParam[id].minSize);
write(d, mParam[id].maxSize);
write(d, mParam[id].numAspectRatios);
for (int32_t i = 0; i < mParam[id].numAspectRatios; ++i)
{
write(d, mParam[id].aspectRatios[i]);
}
write(d, mParam[id].H);
write(d, mParam[id].W);
for (int32_t i = 0; i < 4; ++i)
{
write(d, mParam[id].variance[i]);
}
write(d, mNumPriors[id]);
serializeFromDevice(d, mDeviceWidths[id]);
serializeFromDevice(d, mDeviceHeights[id]);
}
PLUGIN_ASSERT(d == a + getSerializationSize());
}
Weights GridAnchorGenerator::copyToDevice(void const* hostData, size_t count) noexcept
{
void* deviceData;
PLUGIN_CUASSERT(cudaMalloc(&deviceData, count * sizeof(float)));
PLUGIN_CUASSERT(cudaMemcpy(deviceData, hostData, count * sizeof(float), cudaMemcpyHostToDevice));
return Weights{DataType::kFLOAT, deviceData, int64_t(count)};
}
void GridAnchorGenerator::serializeFromDevice(char*& hostBuffer, Weights deviceWeights) const noexcept
{
PLUGIN_CUASSERT(
cudaMemcpy(hostBuffer, deviceWeights.values, deviceWeights.count * sizeof(float), cudaMemcpyDeviceToHost));
hostBuffer += deviceWeights.count * sizeof(float);
}
Weights GridAnchorGenerator::deserializeToDevice(char const*& hostBuffer, size_t count) noexcept
{
Weights w = copyToDevice(hostBuffer, count);
hostBuffer += count * sizeof(float);
return w;
}
bool GridAnchorGenerator::supportsFormat(DataType type, PluginFormat format) const noexcept
{
return (type == DataType::kFLOAT && format == PluginFormat::kLINEAR);
}
char const* GridAnchorGenerator::getPluginType() const noexcept
{
return mPluginName.c_str();
}
char const* GridAnchorGenerator::getPluginVersion() const noexcept
{
return kGRID_ANCHOR_PLUGIN_VERSION;
}
// Set plugin namespace
void GridAnchorGenerator::setPluginNamespace(char const* pluginNamespace) noexcept
{
mPluginNamespace = pluginNamespace;
}
char const* GridAnchorGenerator::getPluginNamespace() const noexcept
{
return mPluginNamespace.c_str();
}
#include <iostream>
// Return the DataType of the plugin output at the requested index
DataType GridAnchorGenerator::getOutputDataType(
int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept
{
PLUGIN_ASSERT(index < mNumLayers);
return DataType::kFLOAT;
}
// Configure the layer with input and output data types.
void GridAnchorGenerator::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
{
PLUGIN_ASSERT(nbOutputs == mNumLayers);
PLUGIN_ASSERT(outputDims[0].nbDims == 3);
}
// Attach the plugin object to an execution context and grant the plugin the access to some context resource.
void GridAnchorGenerator::attachToContext(
cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) noexcept
{
}
// Detach the plugin object from its execution context.
void GridAnchorGenerator::detachFromContext() noexcept {}
void GridAnchorGenerator::destroy() noexcept
{
delete this;
}
IPluginV2Ext* GridAnchorGenerator::clone() const noexcept
{
try
{
auto plugin = std::make_unique<GridAnchorGenerator>(mParam.data(), mNumLayers, mPluginName.c_str());
plugin->setPluginNamespace(mPluginNamespace.c_str());
return plugin.release();
}
catch (std::exception const& e)
{
caughtError(e);
}
return nullptr;
}
GridAnchorBasePluginCreator::GridAnchorBasePluginCreator()
{
mPluginAttributes.clear();
mPluginAttributes.emplace_back(PluginField("minSize", nullptr, PluginFieldType::kFLOAT32, 1));
mPluginAttributes.emplace_back(PluginField("maxSize", nullptr, PluginFieldType::kFLOAT32, 1));
mPluginAttributes.emplace_back(PluginField("aspectRatios", nullptr, PluginFieldType::kFLOAT32, 1));
mPluginAttributes.emplace_back(PluginField("featureMapShapes", nullptr, PluginFieldType::kINT32, 1));
mPluginAttributes.emplace_back(PluginField("variance", nullptr, PluginFieldType::kFLOAT32, 4));
mPluginAttributes.emplace_back(PluginField("numLayers", nullptr, PluginFieldType::kINT32, 1));
mFC.nbFields = mPluginAttributes.size();
mFC.fields = mPluginAttributes.data();
}
char const* GridAnchorBasePluginCreator::getPluginName() const noexcept
{
return mPluginName.c_str();
}
char const* GridAnchorBasePluginCreator::getPluginVersion() const noexcept
{
return kGRID_ANCHOR_PLUGIN_VERSION;
}
PluginFieldCollection const* GridAnchorBasePluginCreator::getFieldNames() noexcept
{
return &mFC;
}
// NOLINTNEXTLINE(readability-function-cognitive-complexity)
IPluginV2Ext* GridAnchorBasePluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept
{
try
{
using namespace std::string_view_literals;
float minScale = 0.2F, maxScale = 0.95F;
int32_t numLayers = 6;
std::vector<float> aspectRatios;
std::vector<int32_t> fMapShapes;
std::vector<float> layerVariances;
PluginField const* fields = fc->fields;
bool const isFMapRect = (kGRID_ANCHOR_PLUGIN_NAMES[1] == mPluginName);
for (int32_t i = 0; i < fc->nbFields; ++i)
{
std::string_view const attrName = fields[i].name;
if (attrName == "numLayers"sv)
{
PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
numLayers = static_cast<int32_t>(*(static_cast<int32_t const*>(fields[i].data)));
}
else if (attrName == "minSize"sv)
{
PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kFLOAT32);
minScale = static_cast<float>(*(static_cast<float const*>(fields[i].data)));
}
else if (attrName == "maxSize"sv)
{
PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kFLOAT32);
maxScale = static_cast<float>(*(static_cast<float const*>(fields[i].data)));
}
else if (attrName == "variance"sv)
{
PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kFLOAT32);
int32_t size = fields[i].length;
layerVariances.reserve(size);
auto const* lVar = static_cast<float const*>(fields[i].data);
for (int32_t j = 0; j < size; j++)
{
layerVariances.push_back(*lVar);
lVar++;
}
}
else if (attrName == "aspectRatios"sv)
{
PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kFLOAT32);
int32_t size = fields[i].length;
aspectRatios.reserve(size);
auto const* aR = static_cast<float const*>(fields[i].data);
for (int32_t j = 0; j < size; j++)
{
aspectRatios.push_back(*aR);
aR++;
}
}
else if (attrName == "featureMapShapes"sv)
{
PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
int32_t size = fields[i].length;
PLUGIN_VALIDATE(!isFMapRect || (size % 2 == 0));
fMapShapes.reserve(size);
int32_t const* fMap = static_cast<int32_t const*>(fields[i].data);
for (int32_t j = 0; j < size; j++)
{
fMapShapes.push_back(*fMap);
fMap++;
}
}
}
// Reducing the number of boxes predicted by the first layer.
// This is in accordance with the standard implementation.
std::vector<float> firstLayerAspectRatios;
PLUGIN_VALIDATE(numLayers > 0);
int32_t const numExpectedLayers = static_cast<int32_t>(fMapShapes.size()) >> (isFMapRect ? 1 : 0);
PLUGIN_VALIDATE(numExpectedLayers == numLayers);
int32_t numFirstLayerARs = 3;
// First layer only has the first 3 aspect ratios from aspectRatios
firstLayerAspectRatios.reserve(numFirstLayerARs);
for (int32_t i = 0; i < numFirstLayerARs; ++i)
{
firstLayerAspectRatios.push_back(aspectRatios[i]);
}
// A comprehensive list of box parameters that are required by anchor generator
std::vector<GridAnchorParameters> boxParams(numLayers);
// One set of box parameters for one layer
for (int32_t i = 0; i < numLayers; i++)
{
int32_t hOffset = (isFMapRect ? i * 2 : i);
int32_t wOffset = (isFMapRect ? i * 2 + 1 : i);
// Only the first layer is different
if (i == 0)
{
boxParams[i] = {minScale, maxScale, firstLayerAspectRatios.data(),
(int32_t) firstLayerAspectRatios.size(), fMapShapes[hOffset], fMapShapes[wOffset],
{layerVariances[0], layerVariances[1], layerVariances[2], layerVariances[3]}};
}
else
{
boxParams[i] = {minScale, maxScale, aspectRatios.data(), (int32_t) aspectRatios.size(),
fMapShapes[hOffset], fMapShapes[wOffset],
{layerVariances[0], layerVariances[1], layerVariances[2], layerVariances[3]}};
}
}
auto obj = std::make_unique<GridAnchorGenerator>(boxParams.data(), numLayers, mPluginName.c_str());
obj->setPluginNamespace(mNamespace.c_str());
return obj.release();
}
catch (std::exception const& e)
{
caughtError(e);
}
return nullptr;
}
IPluginV2Ext* GridAnchorBasePluginCreator::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 GridAnchor::destroy()
auto obj = std::make_unique<GridAnchorGenerator>(serialData, serialLength, mPluginName.c_str());
obj->setPluginNamespace(mNamespace.c_str());
return obj.release();
}
catch (std::exception const& e)
{
caughtError(e);
}
return nullptr;
}
GridAnchorPluginCreator::GridAnchorPluginCreator()
{
mPluginName = kGRID_ANCHOR_PLUGIN_NAMES[0];
}
GridAnchorRectPluginCreator::GridAnchorRectPluginCreator()
{
mPluginName = kGRID_ANCHOR_PLUGIN_NAMES[1];
}
} // namespace nvinfer1::plugin