PyTorchコードを最初から読む--Operators編

18638 ワード

これはPyTorchのソースコードを読んで整理したノートで、後で読むのに便利です.ここでは主にPyTorchのoperatorsの定義がどのように組織されているのか、新しいoperatorを追加する場合はどうすればいいのかを知りたいです.__init__.pysetup.py良い着手点はtorchというモジュールの__init__.pyと実装用のsetup.pyです.2つのファイルを全部閲覧するには大体の概念がある.そして__init__.pyの中で__all__を探して、これらのoperatorsがどのように__all__の中に追加されたのかを観察することで、私たちが使っているすべてのoperatorsがどのように来たのかを知ることができます.
検索__all__で発見された重要な部分は、
__all__ += [name for name in dir(_C)
            if name[0] != '_' and
            not name.endswith('Base')]

この仕事はtorch._Cで定義されているいろいろなものをオンデマンドで__all__に入れることなので、どのoperatorsがどのように来たのか、それともtorch._Cを見に行く必要がありますか.名前から見れば、torch._Cというものは、PyTorchのC/C++などの言語で書かれている部分です.これはこの部分がどのように構築されたのかに関連しています.これはsetup.pyから探します.発見された関連コードセグメントは以下の通りです.
main_sources = ["torch/csrc/stub.cpp"]
extensions = []
packages = find_packages(exclude=('tools', 'tools.*'))
C = Extension("torch._C",
              libraries=main_libraries,
              sources=main_sources,
              language='c++',
              extra_compile_args=main_compile_args + extra_compile_args,
              include_dirs=[],
              library_dirs=library_dirs,
              extra_link_args=extra_link_args + main_link_args + [make_relative_rpath('lib')],
)
extensions.append(C)

基本的にtorch._Cというものはtorch/csrc/というディレクトリの中のたくさんのファイルからコンパイルされたものだと断定できます.torch/csrc/には書類が山積みになっていて、どこから着手するかが問題です.torch._Cはpythonのモジュールなので、pythonのC-bindingでこのモジュールを作成する場所があるに違いありません.これは良い着手点です.このモジュールがどこから作成されたかを見つけるには、grep大法を祭る必要があります.torch/csrc/というディレクトリでこのようなコマンドを実行します.
grep 'torch._C' -r .

Module.cppで発見された最も似た行は、関連するすべてのコード行を得ることができます.
#if PY_MAJOR_VERSION == 2
  ASSERT_TRUE(module = Py_InitModule("torch._C", methods.data()));
#else
  static struct PyModuleDef torchmodule = {
     PyModuleDef_HEAD_INIT,
     "torch._C",
     nullptr,
     -1,
     methods.data()
  };
  ASSERT_TRUE(module = PyModule_Create(&torchmodule));
#endif

この行はこのモジュールを初期化しているように見えますが、ファイル名はModule.cppと呼ばれています.では、次はここから始めましょう.__init__.pysetup.pyも私たちの歴史の舞台を脱退することができます.
また、grep大法を適用する前に、PyTorchをコンパイルしなければならないことに注意してください.PyTorchの多くのコードはコンパイル時に他のファイルに基づいて生成され、生成されたファイルを持って一緒に検索したほうがいいからです.Module.cppとautograd
このファイルを最初から最後までブラウズすると、基本的に初期化モジュールがinitModuleで完了したと断定できます.この関数の中にはたくさんのものが初期化されています.主に具体的などの行がそのoperatorsの初期化を担当しているのかを探します.initModuleには似たようなものがたくさんあることに気づきました
#ifdef USE_CUDA
  torch::cuda::initModule(module);
#endif
  ASSERT_TRUE(THPDoubleStorage_init(module));
  ASSERT_TRUE(THPFloatStorage_init(module));
  ASSERT_TRUE(THPHalfStorage_init(module));
  ASSERT_TRUE(THPLongStorage_init(module));
  ASSERT_TRUE(THPIntStorage_init(module));
  ASSERT_TRUE(THPShortStorage_init(module));
  ASSERT_TRUE(THPCharStorage_init(module));
  ASSERT_TRUE(THPByteStorage_init(module));
  ASSERT_TRUE(THPBoolStorage_init(module));

これは直接スキップして読まないことができます.あなたがこのfeatureの山を有効にしたかどうかにかかわらず、operatorsは存在します.それはあなたがこの山を有効にしなかったときです.名前を見れば関係なく、直接相手にしなくてもいいものがたくさんあります.だから最後に基本的にフィルタリングされたのは、次のように見えます.
PyObject* initModule() {
  HANDLE_TH_ERRORS
  at::init_num_threads();

  C10_LOG_API_USAGE_ONCE("torch.python.import");

#define ASSERT_TRUE(cmd) if (!(cmd)) return nullptr

  THPUtils_addPyMethodDefs(methods, TorchMethods);
  THPUtils_addPyMethodDefs(methods, DataLoaderMethods);
  THPUtils_addPyMethodDefs(methods, torch::autograd::python_functions());
  THPUtils_addPyMethodDefs(methods, torch::multiprocessing::python_functions());
#ifdef USE_CUDA
  THPUtils_addPyMethodDefs(methods, THCPModule_methods());
#endif
#ifdef USE_CUDNN
  THPUtils_addPyMethodDefs(methods, THCUDNN_methods());
#endif
#ifdef USE_DISTRIBUTED
  THPUtils_addPyMethodDefs(methods, THDPModule_methods());
#ifdef USE_C10D
  THPUtils_addPyMethodDefs(methods, torch::distributed::c10d::python_functions());
#endif
#endif

まずTorchMethodsから見て、これはModule.cppに定義されており、operatorsとは関係ないことが一目でわかります.torch::autograd::python_functions()これ、grep大法を使うと、彼の定義はtorch/csrc/autograd/init.cppにあることがわかります.何も望んでいません.THPVariable_initModuleを見続けます.grep大法を使うと、この関数はtorch/csrc/autograd/python_variable.cppの中で定義されていることがわかります.この関数の定義を見てみると、下の行に気づきます.
bool THPVariable_initModule(PyObject *module)
{
  static std::vector 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;
}

関数定義全体でも、operatorsを定義する行に最も似ている行だけがあり、深く掘り下げ続け、grep大法を使用し続けます.
grep variable_methods -r .

この変数はtorch/csrc/の下に定義されたautograd/generated/python_variable_methods.cppの中にあることがわかります.開けてみると、次のような内容が見つかります.
PyMethodDef variable_methods[] = {
  {"__add__", (PyCFunction)THPVariable_add, METH_VARARGS | METH_KEYWORDS, NULL},
  {"__radd__", (PyCFunction)THPVariable_add, METH_VARARGS | METH_KEYWORDS, NULL},
  {"__iadd__", (PyCFunction)THPVariable_add_, METH_VARARGS | METH_KEYWORDS, NULL},
  {"__rmul__", (PyCFunction)THPVariable_mul, METH_VARARGS | METH_KEYWORDS, NULL},
  {"__mul__", (PyCFunction)THPVariable_mul, METH_VARARGS | METH_KEYWORDS, NULL},
  {"__imul__", (PyCFunction)THPVariable_mul_, METH_VARARGS | METH_KEYWORDS, NULL},
  {"__sub__", (PyCFunction)THPVariable_sub, METH_VARARGS | METH_KEYWORDS, NULL},
  {"__isub__", (PyCFunction)THPVariable_sub_, METH_VARARGS | METH_KEYWORDS, NULL},
  {"__div__", (PyCFunction)THPVariable_div, METH_VARARGS | METH_KEYWORDS, NULL},
  {"__truediv__", (PyCFunction)THPVariable_div, METH_VARARGS | METH_KEYWORDS, NULL},
  {"__idiv__", (PyCFunction)THPVariable_div_, METH_VARARGS | METH_KEYWORDS, NULL},
  {"__mod__", (PyCFunction)THPVariable_remainder, METH_VARARGS | METH_KEYWORDS, NULL},
  {"__bool__", (PyCFunction)THPVariable_bool_scalar, METH_NOARGS, NULL},
  {"__float__", (PyCFunction)THPVariable_float_scalar, METH_NOARGS, NULL},
  {"__int__", (PyCFunction)THPVariable_integral_scalar, METH_NOARGS, NULL},
  {"__long__", (PyCFunction)THPVariable_integral_scalar, METH_NOARGS, NULL},
  {"__index__", (PyCFunction)THPVariable_index_scalar, METH_NOARGS, NULL},
  {"__nonzero__", (PyCFunction)THPVariable_bool_scalar, METH_NOARGS, NULL},
  {"__invert__", (PyCFunction)THPVariable_invert, METH_NOARGS, NULL},
  {"__matmul__", (PyCFunction)THPVariable_matmul, METH_VARARGS | METH_KEYWORDS, NULL},
  {"_is_view", (PyCFunction)THPVariable__is_view, METH_NOARGS, NULL},
  {"apply_", (PyCFunction)THPVariable_apply_, METH_O, NULL},
  {"byte", (PyCFunction)THPVariable_byte, METH_NOARGS, NULL},
  {"char", (PyCFunction)THPVariable_char, METH_NOARGS, NULL},
  {"contiguous", (PyCFunction)THPVariable_contiguous, METH_VARARGS | METH_KEYWORDS, NULL},
  {"copy_", (PyCFunction)THPVariable_copy_, METH_VARARGS | METH_KEYWORDS, NULL},
  {"cpu", (PyCFunction)THPVariable_cpu, METH_NOARGS, NULL},
  {"cuda", (PyCFunction)THPVariable_cuda, METH_VARARGS | METH_KEYWORDS, NULL},
  {"dim", (PyCFunction)THPVariable_dim, METH_NOARGS, NULL},
  {"double", (PyCFunction)THPVariable_double, METH_NOARGS, NULL},
  {"element_size", (PyCFunction)THPVariable_element_size, METH_NOARGS, NULL},
  {"float", (PyCFunction)THPVariable_float, METH_NOARGS, NULL},
  {"get_device", (PyCFunction)THPVariable_get_device, METH_NOARGS, NULL},
  {"bool", (PyCFunction)THPVariable_bool, METH_NOARGS, NULL},
  {"half", (PyCFunction)THPVariable_half, METH_NOARGS, NULL},
  {"int", (PyCFunction)THPVariable_int, METH_NOARGS, NULL},
  {"is_contiguous", (PyCFunction)THPVariable_is_contiguous, METH_VARARGS | METH_KEYWORDS, NULL},
  {"item", (PyCFunction)THPVariable_item, METH_NOARGS, NULL},
  {"long", (PyCFunction)THPVariable_long, METH_NOARGS, NULL},
  {"map_", (PyCFunction)THPVariable_map_, METH_VARARGS | METH_KEYWORDS, NULL},
  {"map2_", (PyCFunction)THPVariable_map2_, METH_VARARGS | METH_KEYWORDS, NULL},
  {"ndimension", (PyCFunction)THPVariable_dim, METH_NOARGS, NULL},
  {"nelement", (PyCFunction)THPVariable_numel, METH_NOARGS, NULL},
  {"new", (PyCFunction)THPVariable_new, METH_VARARGS | METH_KEYWORDS, NULL},
  {"new_empty", (PyCFunction)THPVariable_new_empty, METH_VARARGS | METH_KEYWORDS, NULL},
  {"new_full", (PyCFunction)THPVariable_new_full, METH_VARARGS | METH_KEYWORDS, NULL},
  {"new_ones", (PyCFunction)THPVariable_new_ones, METH_VARARGS | METH_KEYWORDS, NULL},
  {"new_tensor", (PyCFunction)THPVariable_new_tensor, METH_VARARGS | METH_KEYWORDS, NULL},
  {"new_zeros", (PyCFunction)THPVariable_new_zeros, METH_VARARGS | METH_KEYWORDS, NULL},
  {"numpy", (PyCFunction)THPVariable_numpy, METH_NOARGS, NULL},
  {"record_stream", (PyCFunction)THPVariable_record_stream, METH_O, NULL},
  {"requires_grad_", (PyCFunction)THPVariable_requires_grad_, METH_VARARGS | METH_KEYWORDS, NULL},
  {"short", (PyCFunction)THPVariable_short, METH_NOARGS, NULL},
  {"size", (PyCFunction)THPVariable_size, METH_VARARGS | METH_KEYWORDS, NULL},
  {"storage", (PyCFunction)THPVariable_storage, METH_NOARGS, NULL},
  {"storage_offset", (PyCFunction)THPVariable_storage_offset, METH_NOARGS, NULL},
  {"storage_type", (PyCFunction)THPVariable_storage_type, METH_NOARGS, NULL},
  {"stride", (PyCFunction)THPVariable_stride, METH_VARARGS | METH_KEYWORDS, NULL},
  {"to", (PyCFunction)THPVariable_to, METH_VARARGS | METH_KEYWORDS, NULL},
  {"tolist", (PyCFunction)THPVariable_tolist, METH_NOARGS, NULL},
  {"type", (PyCFunction)THPVariable_type, METH_VARARGS | METH_KEYWORDS, NULL},
  ${py_method_defs}
  {NULL}
};

これは長いリストで、すべてのoperatorsを列挙し、各operatorはTHPVariable_の先頭の関数に対応して同じファイルに定義され、このファイルの先頭はtools/autograd/templates/python_variable_methods.cppというテンプレートから生成されたことを説明しています.
すべてのTHPVariable_の最初の関数を参照すると、すべての異なるこれらの関数は大きく異なり、基本的にコア部分は以下の(wrap)の内容しかありません.
static PyObject * THPVariable_integral_scalar(PyObject* self, PyObject* args) {
  HANDLE_TH_ERRORS
  jit::tracer::warn("Converting a tensor to a Python integer", jit::tracer::WARN_PYTHON_DATAFLOW);
  auto& self_ = reinterpret_cast(self)->cdata;
  if (isFloatingType(self_.scalar_type())) {
    // we can't dispatch to item here because we want to avoid ATen overflow checks;
    // the python integral type (long in python2) can't overflow.
    return THPUtils_packDoubleAsInt(dispatch_to_CDouble(self_));
  } else {
    return wrap(dispatch_to_CLong(self_));
  }
  END_HANDLE_TH_ERRORS
}

そのうちdispatch_xxxxxxxxxxというoperatorのコア実装部分であるべきである.grep大法で掘り続けるには、operator検索を1つ選ぶだけでいいです.たとえば、次のようにします.
grep dispatch_to_CDouble -r .

検索するとこれらのdispatch_の先頭の関数は、同ディレクトリ以下のpython_variable_methods.cppファイルに定義されています.このファイルを開いて、これらのdispatch関数の定義を参照してください.
static double dispatch_to_CDouble(const Tensor & self) {
  AutoNoGIL no_gil;
  OptionalDeviceGuard device_guard(device_of(self));
  if (self.numel() != 1) {
    throw ValueError("only one element tensors can be converted to Python scalars");
  }
  return self.item();
}

コードから、これらのoperatorは、実際にはTensorというクラスのメンバー関数であることが分かったので、次はTensorというクラスを掘るべきだと知っています.それ以外にも、コード生成の原理を理解することが重要です.これにより、コードジェネレータがこれらのoperatorsの定義をどのように見つけ、これらの関数を生成するかがわかります.
Tensorというクラスの出典はpython_variable_methods_dispatch.hファイルのヘッダで見つけることができます.
namespace torch { namespace autograd {

using at::Tensor;
using at::Scalar;
using at::TensorList;
using at::IntArrayRef;
using at::Generator;
using at::SparseTensorRef;
using at::Storage;

${py_method_dispatch}

}} // namespace torch::autograd

このようにTensorはATenの中で定義されていることから、autogradも基本的に私たちの歴史の舞台を脱退し、ATenが登場する番になったようだ.
ATen
ATenを学ぶのは実はとても簡単で、atenディレクトリの中でめちゃくちゃにひっくり返して、フォルダごとに目を開けて、すべてのREADME.mdを一度読んで、実際にATenの演算子がどのように定義されているのかを発見して、実際には、すでにaten/src/ATen/README.mdファイルの中で、非常に詳しい説明をしました.
README.mdの情報を総合して、簡単にまとめると、PyTorchの演算子は、すべてATenの中に定義されていますが、ATenの中の演算子の実現は、一部は古いLua Torchから継承されています.この部分のコードは、aten/src/TH*というディレクトリにあります.これらは歴史が残した遺産で、継承して直接使用します.PyTorchが最終的に望んだoperatorの実現方式ではない.最終的に「良い」実現方式は、aten/src/ATen/native/ディレクトリにあります.多くの演算子も、このディレクトリの下で、再実現されています.これらの古い演算子のリストは、aten/src/ATen/Declarations.cwrapで定義されています.新しい演算子のリストの定義は、aten/src/ATen/native/native_functions.yamlにあります.本文は新しい演算子の実現方式だけを検討する.
ここまで来ると、新しい演算子が実現し続けるには、native_functions.yamlというファイルがどのように読み込まれたのかを調べる必要があります.PyTorchのルートディレクトリの下でgrep大法を使い続け、キーワードnative_functions.yamlを検索した結果、aten/src/ATen/gen.py:私たちが望んでいるように見える結果は:
native_files = filter_by_extension(options.files, 'native_functions.yaml')
gen.pyというファイルを開くと、次のようなコードが表示されます.
TEMPLATE_PATH = options.source_path + "/templates"
GENERATOR_DERIVED = CodeTemplate.from_file(
    TEMPLATE_PATH + "/GeneratorDerived.h")
TYPE_DERIVED_CPP = CodeTemplate.from_file(TEMPLATE_PATH + "/TypeDerived.cpp")
SPARSE_TYPE_DERIVED_CPP = CodeTemplate.from_file(TEMPLATE_PATH + "/SparseTypeDerived.cpp")
TYPE_DERIVED_H = CodeTemplate.from_file(TEMPLATE_PATH + "/TypeDerived.h")
TYPE_H = CodeTemplate.from_file(TEMPLATE_PATH + "/Type.h")
TYPE_EXTENDED_INTERFACE_H = CodeTemplate.from_file(TEMPLATE_PATH + "/TypeExtendedInterface.h")
TYPE_DEFAULT_H = CodeTemplate.from_file(TEMPLATE_PATH + "/TypeDefault.h")
TYPE_DEFAULT_CPP = CodeTemplate.from_file(TEMPLATE_PATH + "/TypeDefault.cpp")
TYPE_EXTENSION_H = CodeTemplate.from_file(TEMPLATE_PATH + "/TypeExtension.h")
TYPE_EXTENSION_CPP = CodeTemplate.from_file(TEMPLATE_PATH + "/TypeExtension.cpp")

LEGACY_TH_DISPATCHER_H = CodeTemplate.from_file(TEMPLATE_PATH + "/LegacyTHDispatcher.h")
LEGACY_TH_DISPATCHER_CPP = CodeTemplate.from_file(TEMPLATE_PATH + "/LegacyTHDispatcher.cpp")
LEGACY_TH_DISPATCHER_DERIVED_CPP = CodeTemplate.from_file(TEMPLATE_PATH + "/LegacyTHDispatcherDerived.cpp")
LEGACY_TH_DISPATCHER_DERIVED_H = CodeTemplate.from_file(TEMPLATE_PATH + "/LegacyTHDispatcherDerived.h")

REGISTER_CPU_H = CodeTemplate.from_file(TEMPLATE_PATH + "/RegisterCPU.h")
REGISTER_CPU_CPP = CodeTemplate.from_file(TEMPLATE_PATH + "/RegisterCPU.cpp")

REGISTER_CUDA_H = CodeTemplate.from_file(TEMPLATE_PATH + "/RegisterCUDA.h")
REGISTER_CUDA_CPP = CodeTemplate.from_file(TEMPLATE_PATH + "/RegisterCUDA.cpp")

TENSOR_H = CodeTemplate.from_file(TEMPLATE_PATH + "/Tensor.h")
TENSOR_METHODS_H = CodeTemplate.from_file(TEMPLATE_PATH + "/TensorMethods.h")

FUNCTIONS_H = CodeTemplate.from_file(TEMPLATE_PATH + "/Functions.h")

LEGACY_TH_FUNCTIONS_H = CodeTemplate.from_file(TEMPLATE_PATH + "/LegacyTHFunctions.h")
LEGACY_TH_FUNCTIONS_CPP = CodeTemplate.from_file(TEMPLATE_PATH + "/LegacyTHFunctions.cpp")

NATIVE_FUNCTIONS_H = CodeTemplate.from_file(TEMPLATE_PATH + "/NativeFunctions.h")

EXTENSION_BACKEND_REGISTRATION_H = CodeTemplate.from_file(TEMPLATE_PATH + "/ExtensionBackendRegistration.h")

および:
def generate_outputs():
    cwrap_files = filter_by_extension(options.files, '.cwrap')
    nn_files = filter_by_extension(options.files, 'nn.yaml', '.h')
    native_files = filter_by_extension(options.files, 'native_functions.yaml')

    declarations = [d
                    for file in cwrap_files
                    for d in cwrap_parser.parse(file)]

    declarations += nn_parse.run(nn_files)
    declarations += native_parse.run(native_files)
    declarations = preprocess_declarations.run(declarations)

この辺りでコンテキストを探し続けると、ATenのコード生成は、gen.pyなどのPythonスクリプトによって、前に述べたいくつかのリストファイルを解析し、aten/src/ATen/templates/ディレクトリの下のファイルに基づいて生成されることがわかります.これらのファイルは、テンプレートで、長くないので、全部閲覧すればいいです.読んでいるうちに、Tensor.h、 TensorMethods.h 。など、非常に多くの重要な情報が見つかります.
引き続きgrep大法すりを動員し、PyTorchのルートディレクトリでgrepでtensor_method_definitionsを検索すると、基本的に、上の読み終わったら、Tensorクラスがどのように定義されているかがわかります.最後のステップは、次のtensor_method_definitionsがどのように充填されているかを見ることです.function_wrapper.pyファイルに直接ジャンプprocess_native(option,output_options)関数の下に、次のセクションがあります.
if is_method:
    top_env['tensor_method_declarations'].append(
            TENSOR_METHOD_DECLARATION.substitute(env))
    top_env['tensor_method_definitions'].append(
                TENSOR_METHOD_DEFINITION.substitute(env))
    method_of.append('Tensor')

上記のコードのTENSOR_METHOD_DECLARATIONTENSOR_METHOD_DEFINITIONは、同じファイルに定義されています.
# add non-virtual declaration to Tensor.h
TENSOR_METHOD_DECLARATION = CodeTemplate("""\
${return_type} ${api_name}(${method_formals_with_defaults})${const_mark};
""")
# add non-virtual declaration to Tensor.cpp
TENSOR_METHOD_DEFINITION = CodeTemplate("""\
inline ${return_type} Tensor::${api_name}(${method_formals})${const_mark} {
    return dispatch_type().${api_name}(${method_actuals});
}
""")

これで、基本的にはATen全体のコード生成は、基本的にはすり抜けてしまい、具体的に仕事をするときは、これらの関連するファイルに戻って確認すればいいのです.
本文が終わる
備考:転入先 PyTorchコードを最初から読む--Operators編