class PatchLayer(torch.autograd.Function):
@staticmethod
def forward(ctx, net, coords, radius):
""" forward patchify """
ctx.radius = radius
ctx.save_for_backward(net, coords)
patches, = cuda_corr.patchify_forward(net, coords, radius)
return patches
@staticmethod
def backward(ctx, grad):
""" backward patchify """
net, coords = ctx.saved_tensors
grad, = cuda_corr.patchify_backward(net, coords, grad, ctx.radius)
return grad, None, None
def patchify(net, coords, radius, mode='bilinear'):
""" extract patches """
patches = PatchLayer.apply(net, coords, radius)
if mode == 'bilinear':
offset = (coords - coords.floor()).to(net.device)
dx, dy = offset[:,:,None,None,None].unbind(dim=-1)
d = 2 * radius + 1 # 计算了特征块的大小
# 计算了四种加权组合（类似双线性插值的原则）。
x00 = (1-dy) * (1-dx) * patches[...,:d,:d]
x01 = (1-dy) * ( dx) * patches[...,:d,1:]
x10 = ( dy) * (1-dx) * patches[...,1:,:d]
x11 = ( dy) * ( dx) * patches[...,1:,1:]
return x00 + x01 + x10 + x11 # 返回这些组合的加和，得到插值后的特征块。
return patches # 直接返回未插值的 patches
问题1： pytorch 如何自定义操作？
torch.autograd.Function: 是 PyTorch 中创建自定义操作的基本类。需要重写 forward 和 backward 方法来实现前向和反向传播的逻辑。
上下文对象 ctx: 用于在前向传播中保存反向传播需要的数据。通过 ctx.save_for_backward 保存张量，反向传播时通过 ctx.saved_tensors 取出。
CUDA 操作 cuda_corr.patchify_forward 和 cuda_corr.patchify_backward: 这两个函数在代码中并未定义，它们是自定义的 CUDA 扩展，用于执行高效的 patch 提取和梯度计算。
在这段代码中，PatchLayer 继承自 torch.autograd.Function，并定义了自定义的前向和反向传播逻辑。这类函数可以用来创建不使用标准 torch.nn.Module 的自定义操作。
具体解析如下：
1. PatchLayer 类
PatchLayer 类主要实现了一个从网络中提取 patch（图像小块）的自定义操作。它包括 forward 和 backward 两个静态方法。
-
重载
forward静态方法:ctx是 PyTorch 提供的上下文对象，用于保存反向传播需要的信息。net表示输入的特征图。coords是提取 patch 的中心坐标。radius表示从中心点开始提取的 patch 半径，决定了 patch 的大小。ctx.radius = radius和ctx.save_for_backward(net, coords)是在前向传播中保存必要的信息，以便在反向传播时使用patches, _ = cuda_corr.patchify_forward(net, coords, radius)是调用 自定义的 CUDA 操作 来进行 patch 提取。只接收函数返回值中的第一个（假定返回的是一个元组或列表）。这是一个CUDA实现的函数，用于高效地提取特征块。
-
重载
backward静态方法:ctx.saved_tensors是在forward方法中保存的net和coords。grad是从上层损失反向传播来的梯度。cuda_corr.patchify_backward用于计算反向传播的梯度。- 返回的是 grad 与多个 None，表示只有输入的第一个参数有梯度，其他参数没有
cuda_corr.patchify_forward 是 C++ CUDA 实现的。
PatchLayer 类允许开发者在前向和反向传播中插入自定义的 CUDA 操作，从而适应特定的计算需求，同时仍然能够利用 PyTorch 的自动微分能力。
2. patchify 函数
这个函数是用来从输入的 net 张量中提取 patch（小块图像）的。
-
PatchLayer.apply(net, coords, radius): 调用自定义的PatchLayer，提取 patch。返回的patches是从net中提取的 patches。 apply 是 PyTorch 自定义函数调用的固定方式，负责调用 forward 方法。 -
双线性插值:
如果mode设置为'bilinear'，会对提取的 patch 进行双线性插值操作。这部分代码主要处理：coords - coords.floor()计算了坐标相对于整数网格点的偏移，用来决定插值的比例。dx, dy将偏移分成 x 和 y 两个方向。x00, x01, x10, x11分别表示在 4 个邻近点的插值计算结果。- 最终返回加权后的结果，完成双线性插值。
-
返回值: 如果
mode为'bilinear'，返回双线性插值后的 patch；否则直接返回从PatchLayer提取的 patches。
代码工作流程
patchify函数被调用，传入net、coords、radius和插值模式。- 如果
mode='bilinear'，则进行双线性插值，最终得到插值后的 patch。 - 否则，直接返回提取到的 patches。
forward和backward方法定义了如何进行前向和反向传播。
总结：这段代码的主要功能是在网络中提取指定半径（radius）的特征图块。这在深度学习的特征匹配、局部特征提取等任务中很有用。通过自定义的 PatchLayer，实现了一个前向（和对应的反向）CUDA操作，使得该操作可以嵌入到PyTorch的自动微分系统中。
参考资料：
- 定义torch.autograd.Function的子类，自己定义某些操作，且定义反向求导函数
- Extending PyTorch
- PyTorch: Defining new autograd functions