天天看点

pytorch nonzero_一个Tensor的生命历程(Pytorch版)上篇

“OLDPAN博客”,侃侃而谈人工智能 深度酝酿优质原创文!

此篇文章较为硬核,因内容较多,故分为上下两篇

不知道大家是否对Pytorch中的Tensor是如何生成感兴趣,深入理解这个知识可以加深你对深度学习框架的一些印象和操作熟练度。

pytorch nonzero_一个Tensor的生命历程(Pytorch版)上篇

文中涉及到大量的Pytorch的C++源码,版本为1.4.0a,适合有一定Pytorch源码基础的童鞋观看,同时也涉及到一些python中的C/C++拓展的一些基础知识,其中每一段代码的第一行表明了该代码的文件位置。需要注意有些代码是自动生成的,原始工程中并没有,需要编译。

还要注意一点,因为Pytorch仍在积极开发中,所以代码接口变化还是比较频繁,当你看到本文的时候,有可能展示的源码与master版的略有不同,但是大部分的代码逻辑变动不大,我们只需要知道核心工作原理即可。

那开始吧!

现在有一个Tensor,不,是两个,创建两个rand后的tensor然后加起来。

执行后输出:

呃,输出不重要,先将上述代码细分下:

看第一句发生了什么:

_t1 = torch.rand(3, 4)  # 
_t2 = _t1.__getitem__(0)del _t1
_t3 = torch.rand(3, 4)
res = _t2.__add__(_t3)del _t2del _t3
           

其实

torch.rand

torch_C._VariableFunctions

这个模块中,

torch.rand

不是一个python的函数,只是一个模块中方法的名称,通过

torch.rand

调用

torch

模块中的

rand

方法,而这个模块是通过python的C/C++拓展机制生成的,实际中

torch.rand

对应的代码是通过一个yaml文本自动生成的。

这个文件是一个自动生成代码函数的参数列表,Pytorch源码中有很多的代码文件是通过

gen.py

自动生成的,至于为什么要自动生成,是因为很多的函数代码比较相似,重复性较多,通过自动生成可以避免大部分重复的工作量。

通过上述的自动生成代码文件,在如下的代码的

${py_method_defs}

的位置生成

rand

以及其他函数的方法。

将上述的

native_functions.yaml

中的函数参数通过生成机制,在上述代码的

${py_method_defs}

位置,生成新的代码以及新的文件,我们可以看到我们的

"rand"

由上面代码可以看到,

"rand"

对应的绑定函数为

THPVariable_rand

。具体探究这个函数之前,我们首先需要初始化,因为这个函数要绑定在python端,将上述的一堆方法(tp_methods)与类型对象(PyTypeObject)绑定:

然后进行初始化,将上述的类型对象初始化为python中的模块:

这样我们在python端调用的时候会在生成的

torch_C._VariableFunctions

中找这个方法:

好,现在我们具体讨论一下

{"rand", (PyCFunction)(void(*)(void))THPVariable_rand, METH_VARARGS | METH_KEYWORDS | METH_STATIC, NULL}

这个方法对应的函数吧。

// torch/csrc/autograd/generated/python_torch_functions.cpp

static PyObject * THPVariable_rand(PyObject* self_, PyObject* args, PyObject* kwargs){
  HANDLE_TH_ERRORSstatic PythonArgParser parser({"rand(IntArrayRef size, *, DimnameList? names, ScalarType dtype=None, Layout layout=torch.strided, Device device=None, bool pin_memory=False, bool requires_grad=False)","rand(IntArrayRef size, *, Generator generator, DimnameList? names, ScalarType dtype=None, Layout layout=torch.strided, Device device=None, bool pin_memory=False, bool requires_grad=False)","rand(IntArrayRef size, *, Generator generator, Tensor out=None, ScalarType dtype=None, Layout layout=torch.strided, Device device=None, bool pin_memory=False, bool requires_grad=False)","rand(IntArrayRef size, *, Tensor out=None, ScalarType dtype=None, Layout layout=torch.strided, Device device=None, bool pin_memory=False, bool requires_grad=False)",
  }, /*traceable=*/true);

  ParsedArgs<9> parsed_args;
  auto r = parser.parse(args, kwargs, parsed_args);

  if (r.idx == 0) {
    auto size = r.intlist(0);
    auto __names = r.toDimnameListOptional(1);
    c10::optional names = __names ? c10::make_optional(DimnameList(__names.value())) : c10::nullopt;auto dtype = r.scalartype(2);auto device = r.device(4);const auto options = TensorOptions()
        .dtype(dtype)
        .device(device)
        .layout(r.layout(3).layout)
        .requires_grad(r.toBool(6))
        .pinned_memory(r.toBool(5));return wrap(dispatch_rand(size, names, options));
  ...
  } else if (r.idx == 3) { // 最终执行到这一个分支if (r.isNone(1)) {auto size = r.intlist(0);auto dtype = r.scalartype(2);auto device = r.device(4);const auto options = TensorOptions()
          .dtype(dtype)
          .device(device)
          .layout(r.layout(3).layout)
          .requires_grad(r.toBool(6))
          .pinned_memory(r.toBool(5));return wrap(dispatch_rand(size, options));
    } else {
      check_out_type_matches(r.tensor(1), r.scalartype(2), r.isNone(2),
                             r.layout(3), r.isNone(3),
                             r.device(4), r.isNone(4));return wrap(dispatch_rand(r.intlist(0), r.tensor(1)).set_requires_grad(r.toBool(6)));
    }
  }
  Py_RETURN_NONE;
  END_HANDLE_TH_ERRORS
}
           

可以看到上述函数最终实际执行的是

dispatch_rand

,这里需要注意,这个函数释放了GIL锁,这会使当前的执行代码和python中执行的代码互不影响:

然后我们进入

torch::rand

,这里有一点需要注意,在

torch::rand

这个函数中我们最终返回的是

autograd::make_variable

后的tensor,也就是说如果我们不需要differentiable的tensor的话,是可以直接返回

at::rand

这也就是为什么在Pytorch的C++前端中提到如果直接使用

at::rand

构造的Tensor是没有自动求导功能的:

// torch/csrc/autograd/generated/variable_factories.h

inline at::Tensor rand(at::IntArrayRef size, const at::TensorOptions & options = {}) {
  torch::jit::Node* node = nullptr;
  std::shared_ptr<:tracer::tracingstate> tracer_state;if (jit::tracer::isTracing()) { // 这个分支不会进入,因为我们并没有使用Jit
    tracer_state = jit::tracer::getTracingState();
    at::Symbol op_name;
    op_name = jit::Symbol::fromQualString("aten::rand");
    node = tracer_state->graph->create(op_name, /*num_outputs=*/0);
    jit::tracer::recordSourceLocation(node);
    jit::tracer::addInputs(node, "size", size);
    jit::tracer::addInputs(node, "options", options);
    tracer_state->graph->insertNode(node);
    jit::tracer::setTracingState(nullptr);
  }
  at::Tensor tensor = ([&]() {
    at::AutoNonVariableTypeMode non_var_type_mode(true);return at::rand(size, at::TensorOptions(options).is_variable(false));
  })();
  at::Tensor result = 
    autograd::make_variable(std::move(tensor), /*requires_grad=*/options.requires_grad());if (tracer_state) {
    jit::tracer::setTracingState(std::move(tracer_state));
    jit::tracer::addOutput(node, result);
  }return result;
}
           

然后我们继续进入

at::rand

// build/aten/src/ATen/Functions.h

static inline Tensor rand(IntArrayRef size, const TensorOptions & options) {
#ifdef USE_STATIC_DISPATCH
    return TypeDefault::rand(size, options);
#else // 从以下开始执行
    globalLegacyTypeDispatch().initForTensorTypeSet(at::detail::multi_dispatch_tensor_type_set(options));
    static auto table = globalATenDispatch().getOpTable("aten::rand(int[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor");
    return table->callUnboxedconst TensorOptions &>(size, options);#endif
}
           

我们可以看到上述的代码中

getOpTable

函数的参数是我们具体调用函数的一些string,也就是说

getOpTable

方法可以根据字符串类型的表示找到相对应的函数,我们先看看

globalATenDispatch()

是什么:

很像设计模式中的单例模式吧?那

ATenDispatch

又是什么?看到

registerOp

这个方法没,多么的熟悉啊,显然这是一个op注册机制的类。

那么与

std::string

组成map的

ATenOpTable

又是什么呢?下面的介绍已经比较清楚了,这个类储存了不同backend下的实现方法,同时也可以应用于Variables。

好了,回到上面

rand

函数中的最后一句

return table->callUnboxed(size, options);

。我们可以看到table就是

ATenOpTable

类的一个实例,而

callUnboxed

是它的一个方法,这个方法根据传递的模板参数返回了特定的函数:

进入

at::native::rand

进入

native::rand

进入

at::empty

// build/aten/src/ATen/Functions.h

static inline Tensor empty(IntArrayRef size, const TensorOptions & options, c10::optional memory_format) {
#ifdef USE_STATIC_DISPATCH
    switch(tensorTypeIdToBackend(impl::dispatchTypeId(at::detail::multi_dispatch_tensor_type_set(options)))) {
        case Backend::CPU:
            return CPUType::empty(size, options, memory_format);
            break;
        case Backend::SparseCPU:
            return SparseCPUType::empty(size, options, memory_format);
            break;
        default:
            AT_ERROR("empty not implemented for ", at::toString(at::detail::multi_dispatch_tensor_type_set(options)));
    }
#else
    globalLegacyTypeDispatch().initForTensorTypeSet(at::detail::multi_dispatch_tensor_type_set(options));
    static auto table = globalATenDispatch().getOpTable("aten::empty.memory_format(int[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor");
    return table->callUnboxedconst TensorOptions &, c10::optional>(size, options, memory_format);#endif
}
           

这次继续按照之前的方式来找到

"aten::empty.memory_format(int[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor"

这个op。需要注意这个op函数也是自动生成的,对应不同的backend。

所以最终table中Unbox后的函数为:

我们进入

at::native::empty_cpu

// aten/src/ATen/native/TensorFactories.cpp

Tensor empty_cpu(IntArrayRef size, const TensorOptions& options, c10::optional<:memoryformat> optional_memory_format) {
  AT_ASSERT(options.device().type() == DeviceType::CPU);
  AT_ASSERT(!options.is_variable());  // is_variable should have been 'unpacked'  // TODO: remove this when Variable and Tensor are merged
  check_size_nonnegative(size);

  c10::Allocator* allocator;
  if (options.pinned_memory()) {
    allocator = detail::getCUDAHooks().getPinnedMemoryAllocator();
  } else {
    allocator = at::getCPUAllocator(); // 执行这句
  }

  int64_t nelements = prod_intlist(size);
  auto dtype = options.dtype();
  auto storage_impl = c10::make_intrusive(
    dtype,
    nelements,
    allocator->allocate(nelements * dtype.itemsize()),
    allocator,/*resizeable=*/true);auto tensor = detail::make_tensor(std::move(storage_impl), at::TensorTypeId::CPUTensorId);// Default TensorImpl has size [0]if (size.size() != 1 || size[0] != 0) {
    tensor.unsafeGetTensorImpl()->set_sizes_contiguous(size);
  }auto memory_format = optional_memory_format.value_or(MemoryFormat::Contiguous);
  tensor.unsafeGetTensorImpl()->empty_tensor_restride(memory_format);return tensor;
}
           

这个时候因为我们调用的是CPU版的empty,这时需要在CPU申请空间,首先得到正确的申请空间的“方法”函数,赋予

c10::Allocator*

先是这个:

深入:

再深入:

再深入,可以发现这个allocator是

allocator_array

中的一个,在下面的函数

GetAllocator

中根据索引标号来取出:

allocator_array

这个东西是怎么来的?其实它是一个全局变量,用来储存各种allocator,同时配备了

SetAllocator

GetAllocator

来设置和获取相应的分配器:

然后使用

REGISTER_ALLOCATOR

来注册这些alloc。

// c10/core/Allocator.h

template struct AllocatorRegisterer {explicit AllocatorRegisterer(Allocator* alloc) {
    SetAllocator(t, alloc);
  }
};#define REGISTER_ALLOCATOR(t, f)                    \
  namespace {                                       \
  static AllocatorRegisterer g_allocator_d(f); \
  }
           

比如,把

DefaultCPUAllocator

注册为

DeviceType::CPU

,而

DeviceType::CPU

就是枚举成员,对应着一个数字。

DefaultCPUAllocator

就是我们在CPU中开辟空间实际要调用的alloc类,它继承了

at::Allocator

其中实际的开辟函数

alloc_cpu

free_cpu

,这两个函数在开辟空间和删除空间的时候会被调用:

未完待续,请期待下一篇~

pytorch nonzero_一个Tensor的生命历程(Pytorch版)上篇

关注Oldpan博客,同步更新博客最新消息,持续酝酿深度学习质量文。

pytorch nonzero_一个Tensor的生命历程(Pytorch版)上篇

点击下方原文链接获取更多相关文章

如果感觉有收获,不妨分享一下吧~