小白学习pytorch源码（一）：torch包函数如何实现？揭秘__init__.py
小白学习pytorch源码（一）
- 学习目的与计划
- 学习计划
- 学习资源
- torch包函数实现
-
- torch.randint()最详细解读
- pytorch模块结构
- 如何使用C++和cuda编程改变pytorch默认训练函数？
学习目的与计划
考研完后的暑假获得了很长一段的空闲，因为想去的实习都去不了，又不想在无用的实习上浪费时间，所以决定在暑假看懂pytorch源码，并且自己尝试实现一个最简单的深度学习框架。在网上查找资料时，发现相关资料较少，许多都为直接翻译官方pytorch文档原文，且都较为零散，缺少一个完整的学习经验，所以斗胆尝试自己制作一个源码学习文档，包含对官方文档的学习和自己的思考，以便其他学习者更好地理解。
学习计划
将大体按照官网文档的形式分模块阅读，不同的是会加入自己的思考与理解，在一些需要前置知识的要点处会补充知识说明，所有参考文档会在文章最后写出。下面开始学习。
学习资源
由于直接pip install的python包只包含pytorch的python部分和c++部分的头文件，并不包含c++部分的全部内容，所以需要到github原作者处下载源代码，具体地址如下：pytorch源代码
pytorch官方文档为英文且一直在更新中，而中文文档少有能够一直更新的，所以这里推荐看官方文档学习，地址如下：pytorch官方文档
torch包函数实现
事实上，在目前流行的机器学习框架中，pytorch算是结构最优美的框架之一。其模块虽然多，但功能各异，都具有自己的作用。
当我们执行该行指令：
import torch
python会自动执行库文件中的torch包中的_init_.py，该文件中写明了对torch包的功能定义，原文如下：
The torch package contains data structures for multi-dimensional
tensors and defines mathematical operations over these tensors.
Additionally, it provides many utilities for efficient serializing of
Tensors and arbitrary types, and other useful utilities.It has a CUDA counterpart, that enables you to run your tensor
computations on an NVIDIA GPU with compute capability >= 3.0.
即torch包主要功能为多维张量定义数据结构且为这些张量定义数学运算，同时为张量和各种数据类型提供多种功能。
继续观察该文件，我们会发现其还进行了基本的环境变量、cuda文件、DLL文件、python库文件的配置和导入（其中c文件使用ctypes），而C++函数的导入也在该文件中，即：
from torch._C import *
如上文所言，该包还包含许多基础功能和类的定义，因为过多所以不每条都解释了，具体可在torch包官方文档中找到。
相信使用过pytorch的深度学习工作者都调用过torch包中的函数，例如torch.randn()，torch.ones()等等，然而除去torch.is_tensor()等少部分函数在python中有实现外，大部分函数的声明都指向了_C包中的_VariableFunctions.pyi文件，其中包含大量函数和变量的定义。类似于：
def sigmoid(input: Tensor, *, out: Optional[Tensor]=None) -> Tensor: ...
def sigmoid_(input: Tensor) -> Tensor: ...
def sign(input: Tensor, *, out: Optional[Tensor]=None) -> Tensor: ...
def signbit(input: Tensor, *, out: Optional[Tensor]=None) -> Tensor: ...
def sin(input: Tensor, *, out: Optional[Tensor]=None) -> Tensor: ...
def sin_(input: Tensor) -> Tensor: ...
小知识：pyi文件是Python 的存根文件，用于代码检查时的类型提示。
pyi文件是PEP484提案规定的一种用于 Python 代码类型提示（Type Hints）的文件。PEP即Python
Enhancement Proposals，是经过 Python 社区核心开发者讨论并一致同意后，对外发布的一些正式规范文档。
而该文件则是由_VariableFunctions.pyi.in文件由C++编译生成，可以肯定的是，这些函数的实现都是由c++完成的。pytorch使用ctypes调用c语言接口，而对于c++则是使用pybind11来进行绑定，使其能使用c++实现的函数。
但相信读者们更希望能看到pytorch是如何实现这些函数的，这也是本文的目的所在。
笔者经过仔细寻找，在_init_.py文件中找到了如下几行代码：
_C._initExtension(manager_path())
//以及下文中的
_C._init_names(list(torch._storage_classes))
很明显，这些函数的作用就是将c++函数与python进行初始绑定，那么这些函数的位置在哪里呢？
笔者在pytorch源代码的/csrc/Module.cpp文件中找到了这些函数，两个函数cpp代码如下：
static PyObject* THPModule_initExtension(
PyObject* _unused,
PyObject* shm_manager_path) {
HANDLE_TH_ERRORS
if (!THPUtils_checkString(shm_manager_path)) {
THPUtils_setError(
"initialization error - expected bytes/string object as shm_manager_path!");
return nullptr;
}
torch::utils::initializeLayouts();
torch::utils::initializeMemoryFormats();
torch::utils::initializeQSchemes();
torch::utils::initializeDtypes();
torch::tensors::initialize_python_bindings();
std::string path = THPUtils_unpackString(shm_manager_path);
libshm_init(path.c_str());
auto module = THPObjectPtr(PyImport_ImportModule("torch"));
if (!module)
throw python_error();
THPStorage_postInit(module);
THPAutograd_initFunctions();
Py_RETURN_NONE;
END_HANDLE_TH_ERRORS
}
static PyObject* THPModule_initNames(PyObject* self, PyObject* arg) {
static std::vector<std::string> names;
THPObjectPtr types(PySequence_Fast(arg, "expected a sequence"));
if (!types)
return nullptr;
// NOLINTNEXTLINE(bugprone-branch-clone)
auto num_classes = PySequence_Fast_GET_SIZE(types.get());
names.reserve(names.size() + num_classes);
for (Py_ssize_t i = 0; i < num_classes; i++) {
PyObject* obj = PySequence_Fast_GET_ITEM(types.get(), i);
THPUtils_assert(PyType_Check(obj), "expected a PyTypeObject");
PyTypeObject* type = (PyTypeObject*)obj;
THPObjectPtr module_name(PyObject_GetAttrString(obj, "__module__"));
if (!module_name)
return nullptr;
THPUtils_assert(
THPUtils_checkString(module_name.get()),
"expected __module__ to be a string");
std::string name = THPUtils_unpackString(module_name.get());
names.emplace_back(name + "." + type->tp_name);
type->tp_name = names.back().c_str();
}
Py_RETURN_NONE;
}
同时笔者还在/csrc/autograd/python_variable.cpp中找到了如下代码：
bool THPVariable_initModule(PyObject* module) {
THPVariableMetaType.tp_base = &PyType_Type;
if (PyType_Ready(&THPVariableMetaType) < 0)
return false;
Py_INCREF(&THPVariableMetaType);
PyModule_AddObject(module, "_TensorMeta", (PyObject*)&THPVariableMetaType);
static std::vector<PyMethodDef> methods;
THPUtils_addPyMethodDefs(methods, torch::autograd::variable_methods);
THPUtils_addPyMethodDefs(methods, extra_methods);
THPVariableType.tp_methods = methods.data();
if (PyType_Ready(&THPVariableType) < 0)
return false;
Py_INCREF(&THPVariableType);
PyModule_AddObject(module, "_TensorBase", (PyObject*)&THPVariableType);
torch::autograd::initTorchFunctions(module);
torch::autograd::initTensorImplConversion(module);
return true;
}
注意该段代码中的
torch::autograd::initTorchFunctions(module);
该段代码作用即为初始化例如torch.randint()等函数，再进一步了解该段函数，笔者发现其定义在\csrc\autograd\python_torch_functions_manual.cpp文件中，功能包含绑定c++与python函数，检查错误等等。
torch.randint()最详细解读
下面以torch.randint为例，看看pytorch到底用多少步实现这个函数。首先在python_torch_functions_manual.cpp文件中，实现了python函数和c++函数的绑定：
{"randint",
castPyCFunctionWithKeywords(THPVariable_randint),
METH_VARARGS | METH_KEYWORDS | METH_STATIC,
nullptr},
而绑定的THPVariable_randint函数即为c++函数，具体代码如下：
static PyObject* THPVariable_randint(
PyObject* self_,
PyObject* args,
PyObject* kwargs) {
HANDLE_TH_ERRORS
static PythonArgParser parser(
{
"randint(int64_t high, IntArrayRef size, *, Generator generator=None, Tensor out=None, ScalarType dtype=None, Layout layout=torch.strided, Device device=None, bool requires_grad=False)",
"randint(int64_t low, int64_t high, IntArrayRef size, *, Generator generator=None, Tensor out=None, ScalarType dtype=None, Layout layout=torch.strided, Device device=None, bool requires_grad=False)",
},
/*traceable=*/false);
ParsedArgs<9> parsed_args;
auto r = parser.parse(args, kwargs, parsed_args);
if (r.has_torch_function()) {
return handle_torch_function(
r, args, kwargs, THPVariableFunctionsModule, "torch");
}
if (r.idx == 0) {
if (r.isNone(3)) {
auto high = r.toInt64(0);
auto size = r.intlist(1);
auto generator = r.generator(2);
// NOTE: r.scalartype(X) gives the default dtype if r.isNone(X)
auto dtype = r.scalartypeWithDefault(4, at::ScalarType::Long);
auto device = r.device(6);
const auto options = TensorOptions()
.dtype(dtype)
.device(device)
.layout(r.layout(5))
.requires_grad(r.toBool(7));
return wrap(dispatch_randint(high, size, generator, options));
} else {
check_out_type_matches(
r.tensor(3),
r.scalartype(4),
r.isNone(4),
r.layout(5),
r.device(6),
r.isNone(6));
return wrap(dispatch_randint(
r.toInt64(0), r.intlist(1), r.generator(2), r.tensor(3))
.set_requires_grad(r.toBool(7)));
}
} else if (r.idx == 1) {
if (r.isNone(4)) {
auto low = r.toInt64(0);
auto high = r.toInt64(1);
auto size = r.intlist(2);
auto generator = r.generator(3);
// NOTE: r.scalartype(X) gives the default dtype if r.isNone(X)
auto dtype = r.scalartypeWithDefault(5, at::ScalarType::Long);
auto device = r.device(7);
const auto options = TensorOptions()
.dtype(dtype)
.device(device)
.layout(r.layout(6))
.requires_grad(r.toBool(8));
return wrap(dispatch_randint(low, high, size, generator, options));
} else {
check_out_type_matches(
r.tensor(4),
r.scalartype(5),
r.isNone(5),
r.layout(6),
r.device(7),
r.isNone(7));
return wrap(dispatch_randint(
r.toInt64(0),
r.toInt64(1),
r.intlist(2),
r.generator(3),
r.tensor(4))
.set_requires_grad(r.toBool(8)));
}
}
Py_RETURN_NONE;
END_HANDLE_TH_ERRORS
}
吐槽一下即使是pytorch的研发大佬也避免不了用很多个if-else叠加哈哈~
至于为什么有如此多的if-else，其实是为适应多种多样的重载情况，对不同的参数输入准备了不同的dispatch_randint()函数做应对，在这里笔者贴出其中两个dispatch_randint()举例：
inline Tensor dispatch_randint(
int64_t high,
IntArrayRef size,
c10::optional<Generator> generator,
Tensor result) {
pybind11::gil_scoped_release no_gil;
return at::randint_out(result, high, size, generator);
}
inline Tensor dispatch_randint(
int64_t high,
IntArrayRef size,
c10::optional<Generator> generator,
const TensorOptions& options) {
torch::utils::maybe_initialize_cuda(options);
pybind11::gil_scoped_release no_gil;
return torch::randint(high, size, generator, options);
}
可以清楚地看到两个函数名称相同参数输入不同，在经过不同处理后，两函数进行输出。而两个函数在最后又都调用了不同函数，这两个函数的位置与之前的函数不同，并不在torch包中，而是在\pytorch-master\aten\src\ATen\native\TensorFactories.cpp这个文件中，也有多种实现，在这里贴出其中两种
Tensor& randint_out(int64_t high,
IntArrayRef size,
c10::optional<Generator> generator,
Tensor& result) {
result.resize_(size);
return result.random_(0, high, generator);
}
Tensor randint(
int64_t low,
int64_t high,
IntArrayRef size,
c10::optional<Generator> generator,
c10::optional<ScalarType> dtype,
c10::optional<Layout> layout,
c10::optional<Device> device,
c10::optional<bool> pin_memory) {
// See [Note: hacky wrapper removal for TensorOptions]
TensorOptions options = TensorOptions().dtype(dtype).layout(layout).device(device).pinned_memory(pin_memory);
auto result = at::empty(size, options);
return result.random_(low, high, generator);
}
可见到这里后，终于这个函数的调用要接近尾声，而这里的random_函数实质上还在另一文件中，其在pytorch-master\aten\src\ATen\native\DistributionTemplates.h中被实现，具体实现的说明如下：
// The purpose of
update_fromandupdate_tois to find the closest valid int64_t number that can be used as actualfrom.
// The current implementation ofrandom_uses uint64_t arithmetics and casts the result to the target dtype(scalar_t).
// This casting can result in generating numbers that happen to be greater or equal totovalue. For instance:
//
// auto actual = torch::empty({3, 3}, torch::half);
// actual.random_(0, 65504);
//
// If random’s uint64_t arithmetics produces 65503 as a random value after casting to torch::half it becomes 65504
// and violates the requirement that random value must be less thanto. To resolve this issueupdate_fromandupdate_to
// movesfromto the right andtoto the left to the next closest value that won’t go outside [from, to) after casting to
// the target dtype. Forto= 65504 it moves left for (1 << (log2(to) - 11 + 1)) = 32 and becomes 65472, which is previous
// available number for torch::half dtype.
其文中提到的update_from和update_to也在该文件中，感兴趣的朋友可以自行寻找学习。
不得不说pytorch的构造确实精致，randint听起来简单，但要实现起来却涉及到如此多的文件，令人感叹。
pytorch模块结构
在/torch/init.py中还导入了一些作者认为的常见模块，从这些导入中可以一窥pytorch的整体模块结构，具体如下：
import torch.cuda
import torch.autograd
from torch.autograd import no_grad, enable_grad, set_grad_enabled
# import torch.fft # TODO: enable once torch.fft() is removed
import torch.futures
import torch.nn
import torch.nn.intrinsic
import torch.nn.quantized
import torch.optim
import torch.optim._multi_tensor
import torch.multiprocessing
import torch.sparse
import torch.utils.backcompat
import torch.onnx
import torch.jit
import torch.linalg
import torch.hub
import torch.random
import torch.distributions
import torch.testing
import torch.backends.cuda
import torch.backends.mkl
import torch.backends.mkldnn
import torch.backends.openmp
import torch.backends.quantized
import torch.quantization
import torch.utils.data
import torch.__config__
import torch.__future__
如何使用C++和cuda编程改变pytorch默认训练函数？
如果希望能够编写c++文件改变pytorch默认的神经网络训练方式以起到加速作用，那么需要安装pybind11进行c++与python的绑定，具体的操作可以参考官方文档：pybind11官方文档，值得注意的是，必须使用visual studio 2015之后的版本环境才能正确使用pybind11。
当然，也可使用cuda编程在显卡方面加速，在笔者自己学明白cuda编程后会更多文档。
在之后的学习中，我会分模块学习pytorch的源代码以及底层实现，并发表出学习笔记，如果感兴趣的话请关注我，给我继续创作的动力，谢谢！
系列其他文章链接如下，持续更新中~
小白学习pytorch源码（一）：torch包函数如何实现？揭秘__init__.py
小白学习pytorch源码（二）：setup.py最详细解读