手把手教你重写grid_sample函数：当PyTorch转ONNX连mmcv都救不了的时候

从零实现grid_sample：破解PyTorch转ONNX时的算子兼容困局

当你在深夜的显示器前看到那个刺眼的RuntimeError: Exporting the operator grid_sampler to ONNX opset version 11 is not supported时，作为中高级开发者的直觉告诉你——这绝不是换个API调用就能解决的简单问题。特别是在生产环境中，当torch版本被锁定、mmcv库因架构问题无法加载时，我们需要更底层的解决方案。

1. 理解grid_sample的数学本质

grid_sample的核心是双线性插值——这个在计算机图形学中广泛应用的技术，其数学原理远比表面看起来的精妙。想象一下，当我们需要从一个非整数坐标位置采样时，实际上是取其周围四个最近像素的加权平均值。

双线性插值的计算步骤：

1. 将归一化坐标(-1到1)转换为实际像素坐标
2. 确定目标点周围的四个整数坐标点(Q11, Q12, Q21, Q22)
3. 计算目标点与这四个点的相对位置权重
4. 进行加权求和得到最终采样值

def normalize_coordinates(grid, h, w, align_corners): if align_corners: # 角点对齐模式下的坐标转换 x = ((grid[..., 0] + 1) / 2) * (w - 1) y = ((grid[..., 1] + 1) / 2) * (h - 1) else: # 非角点对齐模式 x = ((grid[..., 0] + 1) * w - 1) / 2 y = ((grid[..., 1] + 1) * h - 1) / 2 return x, y

注意：align_corners参数会显著影响边界像素的处理方式。当为True时，-1和1正好对应边界像素的中心；为False时则对应边界像素的边缘。

2. 构建可ONNX导出的采样核心

PyTorch原生的grid_sample无法导出到某些ONNX opset版本，主要是因为其内部实现使用了某些特定操作。我们需要用基础张量操作重建这个过程。

关键步骤分解：

1. 坐标处理：将输入网格坐标转换为实际图像坐标
2. 边界处理：处理超出图像范围的坐标
3. 权重计算：计算四个邻近点的双线性权重
4. 值收集：使用torch.gather获取邻近点像素值
5. 加权求和：组合权重和像素值得到最终结果

def compute_bilinear_weights(x, y): x0 = torch.floor(x).long() y0 = torch.floor(y).long() x1 = x0 + 1 y1 = y0 + 1 # 计算四个权重分量 wa = (x1 - x) * (y1 - y) wb = (x1 - x) * (y - y0) wc = (x - x0) * (y1 - y) wd = (x - x0) * (y - y0) return x0, x1, y0, y1, wa, wb, wc, wd

3. 处理边界条件的艺术

边界处理是图像采样中最容易被忽视却至关重要的环节。我们采用零填充策略，这与PyTorch原生grid_sample的行为保持一致。

边界处理策略对比：

def apply_padding(im, grid): # 零填充一圈像素 im_padded = F.pad(im, pad=[1, 1, 1, 1], mode='constant', value=0) padded_h = im.shape[2] + 2 padded_w = im.shape[3] + 2 # 调整坐标以匹配填充后的图像 x0, x1 = x0 + 1, x1 + 1 y0, y1 = y0 + 1, y1 + 1 # 裁剪到有效范围内 x0 = torch.clamp(x0, 0, padded_w - 1) x1 = torch.clamp(x1, 0, padded_w - 1) y0 = torch.clamp(y0, 0, padded_h - 1) y1 = torch.clamp(y1, 0, padded_h - 1) return im_padded, x0, x1, y0, y1

4. 高效像素收集与组合

使用torch.gather进行高效像素收集是性能关键。这里需要特别注意张量的形状变换和维度对齐。

def gather_pixel_values(im_padded, x0, y0, x1, y1, padded_w): # 将图像展平以便于gather操作 im_flat = im_padded.view(im_padded.size(0), im_padded.size(1), -1) # 计算一维索引 x0_y0 = (x0 + y0 * padded_w).unsqueeze(1) x0_y1 = (x0 + y1 * padded_w).unsqueeze(1) x1_y0 = (x1 + y0 * padded_w).unsqueeze(1) x1_y1 = (x1 + y1 * padded_w).unsqueeze(1) # 收集四个邻近点像素值 Ia = torch.gather(im_flat, 2, x0_y0.expand(-1, im_padded.size(1), -1)) Ib = torch.gather(im_flat, 2, x0_y1.expand(-1, im_padded.size(1), -1)) Ic = torch.gather(im_flat, 2, x1_y0.expand(-1, im_padded.size(1), -1)) Id = torch.gather(im_flat, 2, x1_y1.expand(-1, im_padded.size(1), -1)) return Ia, Ib, Ic, Id

5. 完整实现与ONNX导出

将上述组件组合起来，我们得到一个完整的、可ONNX导出的grid_sample替代方案。

def custom_grid_sample(im, grid, align_corners=False): n, c, h, w = im.shape gn, gh, gw, _ = grid.shape assert n == gn, "Batch size mismatch" # 步骤1：坐标归一化 x, y = normalize_coordinates(grid, h, w, align_corners) x = x.contiguous().view(n, -1) y = y.contiguous().view(n, -1) # 步骤2：计算双线性权重 x0, x1, y0, y1, wa, wb, wc, wd = compute_bilinear_weights(x, y) # 步骤3：边界处理 im_padded, x0, x1, y0, y1 = apply_padding(im, grid) # 步骤4：像素收集 padded_w = w + 2 Ia, Ib, Ic, Id = gather_pixel_values(im_padded, x0, y0, x1, y1, padded_w) # 步骤5：加权求和 wa = wa.unsqueeze(1) wb = wb.unsqueeze(1) wc = wc.unsqueeze(1) wd = wd.unsqueeze(1) output = (Ia * wa + Ib * wb + Ic * wc + Id * wd) return output.reshape(n, c, gh, gw)

在实际项目中遇到的一个棘手问题是：当输入张量在CUDA设备上时，某些操作可能导致隐式的设备同步。特别是在处理大规模图像时，这种同步会显著影响性能。解决方案是确保所有中间张量都显式指定设备：

device = im.device x0 = x0.to(device) y0 = y0.to(device) # 其他张量也需要类似处理

这个自定义实现不仅解决了ONNX导出问题，更重要的是让我们深入理解了图像采样的核心机制。当再次面对类似"opset version not supported"的错误时，我们有了更多底气和工具去应对。