如何让继承torch.autograd.Function的Binarizer类支持pickle序列化？

问题根源 #

你定义的Binarizer类是局部作用域内的嵌套类（从代码缩进能看出），而pickle序列化对象时，要求类必须在全局命名空间中可被访问——局部类无法被pickle定位到，因此会抛出序列化失败的错误。

另外，虽然torch.autograd.Function的子类本身支持序列化，但仅限类处于全局作用域的情况，局部嵌套的类不满足这个条件。

解决步骤 #

1. 优先方案：将类移到全局作用域 #

把Binarizer类从当前的嵌套位置（比如某个函数或类内部）移到脚本最外层，确保它处于全局命名空间下：

import torch

# 全局作用域定义Binarizer
class Binarizer(torch.autograd.Function):
    """Binarizes {0, 1} a real-valued tensor."""

    @staticmethod
    def forward(ctx, inputs, threshold=5e-3):
        outputs = inputs.clone()
        outputs[inputs <= threshold] = 0
        outputs[inputs > threshold] = 1
        return outputs

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output, None 

之后使用时直接调用全局类的apply方法：

# 对GRU生成的掩码进行二值化
binarized_mask = Binarizer.apply(gru_mask_tensor)

2. 替代方案：使用PyTorch官方保存方法 #

PyTorch推荐用torch.save()保存模型检查点，它内部针对PyTorch对象做了专门的序列化处理，比直接用pickle更可靠：

# 保存模型状态字典
torch.save(model.state_dict(), 'model_checkpoint.pth')

# 加载时
model.load_state_dict(torch.load('model_checkpoint.pth'))

3. 特殊场景：必须保留嵌套结构的处理（不推荐） #

如果一定要把Binarizer留在局部作用域，需要给类添加__reduce__方法，手动告诉pickle如何序列化和反序列化它。这种方式复杂度高，容易出错，仅作参考：

class YourParentClass:
    def __init__(self):
        pass

    class Binarizer(torch.autograd.Function):
        @staticmethod
        def forward(ctx, inputs, threshold=5e-3):
            outputs = inputs.clone()
            outputs[inputs <= threshold] = 0
            outputs[inputs > threshold] = 1
            return outputs

        @staticmethod
        def backward(ctx, grad_output):
            return grad_output, None 

        def __reduce__(self):
            # 返回类的全局路径和初始化参数
            return (YourParentClass.Binarizer, ())

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