自定义层 API (TensorRT 2.1) 的简单示例？

Question

我正在使用TensorRT 2.1 并希望实现一个简单的自定义层。 （目标是在嵌入系统上使用TensorRT 运行Single Shot Detector。）

为了练习，我想制作一个Inc 层（只需将输入张量值加 1.0 并保持维度相同）。

我按照sampleFasterRNN.cpp 示例中的class Reshape : public Iplugin 实现Inc 类。除了getOutputDimensions() 以保持相同的尺寸外，我几乎保持了所有内容。 （这看起来不错。）

我应该在哪里实现“添加1.0”部分？我猜应该是在“enqueue()”中。所以，我尝试了

int enqueue(int batchSize, const void*const *inputs, void** outputs, void*, cudaStream_t stream) override
{
  # the below is from the Reshape class. seems to copy from input to output
  CHECK(cudaMemcpyAsync(outputs[0], inputs[0], mCopySize * batchSize, cudaMemcpyDeviceToDevice, stream));
  # add 1.0 to first ten values
  float* foutputs = (float*) outputs[0];
  int i; for (i = 0; i < 10; i++) foutputs[i] += 1.0;   
  return 0;
}

但是，这部分会导致“segmentation fault”错误。

我的问题是：

1. 在哪里以及如何在输入和输出之间实现一些计算？
2. 谁能提供一个简单的例子？

Answer 1

参考文件samples/samplePlugin/samplePlugin.cpp 并查看FCPlugin 类。您的实际计算应该进入enqueue 方法。您可能必须编写一个 CUDA 内核来执行增量。