我在哪里可以找到 torch.unique() 的源代码？

Question

我只能在pytorch源代码(https://github.com/pytorch/pytorch/blob/2367face24afb159f73ebf40dc6f23e46132b770/torch/functional.py#L783)中找到以下函数调用：

_VF.unique_dim() 和 torch._unique2()

但它们不指向目录中的其他任何地方

Answer 1

大部分 pytorch 后端代码是用 C++ and/or CUDA 实现的。要查看它，您需要在源代码中找到适当的入口点。有几种方法可以做到这一点，但我发现最简单的方法是在 github.

上搜索关键字，而无需自己下载所有代码。

例如，如果您转到 github.com 并搜索 unique_dim repo:pytorch/pytorch，然后单击左侧的“代码”选项卡，您应该可以快速找到以下内容。

来自torch/jit/_builtins.py:103

 17: _builtin_ops = [
...
103:    (torch._VF.unique_dim, "aten::unique_dim"),

由此和对代码的进一步分析我们可以得出结论，torch._VF.unique_dim 实际上是从 ATen 库中调用 aten::unique_dim 函数。

像ATen there are multiple implementations of this function. Most ATen functions are registered in aten/src/ATen/native/native_functions.yaml中的大多数函数一样，这里的函数通常会有_cpu和_cuda版本。

回到搜索结果我们可以发现CUDA实现实际上是在aten/src/ATen/native/cuda/Unique.cu:197

调用函数unique_dim_cuda

196: std::tuple<Tensor, Tensor, Tensor>
197: unique_dim_cuda(const Tensor& self, const int64_t dim, const bool sorted, const bool return_inverse, const bool return_counts) {
198:   return AT_DISPATCH_ALL_TYPES_AND2(kBool, kHalf, self.scalar_type(), "unique_dim", [&] {
199:     return unique_dim_cuda_template<scalar_t>(self, dim, false, return_inverse, return_counts);
200:   });
201: }

并且 CPU 实现在 aten/src/ATen/native/Unique.cpp:271

调用函数 unique_dim_cpu

270: std::tuple<Tensor, Tensor, Tensor>
271: unique_dim_cpu(const Tensor& self, const int64_t dim, const bool sorted, const bool return_inverse, const bool return_counts) {
272:   return AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::BFloat16, at::ScalarType::Bool, self.scalar_type(), "unique_dim", [&] {
273:     // The current implementation using `dim` always sorts due to unhashable tensors
274:     return _unique_dim_cpu_template<scalar_t>(self, dim, false, return_inverse, return_counts);
275:   });
276: }

从这一点开始，您应该能够进一步跟踪函数调用以准确了解它们在做什么。

按照类似的搜索字符串，您应该会发现 torch._unique2 分别在 aten/src/ATen/native/cuda/Unique.cu:188 and aten/src/ATen/native/Unique.cpp:264 上针对 CUDA 和 CPU 实现。

我在哪里可以找到 torch.unique() 的源代码？

where can I find the source code for torch.unique()?

pytorch