如何解决TensorRT samples的BufferManager不支持有动态维度的engine的问题

最新推荐文章于 2024-08-28 10:21:15 发布

原创

最新推荐文章于 2024-08-28 10:21:15 发布 · 3.4k 阅读

4 ·

CC 4.0 BY-SA版权

文章标签：

#TensorRT #深度学习

博客详细分析了在使用TensorRT时遇到的内存分配异常问题，特别是在处理具有dynamic batch_size或高度/宽度维度的模型时。问题源于BufferManager类在创建缓冲区时，由于动态维度值为负数导致体积计算错误，进而引发内存分配过大，最终抛出`std::bad_alloc`异常。解决方案是对动态维度进行检查，将其转换为具体数值，避免计算错误。该博客还探讨了如何修改BufferManager构造函数以适应不同尺寸输入数据的推理需求。

TensorRT samples里common的代码是写得不错的用于简化调用TensorRT engine的套路的封装代码，使用这些封装类可以节省些代码，也使得代码更优雅点，但是里面有点问题，例如，有dynamic batch_size或者height/width维度的模型engine，在调用时会发生崩溃:

terminate called after throwing an instance of 'std::bad_alloc'
what(): std::bad_alloc
Aborted (core dumped)

追查了一下这个abort发生的原因是当有dynamic维度，例如batch_size=-1时，崩溃在/usr/src/tensorrt/samples/common/buffers.h里面的代码:

class BufferManager
{
public:
    static const size_t kINVALID_SIZE_VALUE = ~size_t(0);

    //!
    //! \brief Create a BufferManager for handling buffer interactions with engine.
    //!
    BufferManager(std::shared_ptr<nvinfer1::ICudaEngine> engine, const int batchSize = 0,
        const nvinfer1::IExecutionContext* context = nullptr)
        : mEngine(engine)
        , mBatchSize(batchSize)
    {
        // Full Dims implies no batch size.
        assert(engine->hasImplicitBatchDimension() || mBatchSize ==