CANN ops-cv UpsampleTrilinear3dBackward 算子深度解析:三维三线性上采样的反向梯度计算与 aclnn 接口实战
【免费下载链接】ops-cv本项目是CANN提供的图像处理、目标检测相关的算子库,实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-cv
本篇技术指南以 CANN ops-cv 开源算子库中的 UpsampleTrilinear3dBackward 算子文档 为核心,系统讲解三维三线性插值(Trilinear Interpolation)反向算子的数学原理、参数约束、两段式 aclnn 调用方式,并结合该算子目录下的源码与测试用例剖析其在 NPU 上的实现机制。读者读完将能够独立完成该算子的 aclnn 接口调用、理解其梯度传播公式,并掌握其 host 侧算子定义、InferShape、Tiling 与 kernel 侧 matmul 实现的全链路结构。
产品支持情况
根据 README.md 中列出的产品支持矩阵,UpsampleTrilinear3dBackward 算子的支持情况如下:
| 产品 | 是否支持 |
|---|---|
| Ascend 950PR/Ascend 950DT | × |
| Atlas A3 训练系列产品/Atlas A3 推理系列产品 | √ |
| Atlas A2 训练系列产品/Atlas A2 推理系列产品 | √ |
| Atlas 200I/500 A2 推理产品 | × |
| Atlas 推理系列产品 | × |
| Atlas 训练系列产品 | × |
此外,在算子对应的 aclnnUpsampleTrilinear3dBackward 接口文档 中也列示了各产品的支持情况,两处文档在部分产品(如 Atlas 推理系列、Atlas 训练系列)上的表述存在差异,实际使用时建议以与自身环境匹配的版本化文档为准。从 host 侧算子注册文件 upsample_trilinear3d_backward_def.cpp 可以看到,该算子当前注册的 AICore 配置为ascend910b与ascend910_93两个平台。
功能说明:三维三线性上采样反向算子
UpsampleTrilinear3dBackward 是 ResizeUpsampleTrilinear 算子(即 aclnnUpsampleTrilinear3d 接口)的反向计算算子。其输入为正向算子输出的梯度gradOutput,输出为反向传播得到的输入侧梯度gradInput,shape 与正向算子的输入一致。该算子通常用于深度学习训练链路中F.interpolate(..., mode='trilinear')这类三维特征图缩放操作的自动求导。
正向核心算法逻辑
要理解反向计算,首先要弄清正向三线性插值的计算过程,其核心逻辑分为三步:
- 将目标图像的每一个点映射回原图,得到一个带小数点的坐标;
- 根据这个浮点数坐标,计算前后相邻的原始图像的点;
- 分别计算相邻点到对应目标点的权重,按照权重相乘累加即可得到目标点值。
缩放系数与坐标映射公式
缩放方式分为角对齐(alignCorners 为 true)与边对齐(alignCorners 为 false)两种:角对齐表示按照原始图片左上角像素中心点对齐;边对齐表示按照原始图片左上角顶点及两条边对齐。二者在缩放系数和坐标位置的计算上存在差异。对于三维插值点 $(N, C, D, H, W)$,三个维度的缩放系数定义如下:
$$ scale_d =\begin{cases} (inputSize[2]-1) / (outputSize[0]-1) & alignCorners=true \ 1 / scales_d & alignCorners=false&scales_d>0\ inputSize[2] / outputSize[0] & alignCorners=false \end{cases} $$
$$ scale_h =\begin{cases} (inputSize[3]-1) / (outputSize[1]-1) & alignCorners=true \ 1 / scales_h & alignCorners=false&scales_h>0\ inputSize[3] / outputSize[1] & alignCorners=false \end{cases} $$
$$ scale_w =\begin{cases} (inputSize[4]-1) / (outputSize[2]-1) & alignCorners=true \ 1 / scales_w & alignCorners=false&scales_w>0\ inputSize[4] / outputSize[2] & alignCorners=false \end{cases} $$
其中inputSize是正向输入在 D、H、W 维度上的大小,outputSize是正向输出在 D、H、W 维度上的大小(对应到反向场景中,分别是gradInput与gradOutput的空间尺寸),scales_d/scales_h/scales_w为用户显式指定的缩放因子。
对于输出(反向场景下即 gradOutput)某个方向上的点 p(x, y, z),映射回原始图像中的点记为 q(x', y', z'),则有:
$$ x' =\begin{cases} x * scale_d & alignCorners=true \ MAX(0,{(x+0.5)*scale_d-0.5}) & alignCorners=false \end{cases} $$
$$ y' =\begin{cases} y * scale_h & alignCorners=true \ MAX(0,{(y+0.5)*scale_h-0.5}) & alignCorners=false \end{cases} $$
$$ z' =\begin{cases} z * scale_w & alignCorners=true \ MAX(0,{(z+0.5)*scale_w-0.5}) & alignCorners=false \end{cases} $$
三线性插值:八个相邻点的加权累加
得到浮点坐标后,取其整数部分与整数部分加一,得到三个方向上各自的前后相邻点索引,并计算相应权重:
$$ x_{0} =int(x'),x_{1} =int(x')+1, lambda_{0} = x_{1}-x', lambda_{1} = 1-lambda_{0} $$
$$ y_{0} =int(y'),y_{1} =int(y')+1, lambdb_{0} = y_{1}-y', lambdb_{1} = 1-lambdb_{0} $$
$$ z_{0} =int(z'),z_{1} =int(z')+1, lambdc_{0} = z_{1}-z', lambdc_{1} = 1-lambdc_{0} $$
目标点像素值由包围它的 8 个原始点按三维权重加权求和得到:
$$ {V(p_{x, y, z})} = {V(p_{x0, y0, z0})} * {lambda_{0}} * {lambdb_{0}} * {lambdc_{0}} + {V(p_{x0, y0, z1})} * {lambda_{0}} * {lambdb_{0}} * {lambdc_{1}} + {V(p_{x0, y1, z0})} * {lambda_{0}} * {lambdb_{1}} * {lambdc_{0}} + {V(p_{x0, y1, z1})} * {lambda_{0}} * {lambdb_{1}} * {lambdc_{1}} + {V(p_{x1, y0, z0})} * {lambda_{1}} * {lambdb_{0}} * {lambdc_{0}} + {V(p_{x1, y0, z1})} * {lambda_{1}} * {lambdb_{0}} * {lambdc_{1}} + {V(p_{x1, y1, z0})} * {lambda_{1}} * {lambdb_{1}} * {lambdc_{0}} + {V(p_{x1, y1, z1})} * {lambda_{1}} * {lambdb_{1}} * {lambdc_{1}} $$
反向梯度传播公式
反向计算的核心思路是:将 gradOutput 中每个点的梯度,按正向插值时的权重反哺给其参与插值的 8 个原始点。假设正向插值的输出图像 out(x, y, z) 受原图像 input(x_i, y_j, z_k) 影响,则:
$$ gradInput(x_i,y_j,z_k) += gradOutput(x,y,z) * lambda(x_i,y_j,z_k)* lambdb(x_i,y_j,z_k)* lambdc(x_i,y_j,z_k) $$
即每个原始点累积所有“以它为插值源”的目标点梯度乘以其对应的三轴权重。由于一个原始点可能被多个输出点引用,因此梯度使用+=累加语义。这一公式在 kernel 侧被进一步转化为三个方向的连续矩阵乘(matmul)流水实现,具体见后文源码剖析。
参数说明
算子层面的输入、输出与属性定义如下表(对应 README.md 的参数说明章节):
| 参数名 | 输入/输出/属性 | 描述 | 数据类型 | 数据格式 |
|---|---|---|---|---|
| grad_output | 输入 | 表示反向计算的梯度 Tensor,对应公式中的gradOutput。 | FLOAT32、FLOAT16、DOUBLE | NCDHW |
| input_size | 属性 | 表示输出gradInput分别在 N、C、D、H 和 W 维度上的空间大小,对应公式中的inputSize。size 为 5,且各元素均大于零。 | LISTINT | - |
| output_size | 属性 | 表示输入grad_output在 D、H 和 W 维度上的空间大小,对应公式中的outputSize。size 为 3,且各元素均大于零。 | LISTINT | - |
| scales | 可选属性 | 指定沿每个维度的缩放数组,包含 3 个元素:scale_depth、scale_height、scale_width,对应公式中的scales_d、scales_h、scales_w。默认值为空。 | LISTFLOAT | - |
| align_corners | 可选属性 | 决定是否对齐角像素点,对应公式中的alignCorners。为 true 时输入和输出张量的角像素点会被对齐,否则不对齐。默认值为 false。 | BOOL | - |
| y | 输出 | 表示反向计算的输出张量,对应公式中的gradInput。数据类型和数据格式与入参grad_output保持一致。shape 依赖input_size和output_size/scales。 | FLOAT32、FLOAT16、DOUBLE | NCDHW |
在 host 侧算子定义文件 upsample_trilinear3d_backward_def.cpp 中,上述参数被注册为:
grad_output为必选输入,支持DT_BF16、DT_FLOAT16、DT_FLOAT三种数据类型(注意此处与算子文档表格中列的 DOUBLE 并不一致,实际支持的以注册表为准);output_size、input_size为必选属性(ListInt);align_corners为可选属性,默认 false;scale_d、scale_h、scale_w为可选属性(Float);grad_input为必选输出,数据类型与输入一致。
约束说明:README 中明确该算子无额外约束。但在 aclnn 接口层面存在一系列 shape 与内存约束,详见下文「接口约束」小节。
aclnn 两段式接口
CANN 的 aclnn 算子接口采用「两段式」调用模型(参见 two_phase_api 说明):必须先调用aclnnUpsampleTrilinear3dBackwardGetWorkspaceSize获取计算所需的 workspace 大小以及封装了算子计算流程的执行器,再调用aclnnUpsampleTrilinear3dBackward执行计算。
函数原型
aclnnStatus aclnnUpsampleTrilinear3dBackwardGetWorkspaceSize( const aclTensor *gradOut, const aclIntArray *outputSize, const aclIntArray *inputSize, bool alignCorners, double scalesD, double scalesH, double scalesW, aclTensor *gradInput, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnUpsampleTrilinear3dBackward( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)GetWorkspaceSize 参数说明
| 参数名 | 输入/输出 | 描述 | 使用说明 | 数据类型 | 数据格式 | 维度(shape) | 非连续Tensor |
|---|---|---|---|---|---|---|---|
| gradOut(aclTensor*) | 输入 | 表示反向计算的梯度 Tensor,对应公式中的gradOut。 | 不支持空 Tensor;当数据格式为 ND 时,默认按照 NCDHW 格式处理。 | FLOAT32、FLOAT16、BFLOAT16、DOUBLE | NCDHW、NDHWC、ND | 5 | √ |
| outputSize(aclIntArray*) | 输入 | 表示输入gradOut在 D、H 和 W 维度上的空间大小,对应公式中的outputSize。 | size 为 3,且各元素均大于零。 | INT64 | - | - | - |
| inputSize(aclIntArray*) | 输入 | 表示输出gradInput分别在 N、C、D、H 和 W 维度上的空间大小,对应公式中的inputSize。 | size 为 5,且最后三个元素均大于零。 | INT64 | - | - | - |
| alignCorners(bool) | 输入 | 表示是否对齐角像素点,对应公式中的alignCorners。 | 为 true 则输入和输出张量的角像素点会对齐,否则不对齐。 | - | - | - | - |
| scalesD(double) | 输入 | 表示输出gradInput的 depth 维度乘数,对应公式中的scales_d。 | - | - | - | - | - |
| scalesH(double) | 输入 | 表示输出gradInput的 height 维度乘数,对应公式中的scales_h。 | - | - | - | - | - |
| scalesW(double) | 输入 | 表示输出gradInput的 width 维度乘数,对应公式中的scales_w。 | - | - | - | - | - |
| gradInput(aclTensor*) | 输出 | 表示反向计算的输出张量,对应公式中的gradInput。 | 不支持空 Tensor;shape 在 N、C、D、H、W 上的大小需与inputSize一致;数据类型、数据格式、shape 与入参gradOut保持一致。 | FLOAT32、FLOAT16、BFLOAT16、DOUBLE | NCDHW、NDHWC、ND | 5 | √ |
| workspaceSize(uint64_t*) | 输出 | 返回需要在 Device 侧申请的 workspace 大小。 | - | - | - | - | - |
| executor(aclOpExecutor**) | 输出 | 返回 op 执行器,包含了算子计算流程。 | - | - | - | - | - |
需要注意的是,接口文档中注明:Atlas 训练系列产品上,参数gradOut和gradInput的数据类型不支持 BFLOAT16。
返回值与错误码
第一段接口完成入参校验,出现以下场景时报错:
| 返回码 | 错误码 | 描述 |
|---|---|---|
| ACLNN_ERR_PARAM_NULLPTR | 161001 | 传入的 gradOut、outputSize、inputSize 或 gradInput 是空指针。 |
| ACLNN_ERR_PARAM_INVALID | 161002 | gradOut 的数据类型和数据格式不在支持的范围内;或 gradOut 与 gradInput 数据类型不一致;或 gradOut 维度不为 5 维;或 outputSize 的 size 不等于 3;或 outputSize 某元素不大于 0;或 inputSize 的 size 不等于 5;或 inputSize 最后三个元素中某个不大于 0;或 gradOut 与 inputSize 在 N、C 维度上的 size 不同;或 gradOut 在 D、H、W 维度上的 size 与 outputSize[0..2] 不一致;或 gradInput 的 shape 与 inputSize[0..4] 不一致。 |
这些校验逻辑在 aclnn_upsample_trilinear_3d_backward.cpp 的CheckParams中实现:依次检查空指针(CheckNotNull2In2Out)、数据类型合法性(CheckDtypeValid1Out1In)、shape 与 format(CheckUpsampleShape、CheckFormat)、输入元素合法性(CheckInputElement,包括各维 size 大于 0、gradOut/gradInput 的 shape 与 inputSize/outputSize 完全匹配)以及上边界(CheckUplimit,各维不得超过 INT32_MAX)。CheckFormat仅接受 ND、NCDHW、NDHWC 三种存储格式(源码 L35-L44)。
接口约束
shape 约束:
gradOut、gradInput每个维度的取值小于等于 2^20;gradInput的 N 轴和 C 轴与gradOut保持一致;内存占用需小于 60GB。内存占用的计算公式为:$$ N * C * (gradOut_D * gradOut_H * gradOut_W + gradInput_D * gradInput_H * gradInput_W + gradOut_D * gradOut_H * gradInput_W + gradOut_D * gradInput_H * gradInput_W) * sizeof(dtype) < 60 * 1024 * 1024 * 1024 $$
其中 N、C 分别为输入输出的 N 轴与 C 轴,dtype 为输入张量的数据类型。另有
N * C * gradOut_D * gradOut_H < 2^31与gradInput_W * gradInput_H < 2^31两条约束。scales 约束:当输入 scalesD、scalesH、scalesW 取值均大于 0 时,参数 inputSize、outputSize、scales 需满足:
$$ outputSize_D = floor(inputSize_D * scalesD), \quad outputSize_H = floor(inputSize_H * scalesH), \quad outputSize_W = floor(inputSize_W * scalesW) $$
这一“output_size 与 scales 二选一”的约束在 InferShape 中同样有体现:在 upsample_trilinear3d_backward_infershape.cpp 中,若
output_size非空且scales为空,则校验 output_size 的三个元素必须等于 grad_output 的 D、H、W 维;若output_size为空且scales非空,则校验floor(input_size[i] * scales[i])必须等于 grad_output 对应维度;两者同时为空或同时非空则直接报错。确定性计算:aclnnUpsampleTrilinear3dBackward 默认为确定性实现,多次执行结果可复现。
调用示例
仓库中提供了完整的可运行样例 test_aclnn_upsample_trilinear3d_backward.cpp,其编译与执行过程可参考 编译与运行样例。核心调用流程如下:
#include <iostream> #include <vector> #include "acl/acl.h" #include "aclnnop/aclnn_upsample_trilinear_3d_backward.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); // 调用aclrtMalloc申请device侧内存 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); // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上 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); // 计算连续tensor的strides 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]; } // 调用aclCreateTensor接口创建aclTensor *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_NCDHW, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { // 1. (固定写法)device/stream初始化,参考acl API手册 // 根据自己的实际device填写deviceId 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); // 2. 构造输入与输出,需要根据API的接口自定义构造 std::vector<int64_t> gradOutShape = {2, 2, 2, 2, 2}; std::vector<int64_t> gradInputShape = {2, 2, 1, 1, 1}; void* gradOutDeviceAddr = nullptr; void* gradInputDeviceAddr = nullptr; aclTensor* gradOut = nullptr; aclTensor* gradInput = nullptr; std::vector<float> gradOutHostData = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32}; std::vector<float> gradInputHostData = {2.0, 2, 2, 2}; std::vector<int64_t> outputSizeData = {2, 2, 2}; std::vector<int64_t> inputSizeData = {2, 2, 1, 1, 1}; bool alignCorners = false; double scalesD = 0.0; double scalesH = 0.0; double scalesW = 0.0; // 创建gradOut aclTensor ret = CreateAclTensor(gradOutHostData, gradOutShape, &gradOutDeviceAddr, aclDataType::ACL_FLOAT, &gradOut); CHECK_RET(ret == ACL_SUCCESS, return ret); // 创建gradInput aclTensor ret = CreateAclTensor(gradInputHostData, gradInputShape, &gradInputDeviceAddr, aclDataType::ACL_FLOAT, &gradInput); CHECK_RET(ret == ACL_SUCCESS, return ret); const aclIntArray* outputSize = aclCreateIntArray(outputSizeData.data(), outputSizeData.size()); CHECK_RET(outputSize != nullptr, return ACL_ERROR_INTERNAL_ERROR); const aclIntArray* inputSize = aclCreateIntArray(inputSizeData.data(), inputSizeData.size()); CHECK_RET(inputSize != nullptr, return ACL_ERROR_INTERNAL_ERROR); // 3. 调用CANN算子库API,需要修改为具体的API uint64_t workspaceSize = 0; aclOpExecutor* executor; // 调用aclnnUpsampleTrilinear3dBackward第一段接口 ret = aclnnUpsampleTrilinear3dBackwardGetWorkspaceSize(gradOut, outputSize, inputSize, alignCorners, scalesD, scalesH, scalesW, gradInput, &workspaceSize, &executor); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnUpsampleTrilinear3dBackwardGetWorkspaceSize failed. ERROR: %d\n", ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 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); } // 调用aclnnUpsampleTrilinear3dBackward第二段接口 ret = aclnnUpsampleTrilinear3dBackward(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnUpsampleTrilinear3dBackward failed. ERROR: %d\n", ret); return ret); // 4. (固定写法)同步等待任务执行结束 ret = aclrtSynchronizeStream(stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret); // 5. 获取输出的值,将device侧内存上的结果拷贝至host侧,需要根据具体API的接口定义修改 auto size = GetShapeSize(gradInputShape); std::vector<float> resultData(size, 0); ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), gradInputDeviceAddr, size * sizeof(float), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret); for (int64_t i = 0; i < size; i++) { LOG_PRINT("result[%ld] is: %f\n", i, resultData[i]); } // 6. 释放aclTensor和aclScalar,需要根据具体API的接口定义修改 aclDestroyTensor(gradOut); aclDestroyTensor(gradInput); // 7. 释放device资源,需要根据具体API的接口定义修改 aclrtFree(gradOutDeviceAddr); aclrtFree(gradInputDeviceAddr); if (workspaceSize > 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例的调用流程可归纳为七步:设备与流初始化 → 构造输入输出 aclTensor 与 aclIntArray → 调用第一段接口获取 workspace 与 executor → 申请 workspace 设备内存 → 调用第二段接口执行计算 → 同步流并回拷结果 → 释放资源。示例中gradOutshape 为{2, 2, 2, 2, 2}、gradInputshape 为{2, 2, 1, 1, 1},对应一次2x2x2 → 1x1x1的下采样反向,alignCorners=false且三个 scales 传 0(此时完全依据 inputSize/outputSize 计算缩放系数)。
源码级实现剖析
1. 接口层:从 NDHWC 到 NCDHW 的自动转置
在 aclnn_upsample_trilinear_3d_backward.cpp 中,第一段接口完成参数校验后会执行以下关键逻辑:
- 对
gradOut调用l0op::Contiguous做连续性整理; - 若
scalesD > 0 && scalesH > 0 && scalesW > 0,则把三个 scale 组装进scalesList,并把outputSize置为空数组(L203-L209),即「显式指定 scales 时不再需要 outputSize」,与 InferShape 的二选一约束完全对应;同时生成castScales数组保存1/scalesD、1/scalesH、1/scalesW,对应公式中scale = 1/scales的分支; - 核心计算函数
upsampleTrilinear3dBackwardCompute(L149-L174)中,若gradOut为 NDHWC 格式,会先按{0, 4, 1, 2, 3}置换为 NCDHW 调用底层l0op::UpsampleTrilinear3dGradNcdhw,再按{0, 2, 3, 4, 1}置换回 NDHWC 输出;ND 与 NCDHW 格式则直接计算; - 最后通过
l0op::ViewCopy将中间结果写入用户提供的gradInput,并将整个执行图所需的 workspace 大小通过executor->GetWorkspaceSize()返回(L230-L235)。
第二段接口aclnnUpsampleTrilinear3dBackward则直接调用CommonOpExecutorRun(workspace, workspaceSize, executor, stream)完成图执行。
2. InferShape:输出 shape 推导与约束校验
upsample_trilinear3d_backward_infershape.cpp 中通过IMPL_OP_INFERSHAPE(UpsampleTrilinear3dGrad)注册 shape 推导:要求 grad_output 为 5 维、input_size 的 size 必须为 5,随后按上文所述的「output_size 与 scales 二选一」规则校验 grad_output 的 D、H、W 维,最终输出 shape 的 N、C、D、H、W 全部取自input_size。
3. Tiling:分核策略与 workspace 规划
Tiling 逻辑位于 upsample_trilinear3d_backward_tiling_func.cpp,其中几个关键常量与算法值得注意:
SLIDE_SIZE = 16:每个核心按 16 为滑动步长切分数据;MIN_SUPPORT_SCALE = 0.02f:三个方向实际缩放系数(ComputeScales)不得小于 0.02,否则 tiling 失败;AreaPixelComputeScale(L228-L239)与数学公式一致:outputSize == inputSize时缩放系数为 1(该方向无需 resize),alignCorners时取(inputSize-1)/(outputSize-1),否则取scale或inputSize/outputSize;- 三个方向(D/H/W)分别计算各自的单核 K 维度、分核数量与 matmul tiling,最终
needCoreNum取三个方向的最大值(L241-L267); - workspace 由「W 方向中间矩阵 + H 方向中间矩阵 + 各核权重矩阵 + 32MB 预留空间」组成(GetWorkSpace),其中权重矩阵大小为
SLIDE_SIZE * singleCoreK,singleCoreK = min(ceil((16+4)/scale)+4, outputSize); - Tiling 数据通过 upsample_trilinear3d_backward_tiling.h 中
BEGIN_TILING_DATA_DEF定义的结构体传递,包含数据类型、batch、三维输入输出 shape、三轴 scale、alignCorners、各方向是否需要 resize、滑动/分核参数以及 W/H/D 三个TCubeTiling。
4. Kernel:三方向 matmul 流水
kernel 入口 upsample_trilinear3d_backward.cpp 按dataType(half / float / bfloat16)实例化UpsampleTrilinear3dBackwardND<T>模板类。该类定义于 upsample_trilinear3d_backward.h,其核心设计是:
- 将「梯度按权重散射回原图」转化为W → H → D 三个方向的连续 matmul,类内持有
matmulW、matmulH、matmulD三个 Cube 侧 matmul 对象; CalculateRadioTensor计算每个方向的权重矩阵(即各输出点对应的 lambda/lambdb/lambdc),坐标映射逻辑AreaPixelComputeSourceIndex(L256-L263)与公式完全对应:alignCorners 时scale * dstIndex,否则max(scale*(dst+0.5)-0.5, 0);Process(L141-L162)按 W → H → D 顺序依次执行DirectionExpansion,每步之间以SyncAll同步多核,中间结果写入 GM workspace;某方向无需 resize(needResizeX == false)时跳过该方向的 matmul,并把上一阶段结果直接作为输出或下一阶段输入;- 在只对一个方向 resize 的简化场景下,可直接将 matmul 结果写入最终输出
outTensorsGM,避免额外的中间矩阵搬运。
整体上,该 kernel 利用 Cube 算力将逐点插值权重计算转化为矩阵乘,属于「权重矩阵 × 数据矩阵」的经典算子实现范式。
测试与验证
仓库为该算子提供了多层次的测试支撑:
- ST 级用例:在 atk_aclnnUpsampleTrilinear3dBackward.json 中配置了上百组测试场景,对应 PyTorch 的
torch.nn.functional.interpolate(mode='trilinear'),覆盖 bf16/fp16 多种数据类型、align_corners两种取值、多样化的 shape(含[1,64,8,128,128]等大 shape)以及 ±inf 等边界输入值,精度标准采用double_benchmark双精度基准比对; - UT 级用例:
tests/ut/op_host/op_api/test_aclnn_upsample_trilinear_3d_backward.cpp验证接口层行为,tests/ut/op_host/test_upsample_trilinear3d_backward_tiling.cpp验证 tiling 数据生成,tests/ut/op_host/test_upsample_trilinear3d_grad_infershape.cpp验证 shape 推导,tests/ut/op_kernel/test_upsample_trilinear3d_backward.cpp配合gen_data.py/compare_data.py完成 kernel 数值比对。
小结
UpsampleTrilinear3dBackward 是 CANN ops-cv 图像算子库中三维三线性上采样(F.interpolate(..., mode='trilinear'))的反向梯度算子。本文从产品支持、数学原理(角对齐/边对齐缩放、8 邻域加权插值、梯度散射公式)、算子参数、aclnn 两段式接口与错误码、完整调用示例,到 host 侧算子定义 / InferShape / Tiling 与 kernel 侧三方向 matmul 流水做了全链路解析。开发者可直接参照 测试样例 完成调用,并依据 接口文档 中的 shape、内存与 scales 约束合理设计输入规模。
【免费下载链接】ops-cv本项目是CANN提供的图像处理、目标检测相关的算子库,实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-cv
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考