cuda 与 cupy 和 tensorRT 的流同步问题

Posted 2023-02-16

【中文标题】cuda 与 cupy 和 tensorRT 的流同步问题

【英文标题】：Cuda streams synchronization issue with cupy and tensorRT

【发布时间】：2021-11-16 15:37:43

【问题描述】：

我正在使用 TensorRT 和 cupy。如果我设置了cp.cuda.Stream(non_blocking=True)，以下代码也不会等待执行 cuda 调用，而它与non_blocking=False 完美配合。 为什么它不能与non_blocking=True 一起使用？我检查了输入数据，没问题。但代码最终以我的模型返回随机检测（随机数据），这意味着存在一些同步问题。

        # Select stream
        stream.use()
        # Copy cupy array to the buffer
        input_images = cp.array(batch_input_image)
        cp.copyto(cuda_inputs[0], input_images)
        # Run inference.
        context.execute_async(bindings=bindings, stream_handle=stream.ptr, batch_size=len(batch_input_image))
        # Copy results from the buffer
        output_images = cuda_outputs[0].copy()
        # Split results into batch
        list_output = cp.split(output_images, indices_or_sections=len(batch_input_image), axis=0)
        # Squeeze output arrays to remove axis of length one
        list_output = [cp.squeeze(array) for array in list_output]
        # Synchronize the stream
        stream.synchronize()

【问题讨论】：

我对cupy不熟悉，但通常在CUDA中这个标志意味着流正在与“主”流同步。因此，如果一些操作在主流上运行，而另一些在另一个上运行，那么基于标志的不同行为是有意义的。

这可能是 TensorRT 和/或 bindings 的问题，如：github.com/cupy/cupy/issues/6104

【参考方案1】：

在获得 Nvidia 的一些支持后，我可以确认这不是 cupy 问题。 TensorRT 模型的 C++ 代码似乎有问题，如下所述：github.com/cupy/cupy/issues/6104。