Python中‘ctx’与‘self’的区别？PyTorch自定义Function场景疑问解析

问题描述 #

在使用深度学习库PyTorch时，我遇到了如下自定义Function的代码定义。请问ctx的行为是否与self相同？

class LinearFunction(Function):
    @staticmethod
    def forward(ctx, input, weight, bias=None):
        ctx.save_for_backward(input, weight, bias)
        output = input.mm(weight.t())
        if bias is not None:
            output += bias.unsqueeze(0).expand_as(output)
        return output
    @staticmethod
    def backward(ctx, grad_output):
        input, weight, bias = ctx.saved_variables
        grad_input = grad_weight = grad_bias = None
        # 后续反向传播逻辑...

解答 #

这是个非常关键的问题，搞清楚ctx和self的区别对理解PyTorch自定义算子的反向传播逻辑很重要！答案是：两者的行为完全不同，具体差异如下：

• 根本性质不同：你注意到代码里forward和backward都被@staticmethod装饰了，这意味着这两个方法是类的静态方法，不绑定任何类实例——所以这里根本不存在self参数（静态方法不会自动接收实例引用），而ctx是PyTorch框架自动传入的上下文对象，专门用于前向与反向传播之间的状态传递。
• 核心作用不同：ctx的唯一职责是作为“数据中转站”：在forward中通过save_for_backward()保存反向传播需要的张量，在backward中再通过saved_variables（或者新版本的saved_tensors）取出这些数据。除此之外，你还可以给ctx添加自定义属性来存储非张量类型的辅助信息（比如ctx.batch_size = input.size(0)），反向时直接读取即可。
• 使用场景不同：普通类中的self用于访问实例的属性和方法，是面向对象编程中实例与方法绑定的载体；但在PyTorch的Function体系里，因为静态方法的设计，所有跨阶段的状态传递都必须依赖ctx来完成，它是PyTorch自动管理的特殊对象，和普通类的实例引用self没有任何关系。

简单来说，ctx是PyTorch为自定义算子量身打造的上下文传递工具，和普通类里的self完全不是一回事~

内容的提问来源于stack exchange，提问作者Peri Javia