有关yolov4-tiny chunk部分的代码 about yolo-tensorrt HOT 10 CLOSED

#include <cmath>
#include <stdio.h>
#include <cassert>
#include <iostream>
#include "chunk_V2.h"
#include <cuda_runtime.h>
#define ASSERT(assertion)                                        \
    {                                                            \
        if (!(assertion))                                        \
        {                                                        \
            std::cout<<"ASSERTION FAILED in "                    \
                     <<__FILE__<<":"<<__LINE__                   \
                     <<std::endl;                                \
            abort();                                             \
        }                                                        \
    }

#define NV_CUDA_CHECK(status)                                                                      \
    {                                                                                              \
        if (status != 0)                                                                           \
        {                                                                                          \
            std::cout << "Cuda failure: " << cudaGetErrorString(status) << " in file " << __FILE__ \
                      << " at line " << __LINE__ << std::endl;                                     \
            abort();                                                                               \
        }                                                                                          \
    }

namespace nvinfer1
{
	ChunkV2::ChunkV2()
	{

	}
	ChunkV2::ChunkV2(const void* buffer, size_t size) 
	{
		assert(size == sizeof(_n_size_split));
		_n_size_split = *reinterpret_cast<const int*>(buffer);
	}
	ChunkV2::~ChunkV2()
	{

	}

	int ChunkV2::getNbOutputs() const
	{
		return 2;
	}

	Dims ChunkV2::getOutputDimensions(int index, const Dims* inputs, int nbInputDims)
	{
		assert(nbInputDims == 1);
		assert(index == 0 || index == 1);
		return Dims3(inputs[0].d[0] / 2, inputs[0].d[1], inputs[0].d[2]);
	}

	int ChunkV2::initialize()
	{
		return 0;
	}

	void ChunkV2::terminate()
	{
	}

	size_t ChunkV2::getWorkspaceSize(int maxBatchSize) const
	{
		return 0;
	}

	size_t ChunkV2::getSerializationSize() const
	{
		return sizeof(_n_size_split);
	}

	void ChunkV2::serialize(void *buffer)const
	{
		*reinterpret_cast<int*>(buffer) = _n_size_split;
	}

	bool ChunkV2::supportsFormat(DataType type, PluginFormat format) const
	{
		return ((type == DataType::kFLOAT || type == DataType::kHALF || type == DataType::kINT8) &&
				(format == PluginFormat::kNCHW));
	}

	// Set plugin namespace
	void ChunkV2::setPluginNamespace(const char* pluginNamespace)
	{
		_s_plugin_namespace = pluginNamespace;
	}

	const char* ChunkV2::getPluginNamespace() const
	{
		return _s_plugin_namespace.c_str();
	}

	// Configure the layer with input and output data types.
	void ChunkV2::configureWithFormat(
			const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs,
			DataType type, PluginFormat format, int maxBatchSize){
		_n_size_split = inputDims->d[0] / 2 * inputDims->d[1] * inputDims->d[2] *sizeof(float);
        std::cerr << _n_size_split << std::endl;
		ASSERT(format == PluginFormat::kNCHW);
		ASSERT(type == DataType::kFLOAT || type == DataType::kHALF);
		//mDataType = type;
		ASSERT(inputDims[0].nbDims >= 1); // number of dimensions of the input tensor must be >=1

	}

	const char* ChunkV2::getPluginType()const
	{
		return "CHUNK_TRT";
	}

	const char* ChunkV2::getPluginVersion() const
	{
		return "1.0";
	}

	void ChunkV2::destroy()
	{
		delete this;
	}


	// Clone the plugin
	IPluginV2* ChunkV2::clone() const
	{
		ChunkV2 *p = new ChunkV2();
		p->_n_size_split = _n_size_split;
		p->setPluginNamespace(_s_plugin_namespace.c_str());
		return p;
	}

	int ChunkV2::enqueue(int batchSize,
		const void* const* inputs,
		void** outputs,
		void* workspace,
		cudaStream_t stream)
	{
		for (int b = 0; b < batchSize; ++b)
		{
			NV_CUDA_CHECK(cudaMemcpy((char*)outputs[0] + b * _n_size_split, (char*)inputs[0] + b * 2 * _n_size_split, _n_size_split, cudaMemcpyDeviceToDevice));
			NV_CUDA_CHECK(cudaMemcpy((char*)outputs[1] + b * _n_size_split, (char*)inputs[0] + b * 2 * _n_size_split + _n_size_split, _n_size_split, cudaMemcpyDeviceToDevice));
		}
		return 0;
	}


    PluginFieldCollection ChunkV2PluginCreator::_fc{};
    std::vector<PluginField> ChunkV2PluginCreator::_vec_plugin_attributes;

    ChunkV2PluginCreator::ChunkV2PluginCreator()
    {
        _vec_plugin_attributes.clear();
        _fc.nbFields = _vec_plugin_attributes.size();
        _fc.fields = _vec_plugin_attributes.data();
    }

    const char* ChunkV2PluginCreator::getPluginName() const
    {
        return "CHUNK_TRT";
    }

    const char* ChunkV2PluginCreator::getPluginVersion() const
    {
        return "1.0";
    }

    const PluginFieldCollection* ChunkV2PluginCreator::getFieldNames()
    {
        return &_fc;
    }

	IPluginV2* ChunkV2PluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc)
    {
		ChunkV2* obj = new ChunkV2();
        obj->setPluginNamespace(_s_name_space.c_str());
        return obj;
    }

	IPluginV2* ChunkV2PluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength)
    {
		ChunkV2* obj = new ChunkV2(serialData,serialLength);
        obj->setPluginNamespace(_s_name_space.c_str());
        return obj;
    }

    void ChunkV2PluginCreator::setPluginNamespace(const char* libNamespace)
    {
        _s_name_space = libNamespace;
    }

    const char* ChunkV2PluginCreator::getPluginNamespace() const
    {
        return _s_name_space.c_str();
    }

	REGISTER_TENSORRT_PLUGIN(ChunkV2PluginCreator);
}//namespace nvinfer1