PyTorch自定义Tensor类引发JIT追踪警告的解决方法咨询

方案1：脚本化自定义张量包装类 #

PyTorch JIT对带类型标注的脚本类支持良好，把你的MyTensor改成脚本化类，就能让追踪器识别属性访问逻辑：

import torch

@torch.jit.script
class MyTensor:
    tensor: torch.Tensor

    def __init__(self, tensor: torch.Tensor):
        self.tensor = tensor

    # 按你的7个可解释值定义对应属性，示例如下
    @property
    def loss_term1(self) -> torch.Tensor:
        return self.tensor[..., 0]
    
    @property
    def loss_term2(self) -> torch.Tensor:
        return self.tensor[..., 1]
    
    # 继续定义到loss_term7

在MyModule的forward中直接返回这个类的实例，torch.jit.trace就能正常追踪，同时保留属性访问的可读性。

方案2：用namedtuple替代自定义类 #

如果不想写脚本类，Python原生的namedtuple是JIT完全兼容的结构，字段访问和属性一样直观：

from collections import namedtuple
import torch

# 按你的7个值命名字段，比如和原来的属性名一致
MyTensor = namedtuple('MyTensor', ['loss_term1', 'loss_term2', 'loss_term3', 'loss_term4', 'loss_term5', 'loss_term6', 'loss_term7'])

class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = ...  # 你的骨干网络

    def forward(self, x):
        output = self.backbone(x)  # 输出形状(batch_size, 3, 512, 7)
        return MyTensor(
            output[..., 0],
            output[..., 1],
            output[..., 2],
            output[..., 3],
            output[..., 4],
            output[..., 5],
            output[..., 6]
        )

这种方式无需额外装饰，代码更简洁，同时完全避免JIT追踪警告。

方案3：模块内封装访问方法（备选） #

如果不想改变返回值类型，也可以在MyModule中直接提供获取各损失项的方法，不过需要在forward后保留输出引用：

class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = ...
        self._last_output = None

    def forward(self, x):
        self._last_output = self.backbone(x)
        return self._last_output

    def get_loss_term1(self):
        return self._last_output[..., 0]
    
    # 继续定义get_loss_term2到get_loss_term7

这种方式适合不想改变返回值类型的场景，但使用时需要先调用forward，再调用对应的get方法。

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