PyTorch中自定义数据类型替代浮点型实现模型微调/测试的方法问询

可以实现，而且无需逐一替换模型内的函数调用，PyTorch支持通过自定义张量类型（Tensor subclass） 和运算符重载来让内置函数自动适配你的自定义数据类型。以下是具体实现步骤：

1. 自定义张量子类 #

继承torch.Tensor实现自定义数据类型，核心是重载__new__方法初始化张量存储，以及重载基础算术运算方法，确保运算后仍返回自定义类型，避免类型退化。

示例框架：

import torch

class CustomTensor(torch.Tensor):
    # 初始化自定义张量，可在这里处理自定义数据编码逻辑
    @classmethod
    def __new__(cls, data, *args, **kwargs):
        # 底层可复用PyTorch张量存储，也可自定义Storage实现特殊数据表示
        return super().__new__(cls, data, *args, **kwargs)
    
    # 重载加法运算
    def __add__(self, other):
        # 替换为你的自定义加法逻辑
        result = self.data + other.data  # 示例逻辑，需修改为你的规则
        return CustomTensor(result)
    
    # 重载减法运算
    def __sub__(self, other):
        result = self.data - other.data
        return CustomTensor(result)
    
    # 重载乘法运算
    def __mul__(self, other):
        result = self.data * other.data
        return CustomTensor(result)
    
    # 重载除法运算
    def __truediv__(self, other):
        result = self.data / other.data
        return CustomTensor(result)
    
    # 重载原地运算（如+=），保证原地操作后类型不变
    def __iadd__(self, other):
        self.data += other.data
        return self

2. 适配PyTorch内置函数与模块 #

PyTorch的大部分内置函数（如torch.add、torch.sum）和层（如nn.Linear、nn.Conv2d）会自动调用张量的重载方法，部分复杂操作可通过以下两种方式适配：

方法一：实现__torch_function__拦截函数调用 #

通过重载__torch_function__拦截PyTorch内置函数，将其路由到自定义逻辑：

@classmethod
def __torch_function__(cls, func, types, args=(), kwargs=None):
    if kwargs is None:
        kwargs = {}
    
    # 针对torch.sum实现自定义逻辑
    if func is torch.sum:
        dim = kwargs.get("dim", None)
        result = self.data.sum(dim=dim)
        return CustomTensor(result)
    # 复用已重载的运算处理torch.add
    elif func is torch.add:
        return args[0] + args[1]
    
    # 未自定义的函数，尝试用默认逻辑处理后转换为CustomTensor
    try:
        result = super().__torch_function__(func, types, args, kwargs)
        return CustomTensor(result)
    except TypeError:
        raise NotImplementedError(f"CustomTensor暂不支持函数{func}")

方法二：注册底层算子实现（进阶） #

对于卷积、矩阵乘法等底层操作，可通过torch.library注册自定义实现，让内置层自动调用：

import torch.library

# 为CustomTensor注册自定义线性层运算
@torch.library.impl("aten::linear", "CustomTensor")
def custom_linear(input, weight, bias=None):
    # 替换为你的自定义矩阵乘法+偏置逻辑
    output = input.data @ weight.T
    if bias is not None:
        output += bias
    return CustomTensor(output)

# 注册后，nn.Linear会自动使用该实现处理CustomTensor输入

3. 模型适配与测试 #

• 只需将输入数据转换为CustomTensor类型，模型的所有运算会自动触发自定义逻辑：

  model = torch.nn.Linear(10, 2)
  input_data = CustomTensor(torch.randn(32, 10))
  output = model(input_data)  # 自动使用自定义linear实现
  

• 若优化器出现类型兼容问题，可在优化步骤前将参数临时转为浮点型，更新后再转回CustomTensor。

关键注意事项 #

• 所有重载运算必须返回CustomTensor类型，避免运算后退化回普通张量
• 复杂操作（如卷积、池化）需仔细实现自定义逻辑，或基于已有重载运算组合实现
• 测试时优先验证基础运算，再逐步覆盖模型中的复杂层

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