pytorch - 如何在源代码中查找从torch._C导入的函数

Question

我正在尝试追踪源代码中 torch.nn.NLLLoss 的实现。我在文件 torch.nn.function 的函数 nll_loss 中调用了 torch._C.nll_loss。但我找不到创建_C的地方。

有人有这方面的信息吗？

Answer 1

看看A Tour of PyTorch Internals在 PyTorch 博客上。相关摘录:

PyTorch defines a new package torch. In this post we will consider the ._C module. This module is known as an “extension module” - a Python module written in C. Such modules allow us to define new built-in object types (e.g. the Tensor) and to call C/C++ functions.

The ._C module is defined in torch/csrc/Module.cpp. The init_C() / PyInit__C() function creates the module and adds the method definitions as appropriate. This module is passed around to a number of different __init() functions that add further objects to the module, register new types, etc.

Part II那篇文章详细介绍了构建系统。在有关神经网络模块的部分中，它说

Briefly, let’s touch on the last part of the build_deps command: generate_nn_wrappers(). We bind into the backend libraries using PyTorch’s custom cwrap tooling, which we touched upon in a previous post. For binding TH and THC we manually write the YAML declarations for each function. However, due to the relative simplicity of the THNN and THCUNN libraries, we auto-generate both the cwrap declarations and the resulting C++ code.

The reason we copy the THNN.h and THCUNN.h header files into torch/lib is that this is where the generate_nn_wrappers() code expects these files to be located. generate_nn_wrappers() does a few things:

• Parses the header files, generating cwrap YAML declarations and writing them to output .cwrap files
• Calls cwrap with the appropriate plugins on these .cwrap files to generate source code for each
• Parses the headers a second time to generate THNN_generic.h - a library that takes THPP Tensors, PyTorch’s “generic” C++ Tensor Library, and calls into the appropriate THNN/THCUNN library function based on the dynamic type of the Tensor

如果没有上下文，也许没有多大帮助，但我认为我不应该在此处复制整个帖子。

当我在没有阅读这些帖子的情况下试图追踪 NLLLoss 的定义时，我最终得到了 aten/src/THNN/generic/ClassNLLCriterion.c ，通过aten/src/ATen/nn.yaml 。后者可能是第二篇文章谈到的YAML，但我没有检查过。