[Вывод модели] Метод реализации плагина пользовательского оператора TensorRT8

искусственный интеллект алгоритм
[Вывод модели] Метод реализации плагина пользовательского оператора TensorRT8

  «Это первый день моего участия в первом испытании обновлений 2022 года. Подробную информацию о мероприятии см.:Вызов первого обновления 2022 г."

 欢迎关注我的公众号 [极智视界]

  Всем привет, меня зовут Jizhi Vision В этой статье представлен метод реализации подключаемого модуля пользовательского оператора TensorRT8.

  Когда мы ускоряем модель глубокого обучения с помощью Nvidia TensorRT, нам необходимо сериализовать модель из исходной модели, такой как pytorch, tensorflow, caffe или darkent, в TensorRT в структуру модели Engine, а затем десериализовать ее для вывода. Сложность в этом процессе заключается в том, что реализация новых операторов TensorRT часто требует рукописного ввода.Обычно существует два метода рукописного ввода операторов:

  (1) Использовать детализированные функции, уже поддерживаемые TensorRT для сплайсинга, и объединить их в необходимые новые операторы;

  (2) Используйте плагин для разработки пользовательских операторов, и операторы реализованы в cuda c;

  В этом разделе в основном представлен метод использования плагина для разработки пользовательских операторов в TensorRT 8. Поскольку TensorRT8 внес серьезные изменения в разработку плагина, его необходимо представить.

  Чтобы реализовать плагин в TensorRT8, вам нужно создать два класса: custom_plugin, create_plugin и вызвать custom_plugin в create_plugin. Начните ниже.

1. Создайте custom_plugin

  Класс custom_plugin наследуется от IPluginV2Ext, IPluginV2 или IPluginV2DynamicExt.

class ClipPlugin : public IPluginV2
{
public:
    ClipPlugin(const std::string name, float clipMin, float clipMax);
    ClipPlugin(const std::string name, const void* data, size_t length);
    // ClipPlugin 不带参是没有意义的，所以删除了默认构造.
    ClipPlugin() = delete;
    int getNbOutputs() const noexcept override;
    Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) noexcept override;
    int initialize() noexcept override;
    void terminate() noexcept override;
    size_t getWorkspaceSize(int) const noexcept override
    {
        return 0;
    };
    int enqueue(int batchSize, const void* const* inputs, void* const* outputs, void* workspace,
        cudaStream_t stream) noexcept override;   // 通过 cuda c 实现
    size_t getSerializationSize() const noexcept override;
    void serialize(void* buffer) const noexcept override;
    void configureWithFormat(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, DataType type,
        PluginFormat format, int maxBatchSize) noexcept override;
    bool supportsFormat(DataType type, PluginFormat format) const noexcept override;
    const char* getPluginType() const noexcept override;
    const char* getPluginVersion() const noexcept override;
 
    void destroy() noexcept override;
    nvinfer1::IPluginV2* clone() const noexcept override;   // addPluginV2时会调用，该方法会去调用构造函数
    void setPluginNamespace(const char* pluginNamespace) noexcept override;
    const char* getPluginNamespace() const noexcept override;
private:   // 算子参数
    const std::string mLayerName;
    float mClipMin, mClipMax;
    size_t mInputVolume;
    std::string mNamespace;
};

2. Создайте create_plugin

class ClipPluginCreator : public IPluginCreator
{
public:
    ClipPluginCreator();
    const char* getPluginName() const noexcept override;
    const char* getPluginVersion() const noexcept override;
    const PluginFieldCollection* getFieldNames() noexcept override;
    IPluginV2* createPlugin(const char* name, const PluginFieldCollection* fc) noexcept override; //实现该方法时调用 custom_plugin
    IPluginV2* deserializePlugin(const char* name, const void* serialData, size_t serialLength) noexcept override;
    void setPluginNamespace(const char* pluginNamespace) noexcept override;
    const char* getPluginNamespace() const noexcept override;
private:
    static PluginFieldCollection mFC;
    static std::vector<PluginField> mPluginAttributes;
    std::string mNamespace;
};

3. Звонок

  Официальным методом является вызов custom_plugin через create_plugin, что по сути является вызовом custom_plugin, который можно записать следующим образом:

nvinfer1::IPluginV2 *clip = new ClipPlugin(scale, 512, Dtype);
nvinfer1::IPluginV2Layer *Clip = m_network->addPluginV2(&Layers[inputName], 1, *clip);

  где ClipPlugin — это custom_plugin.

  Выше был описан метод реализации плагина пользовательского оператора в TensorRT 8. Я надеюсь, что мой обмен поможет вам в вашем обучении.

【Передача по общему номеру】 "[Вывод модели] Метод реализации плагина пользовательского оператора TensorRT8》

Категории