[TensorRT] Write your plugin which support FP16

最新推荐文章于 2023-08-25 11:13:51 发布

xuke_2022

最新推荐文章于 2023-08-25 11:13:51 发布

阅读量694

点赞数

分类专栏：算法文章标签：机器学习

本文链接：https://blog.youkuaiyun.com/xuke_2018/article/details/105292149

版权

本文详细介绍了如何在TensorRT中编写支持FP16的插件，包括在插件中支持FP16的3个关键接口，使用cuda_fp16.h库编写内核函数，以及在引擎中应用FP16的步骤。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

Total Steps 3:

1. In the plugin, 3 Interface need support FP16

2. Use cuda_fp16.h to write Kernel Function

3. Use FP16 in the Engine

===========================================

1. In the plugin, 3 Interface need support FP16

class FCPlugin : public nvinfer1::IPluginExt
{
public:
    ...

    bool supportsFormat(nvinfer1::DataType type, nvinfer1::PluginFormat format) const override
    {
        int device;
        CHECK(cudaGetDevice(&device));
        cudaDeviceProp props{};
        cudaGetDeviceProperties(&props, device);
        int smVersion = props.major << 8 | props.minor;
        // Half p