记录 jit 编译的 PyTorch Class 方法（Sphinx）

Question

我在尝试使用 Sphinx 记录自定义 PyTorch 模型时遇到问题：jit 编译的方法在文档中显示时没有文档字符串。我该如何解决？我检查了 Python Sphinx autodoc and decorated members and How to autodoc decorated methods with sphinx? 但建议的解决方案似乎不起作用。当我尝试使用 ..automethod 我得到

AttributeError: '_CachedForward' object has no attribute '__getattr__'

这是一个MWE；目前我通过编写 my_forward 并在 forward.

中调用它来规避问题

from torch import jit, Tensor


class DummyModel(jit.ScriptModule):
    """My dummy model"""

    def __init__(self, const: float):
        super(DummyModel, self).__init__()
        self.const = Tensor(const)

    @jit.script_method
    def forward(self, x: Tensor) -> Tensor:
        """This method does show as a :undoc-member: in the documentation"""
        return self.my_forward(x)

    def my_forward(self,  x: Tensor) -> Tensor:
        """This method shows up as a :member: in the documentation"""
        return x + self.const

Answer 1

当运行 Sphinx 时，将 PYTORCH_JIT 环境变量设置为 0。这将禁用脚本和跟踪注释（装饰器）。

见https://pytorch.org/docs/stable/jit.html#disable-jit-for-debugging。

记录 jit 编译的 PyTorch Class 方法（Sphinx）

Documenting jit-compiled PyTorch Class Method (Sphinx)

pytorch

python-sphinx

autodoc

torchscript