aclnnRoiPoolingWithArgMax

产品支持情况

产品	是否支持
Ascend 950PR/Ascend 950DT	√
Atlas A3 训练系列产品/Atlas A3 推理系列产品	×
Atlas A2 训练系列产品/Atlas A2 推理系列产品	×
Atlas 200I/500 A2 推理产品	×
Atlas 推理系列产品	×
Atlas 训练系列产品	×

功能说明

接口功能：对输入特征图按 ROI（感兴趣区域）进行池化，在每个 ROI 内按空间划分为 pooled_h × pooled_w 个格子，对每个格子做最大池化，并输出池化结果及最大值在通道内的一维索引（argmax）。

计算公式：

输入特征图 $x$ 的 shape 为 $(N, C, H, W)$ ，ROI 张量 $rois\text{rois}$ 的 shape 为 $(num_rois,5)(\text{num\_rois}, 5)$ ，每行表示 $b_n, x_1, y_1, x_2, y_2)$ 。标量参数为 $s_h$ 、 $s_w$ （spatial_scale）以及 $pooled_h\text{pooled\_h}$ 、 $pooled_w\text{pooled\_w}$ 。下标 $n$ 表示 ROI 索引， $c$ 表示通道， $(ph,pw)(\text{ph}, \text{pw})$ 表示池化格点。

ROI 映射到特征图：将 ROI 坐标乘以 spatial_scale 得到特征图上的浮点区间：
$x~1=x1sw,y~1=y1sh,x~2=(x2+1)sw,y~2=(y2+1)sh\tilde{x}_1 = x_1 s_w,\quad \tilde{y}_1 = y_1 s_h,\quad \tilde{x}_2 = (x_2+1)s_w,\quad \tilde{y}_2 = (y_2+1)s_h$ $Wroi=x~2−x~1,Hroi=y~2−y~1W_{\text{roi}} = \tilde{x}_2 - \tilde{x}_1,\qquad H_{\text{roi}} = \tilde{y}_2 - \tilde{y}_1$
若 $Wroi≤0W_{\text{roi}} \le 0$ 或 $Hroi≤0H_{\text{roi}} \le 0$ ，该 ROI 的 $y$ 全为 0， $argmax\text{argmax}$ 全为 -1。
Bin 步长与区间：每个池化格 (ph, pw) 对应 ROI 内一个 bin，步长与浮点区间为：
$Δw=Wroipooled_w,Δh=Hroipooled_h\Delta w = \frac{W_{\text{roi}}}{\text{pooled\_w}},\qquad \Delta h = \frac{H_{\text{roi}}}{\text{pooled\_h}}$ $w~1=pw⋅Δw+x~1,w~2=(pw+1)⋅Δw+x~1\tilde{w}_1 = \text{pw} \cdot \Delta w + \tilde{x}_1,\quad \tilde{w}_2 = (\text{pw}+1) \cdot \Delta w + \tilde{x}_1$ $h~1=ph⋅Δh+y~1,h~2=(ph+1)⋅Δh+y~1\tilde{h}_1 = \text{ph} \cdot \Delta h + \tilde{y}_1,\quad \tilde{h}_2 = (\text{ph}+1) \cdot \Delta h + \tilde{y}_1$
取整并裁剪到 $\times [0,H)$ ：
$W)w_1 = \text{clip}(\lfloor\tilde{w}_1\rfloor,\, 0,\, W),\quad w_2 = \text{clip}(\lceil\tilde{w}_2\rceil,\, 0,\, W)$ $H)h_1 = \text{clip}(\lfloor\tilde{h}_1\rfloor,\, 0,\, H),\quad h_2 = \text{clip}(\lceil\tilde{h}_2\rceil,\, 0,\, H)$
其中 $clip(a,l,u)=min⁡(max⁡(a,l),u)\text{clip}(a,l,u) = \min(\max(a,l), u)$ 。若 $w2≤w1w_2 \le w_1$ 或 $h2≤h1h_2 \le h_1$ ，该 bin 为空： $y = 0$ ， $argmax=−1\text{argmax}=-1$ 。
池化输出与 Argmax：记 $\text{rois}[n,0]$ ，bin 区域 $\{(h,w) : h_1 \le h < h_2,\, w_1 \le w < w_2\}$ ，则
$y[n,c,ph,pw]=max⁡(h,w)∈Rx[b,c,h,w]y[n,c,\text{ph},\text{pw}] = \max_{(h,w) \in R} x[b,c,h,w]$
（空 $R$ 时为 0。）
$argmax[n,c,ph,pw]=h∗W+w∗\text{argmax}[n,c,\text{ph},\text{pw}] = h^* W + w^*$
$h^*, w^*)$ 为 bin 内最大值位置（多解取第一个）；空 $R$ 为 -1。

输出 Shape：

输出	Shape	数据类型
$y$	$(num_rois, C, pooled_h, pooled_w)(\text{num\_rois},\, C,\, \text{pooled\_h},\, \text{pooled\_w})$	与 $x$ 一致
$argmax\text{argmax}$	同上	INT32

函数原型

每个算子分为两段式接口，必须先调用“aclnnRoiPoolingWithArgMaxGetWorkspaceSize”接口获取计算所需 workspace 大小以及包含了算子计算流程的执行器，再调用“aclnnRoiPoolingWithArgMax”接口执行计算。

aclnnStatus aclnnRoiPoolingWithArgMaxGetWorkspaceSize(
  const aclTensor   *x,
  const aclTensor   *rois,
  int64_t            pooled_h,
  int64_t            pooled_w,
  float              spatial_scale_h,
  float              spatial_scale_w,
  aclTensor         *y,
  aclTensor         *argmax,
  uint64_t          *workspaceSize,
  aclOpExecutor    **executor)

aclnnStatus aclnnRoiPoolingWithArgMax(
  void          *workspace,
  uint64_t       workspaceSize,
  aclOpExecutor *executor,
  aclrtStream    stream)

aclnnRoiPoolingWithArgMaxGetWorkspaceSize

参数说明

参数名	输入/输出	描述	使用说明	数据类型	数据格式	维度(shape)	非连续Tensor
x（aclTensor*）	输入	输入特征图，格式为 NCHW，(N, C, H, W)。	不支持空 Tensor。输入维度必须为 4 维。 N需要为16的倍数。 N小于等于1024。	FLOAT32、FLOAT16	ND	4	√
rois（aclTensor*）	输入	ROI 框，每行 5 个元素：batch_idx, x1, y1, x2, y2。	shape 为（num_rois，5），不支持空 Tensor。第0维小于等于1024。 x1, y1, x2, y2大于等于0.0。	FLOAT32、FLOAT16	ND	2	√
pooled_h（int64_t）	输入	池化输出高度。	必须大于 0。	-	-	-	-
pooled_w（int64_t）	输入	池化输出宽度。	必须大于 0。	-	-	-	-
spatial_scale_h（float）	输入	ROI 坐标映射到特征图时在高度方向的缩放比例。	必须大于 0。	-	-	-	-
spatial_scale_w（float）	输入	ROI 坐标映射到特征图时在宽度方向的缩放比例。	必须大于 0。	-	-	-	-
y（aclTensor*）	输出	池化结果，shape 为（num_rois，C，pooled_h，pooled_w）。	不支持空 Tensor。数据类型与 x 一致。	FLOAT32、FLOAT16	ND	4	√
argmax（aclTensor*）	输出	每个池化格点最大值在通道内的线性偏移索引。	不支持空 Tensor。 shape 与 y 一致。	INT32	ND	4	√
workspaceSize（uint64_t*）	输出	返回需要在 Device 侧申请的 workspace 大小。	-	-	-	-	-
executor（aclOpExecutor**）	输出	返回 op 执行器，包含了算子计算流程。	-	-	-	-	-

返回值

aclnnStatus：返回状态码，具体参见aclnn返回码。

第一段接口完成入参校验，出现以下场景时报错：

返回码	错误码	描述
ACLNN_ERR_PARAM_NULLPTR	161001	如果传入参数是必选输入、输出或必选属性，且是空指针。
ACLNN_ERR_PARAM_INVALID	161002
		x、rois、y、argmax 的数据类型或格式不在支持范围内。
		x 的 shape 不是 4 维（NCHW）。
		rois 的 shape 第二维不是 5。
		pooled_h、pooled_w、spatial_scale_h、spatial_scale_w 不大于 0。

aclnnRoiPoolingWithArgMax

参数说明

参数名	输入/输出	描述
workspace	输入	在 Device 侧申请的 workspace 内存地址。
workspaceSize	输入	在 Device 侧申请的 workspace 大小，由第一段接口 aclnnRoiPoolingWithArgMaxGetWorkspaceSize 获取。
executor	输入	op 执行器，包含了算子计算流程。
stream	输入	指定执行任务的 Stream。

返回值

aclnnStatus：返回状态码，具体参见aclnn返回码。

约束说明

确定性计算：
- aclnnRoiPoolingWithArgMax 默认确定性实现。
x、argmax、rois的shape[0]相等

调用示例

示例代码如下，仅供参考，具体编译和执行过程请参考编译与运行样例。实际调用时需先通过 opgen 生成 aclnnop/aclnn_roi_pooling_with_arg_max.h，若生成的头文件或接口签名不同，请以生成接口为准。

#include <iostream>
#include <vector>
#include "acl/acl.h"
#include "aclnnop/aclnn_roi_pooling_with_arg_max.h"

#define CHECK_RET(cond, return_expr) \
  do {                               \
    if (!(cond)) {                   \
      return_expr;                   \
    }                                \
  } while (0)

#define LOG_PRINT(message, ...)     \
  do {                              \
    printf(message, ##__VA_ARGS__); \
  } while (0)

int64_t GetShapeSize(const std::vector<int64_t>& shape) {
  int64_t shape_size = 1;
  for (auto i : shape) {
    shape_size *= i;
  }
  return shape_size;
}

int Init(int32_t deviceId, aclrtStream* stream) {
  auto ret = aclInit(nullptr);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
  ret = aclrtSetDevice(deviceId);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
  ret = aclrtCreateStream(stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
  return 0;
}

template <typename T>
int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
                    aclDataType dataType, aclTensor** tensor) {
  auto size = GetShapeSize(shape) * sizeof(T);
  auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
  ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
  std::vector<int64_t> strides(shape.size(), 1);
  for (int64_t i = shape.size() - 2; i >= 0; i--) {
    strides[i] = shape[i + 1] * strides[i + 1];
  }
  *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
                            shape.data(), shape.size(), *deviceAddr);
  return 0;
}

template <typename T>
int CreateAclTensorOutput(const std::vector<int64_t>& shape, void** deviceAddr, aclDataType dataType,
                          aclTensor** tensor) {
  auto size = GetShapeSize(shape) * sizeof(T);
  auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
  std::vector<int64_t> strides(shape.size(), 1);
  for (int64_t i = shape.size() - 2; i >= 0; i--) {
    strides[i] = shape[i + 1] * strides[i + 1];
  }
  *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
                            shape.data(), shape.size(), *deviceAddr);
  return 0;
}

int main() {
  int32_t deviceId = 0;
  aclrtStream stream;
  auto ret = Init(deviceId, &stream);
  CHECK_RET(ret == 0, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);

  std::vector<int64_t> xShape = {2, 16, 25, 42};
  std::vector<int64_t> roisShape = {2, 5};
  std::vector<int64_t> yShape = {2, 16, 3, 3};
  std::vector<int64_t> argmaxShape = {2, 16, 3, 3};

  void* xDeviceAddr = nullptr;
  void* roisDeviceAddr = nullptr;
  void* yDeviceAddr = nullptr;
  void* argmaxDeviceAddr = nullptr;
  aclTensor* x = nullptr;
  aclTensor* rois = nullptr;
  aclTensor* y = nullptr;
  aclTensor* argmax = nullptr;

  int64_t xSize = GetShapeSize(xShape);
  int64_t roisSize = GetShapeSize(roisShape);
  std::vector<float> xHostData(xSize, 1.0f);
  std::vector<float> roisHostData(roisSize, 0.0f);
  roisHostData[0] = 0.0f;
  roisHostData[1] = 0.0f;
  roisHostData[2] = 0.0f;
  roisHostData[3] = 24.0f;
  roisHostData[4] = 41.0f;
  roisHostData[5] = 1.0f;
  roisHostData[6] = 0.0f;
  roisHostData[7] = 0.0f;
  roisHostData[8] = 24.0f;
  roisHostData[9] = 41.0f;

  ret = CreateAclTensor(xHostData, xShape, &xDeviceAddr, aclDataType::ACL_FLOAT, &x);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensor(roisHostData, roisShape, &roisDeviceAddr, aclDataType::ACL_FLOAT, &rois);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensorOutput<float>(yShape, &yDeviceAddr, aclDataType::ACL_FLOAT, &y);
  CHECK_RET(ret == ACL_SUCCESS, return ret);
  ret = CreateAclTensorOutput<int32_t>(argmaxShape, &argmaxDeviceAddr, aclDataType::ACL_INT32, &argmax);
  CHECK_RET(ret == ACL_SUCCESS, return ret);

  int64_t pooledH = 3;
  int64_t pooledW = 3;
  float spatialScaleH = 1.0f;
  float spatialScaleW = 1.0f;

  uint64_t workspaceSize = 0;
  aclOpExecutor* executor = nullptr;
  ret = aclnnRoiPoolingWithArgMaxGetWorkspaceSize(x, rois, pooledH, pooledW, spatialScaleH, spatialScaleW,
                                                  y, argmax, &workspaceSize, &executor);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnRoiPoolingWithArgMaxGetWorkspaceSize failed. ERROR: %d\n", ret);
            return ret);

  void* workspaceAddr = nullptr;
  if (workspaceSize > 0) {
    ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
  }

  ret = aclnnRoiPoolingWithArgMax(workspaceAddr, workspaceSize, executor, stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnRoiPoolingWithArgMax failed. ERROR: %d\n", ret); return ret);

  ret = aclrtSynchronizeStream(stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);

  int64_t yElem = GetShapeSize(yShape);
  std::vector<float> yResult(yElem, 0.0f);
  ret = aclrtMemcpy(yResult.data(), yElem * sizeof(float), yDeviceAddr, yElem * sizeof(float),
                    ACL_MEMCPY_DEVICE_TO_HOST);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy y from device to host failed. ERROR: %d\n", ret); return ret);

  int64_t argmaxElem = GetShapeSize(argmaxShape);
  std::vector<int32_t> argmaxResult(argmaxElem, 0);
  ret = aclrtMemcpy(argmaxResult.data(), argmaxElem * sizeof(int32_t), argmaxDeviceAddr,
                    argmaxElem * sizeof(int32_t), ACL_MEMCPY_DEVICE_TO_HOST);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy argmax from device to host failed. ERROR: %d\n", ret); return ret);

  aclDestroyTensor(x);
  aclDestroyTensor(rois);
  aclDestroyTensor(y);
  aclDestroyTensor(argmax);
  aclrtFree(xDeviceAddr);
  aclrtFree(roisDeviceAddr);
  aclrtFree(yDeviceAddr);
  aclrtFree(argmaxDeviceAddr);
  if (workspaceSize > 0) {
    aclrtFree(workspaceAddr);
  }
  aclrtDestroyStream(stream);
  aclrtResetDevice(deviceId);
  aclFinalize();
  return 0;
}