PyTorch源码学习系列 - 2. Tensor

本系列文章会优先发布于微信公众号和知乎专栏,欢迎大家关注

微信公众号:小飞怪兽屋

知乎:  PyTorch源码学习 - 知乎 (zhihu.com)


目录

前言

Tensor是什么?

Tensor如何存储?

PyTorch眼中的Tensor

Tensor的创建过程

Tensor的函数与方法


前言

如果有人问我刚开始接触PyTorch该从何处学起时,我的回答一定不会是举世瞩目,让人惊叹不已的神经网络(NN, Neural Network),也不会是不明觉厉,精巧设计的梯度自动求导系统(Autograd),反而恰恰是平淡无奇,却无处不在的张量(Tensor)。正如所有人学习编程所做的第一件事情就是在控制台输出“Hello World”一样,Tensor就是PyTorch的“Hello World” ,每一位初学者在接触PyTorch时所做的第一件事就是用torch.tensor函数创建一个属于自己的Tensor。

import torch
torch.tensor([1,2,3])

当我们写下上面这段代码的时候,我们就已经开始走进PyTorch的宏观世界,我们利用PyTorch提供的函数创建了一个Tensor对象。但是,不知你们是否曾好奇过Tensor是如何创建的,是否曾想过Tesnor是如何存储的,是否曾探究过Tensor是如何设计的 ?今天,就让我们深入研究上面这段代码,一起走入Tensor的微观世界。

Tensor是什么?

在深入研究之前让我们先来简单回顾下Tensor是什么?从数学角度上来说,Tensor本质上是一个多维向量。在数学里,一个数我们称之为标量,一维数据我们称之为向量,二维数据我们称之为矩阵,到三维及更高维度的时候,我们并没有为每一个维度数据都提供一个专有名称,我们将这些包含维度信息的数据统称为张量(Tensor)。

PyTorch源码学习系列 - 2. Tensor_第1张图片 Tensor

从图上可以看出,标量就是0维的Tensor,向量是1维的Tesnor,矩阵是二维的Tensor。在介绍Tensor的时候我们引入一个叫做维度的抽象概念,维度本质上是衡量事物的一种方式。比如时间就是一种维度,假如我们在网上卖东西,每天晚上都要汇总今天的销售额,我们就可以得到销售额相对于时间的关系,这里销售额就是基于时间的一维Tensor。当我们觉得销售额不仅仅和时间有关系,可能还和当天的客户访问量有关系,我们就又引入了一种维度,可以得到销售额与(时间,访问量)之间的关系,我们就建立一个二维的Tensor。Tensor就是用来表示这种多维数据的一种抽象概念,在不同的场景下具有不同的物理含义。在空间上,它是表示长宽高的三维tensor;在RGB颜色上,它是表示红绿蓝的三维Tensor;在语音通话中,它是表示时间与语音信息的二维Tensor。

Tensor如何存储?

当我们了解了Tensor的概念后,那该如何在计算机中存储包含多维信息的Tensor呢?当程序在计算机中运行时,我们需要将程序代码,程序运行需要的数据以及程序运行过程中生成的数据加载到内存中,所以我们可以得出一个结论,存储Tensor的物理媒介是内存(GPU上是显存),内存是一块可供寻址的存储单元。在我们设计Tensor的存储方案时,我们需要先弄清Tensor的特性:

  1. Tensor需要支持随机访问。这很重要,因为在Python中Tensor是支持索引(index)和切片(slice)操作的,支持随机访问可以让我们快速的找到需要的数据,提高我们数据访问效率
  2. Tensor的大小是固定的。当我们创建了一个Tensor的时候,它的大小就已经确定了,我们可以修改其上的任意元素,就是不能增加元素来修改它的大小。如果要修改其大小,我们必须根据其大小创建一个新的Tensor,将旧的Tensor上的数据复制过去

说到这里大家是不是觉得很熟悉?没错,就是在数据结构课上学到的第一个数据结构,数组。当我们创建一个数组的时候,我们会向内存申请一块指定大小的连续存储空间,而这恰恰就是PyTorch中Strided Tensor的存储方式。

PyTorch源码学习系列 - 2. Tensor_第2张图片 Tensor逻辑表示与物理存储的映射关系

PyTorch引入了一个叫做步伐(Stride)的概念,其本质上是逻辑索引的一个相对距离,表明当你从某个元素沿着某一维度移动一个元素时的距离。图中是一个二维矩阵,所以其Stride是个size为2的一维向量。当我们从左往右移动的时候(1.0 -> 1.1),由于这两个数在内存中紧挨着,所以我们只移动了一次,Stride在这一维度的值是1;当我们从上往下移动的时候(1.0 -> 2.0),在内存中这两个数之间隔着2个数,所以我们移动了3次,Stride在这一维度的值是3。最终对于这个Tensor来说,它的Stride是(3,1)。

PyTorch源码学习系列 - 2. Tensor_第3张图片 Tensor元素索引

当我们有了Stride后就可以快速的计算出一个元素的物理地址。熟悉C/C++的朋友会发现这其实和C/C++里的多级指针寻址方式非常类似。当我们在Tensor中查找索引是[1,1]的数时,我们通过计算索引和Stride的点积就可以得到逻辑表示的相对位置。而此元数的物理地址就是Tensor的第一个元素地址加上逻辑相对位置与字节数的乘积。最终我们计算出Tensor [1,1] 的物理地址是0x3f, 值是2.1。当然Stride在PyTorch里不仅仅是用来为某一个特定元素寻址。

PyTorch源码学习系列 - 2. Tensor_第4张图片 Tensor视图

前面也说了Tensor是支持Python的切片操作,所以我们经常会有个需求是查看当前Tensor里的部分元素,当我们使用切片操作时就会获得一个新的Tensor。如果我们为每个切片操作都开辟一个新的内存空间,那势必会降低程序的运行效率以及造成内存空间浪费,所以PyTorch引入了一个叫做视图的概念,所有的视图都会共享相同的内存空间,而Stride就是创建新Tesnor视图的关键所在。当我们想要访问Tensor[0,1:] 时,我们创建了一个包含2个元素的一维Tensor视图,这两个元数在内存中是相邻的,所以其Stride是1。但当我们想要访问Tensor[:, 0] 时,我们虽然也创建了一个包含2个元素的一维Tensor,但是这两个数在内存中却相隔2个元素,所以其Stride是3. 正是由于Stride的存在,PyTorch才可以便捷地使所有Tensor视图共享一块内存空间。那PyTorch是怎么做到这些的呢?

PyTorch眼中的Tensor

在我们用户眼中,Tensor就是一个Python类,但是在PyTorch的眼中,事实却绝不仅仅如此简单。

PyTorch源码学习系列 - 2. Tensor_第5张图片 Tensor 类关系

PyTorch将Tensor的物理存储抽象成一个Storage类,与逻辑表示类Tensor解耦,这样我们就可以建立Tensor视图和物理存储Storage之间多对一的联系。Storage是一个声明类,其具体实现在其实现类StorageImpl中。StorageImp中有两个核心的成员:

  1. data_ptr。其指向数据实际存储的内存空间。在类DataPtr中包含了Device相关的成员变量
  2. allocator_。其是一个内存分配器。Allocator是一个抽象类,所有派生类必须实现allocate和raw_deleter两个抽象函数

PyTorch的Tensor除了用Storage类来管理物理存储外,还在Tensor中定义了很多相关的元信息。比如我们前面说到的size,stride以及dtye,这些信息都存在TensorImpl类中的sizes_and_strides_以及data_type_中。key_set_中保存的是PyTorch对Tensor的layout,device以及dtype相关的调度信息。我们在前一篇文章中的架构图里还介绍过PyTorch在C++实现层中实现了算子(Operator),那Operator是如何和Tensor绑定的了?其实PyTorch创建了一个TensorBody.h的模板文件,在该文件中创建了一个继承基类TesnorBase的类Tensor。TensorBase基类中封装了所有和Tensor存储相关的细节,在类Tensor中,PyTorch使用代码自动生成工具将aten/src/ATen/native/native_functions.yaml 中声明的函数替换此处的宏${tensor_method_declarations}

class TORCH_API Tensor: public TensorBase {
...
 public:
  ${tensor_method_declarations}
...
}

当我们了解了C++实现层中的Tensor后,让我们来看看Python中的Tensor是如何与C++中的Tensor绑定的。Python中的Tensor继承于基类_TensorBase,该类并不是用Python语言定义的,其是用Python C API绑定的一个C++类

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;
}

从代码中可以看到Python中的_TensorBase类就是C++中的THPVariableType对象。THPVariable_initModule函数除了声明一个_TensorBase python类之外,还通过 torch::autograd::initTorchFunctions(module) 函数声明了Python Tensor相关的函数,我们常用的 torch.tensor 函数就是在此函数中声明。THPVariableType对象的类型是类PyTypeObject,该类由Python C API 提供,不熟悉的朋友可以自己网上找下相关资料了解下,这里就不详细展开。

PyTypeObject THPVariableType = {
    PyVarObject_HEAD_INIT(
        &THPVariableMetaType,
        0) "torch._C._TensorBase", /* tp_name */
    sizeof(THPVariable), /* tp_basicsize */
    ...
    THPVariable_properties, /* tp_getset */
    ...
    THPVariable_pynew, /* tp_new */
};

这里主要需要了解的就是THPVariable_properties THPVariable_pynew 两个函数。THPVariable_properties主要声明Python Tensor中的属性,如detype,device,layout,shape。THPVariable_pynew 就是Python Tensor通过new语法创建对象时候调用的函数,可以理解为类的构造函数。下面的python代码就会调用该函数

torch.Tensor([1,2,3])

直接通过torch.Tensor创建Python Tensor对象与直接用torch.tensor创建Python Tensor对象是等效的。唯一的区别是前者调用的是类自身的new方法,而后者是用函数(在python中类里的函数称之为方法,类外的函数称之为函数,与C++习惯不同)去创建一个Python Tensor。

Tensor的创建过程

现在我们终于可以来看看当我们用torch.tensor去创建Python Tensor的时候,PyTorch到底做了什么。

import torch
torch.tensor([1,2,3])

torch.Tensor 会调用C++的THPVariable_tensor 函数,该函数在文件 torch/csrc/autograd/python_torch_functions_manual.cpp 中。在经过一系列的参数检测之后,在函数结束之前调用了torch::utils::tensor_ctor 函数。

// implemented on python object to allow torch.tensor to be constructed with
// arbitrarily nested python objects - list, tuple, np array, scalar, etc.
static PyObject* THPVariable_tensor(
    PyObject* self,
    PyObject* args,
    PyObject* kwargs) 
  HANDLE_TH_ERRORS
  ...
  return THPVariable_Wrap(torch::utils::tensor_ctor(
      torch::tensors::get_default_dispatch_key(),
      torch::tensors::get_default_scalar_type(),
      r));
  END_HANDLE_TH_ERRORS
}

torch::utils::tensor_ctor 在文件torch/csrc/utils/tensor_new.cpp 中,该文件包含了创建Tensor的一些工具函数。在该函数中调用了internal_new_from_data 函数创建Tensor。

Tensor tensor_ctor(
    c10::DispatchKey dispatch_key,
    at::ScalarType scalar_type,
    PythonArgs& r) {
    ...
    auto new_tensor = internal_new_from_data(
        typeIdWithDefault(r, 2, dispatch_key),
        r.scalartypeWithDefault(1, scalar_type),
        r.deviceOptional(2),
        data,
        /*copy_variables=*/true,
        /*copy_numpy=*/true,
        /*type_inference=*/type_inference,
        pin_memory);
    ...
    return new_tensor;
}

internal_new_from_data 函数中有3个比较重要的函数

  1. empty。该函数主要根据参数的大小来为Tensor在CPU内存上申请一块空间用来后面存储元素。
  2. recursive_store。一个递归函数用来遍历传进来的每个元素并将其保存到对应位置
  3. to。根据设备参数来决定Tensor放到哪个设备中
Tensor internal_new_from_data(
    c10::TensorOptions options,
    at::ScalarType scalar_type,
    c10::optional device_opt,
    PyObject* data,
    bool copy_variables,
    bool copy_numpy,
    bool type_inference,
    bool pin_memory = false) {
  ...

  Tensor tensor;
  {
        ...
        tensor = at::empty(sizes, opts.pinned_memory(pin_memory));
        if (c10::multiply_integers(tensor.sizes()) != 0) {
          recursive_store(
              (char*)tensor.data_ptr(),
              tensor.sizes(),
              tensor.strides(),
              0,
              inferred_scalar_type,
              tensor.dtype().itemsize(),
              data);
    }

    tensor = tensor.to(
        device, inferred_scalar_type, /*non_blocking=*/false, /*copy=*/false);
  ...
  return at::lift_fresh(tensor);
}

最后我们来看看recursive_store 这个函数。该函数的核心就在于

  1. 递归函数的退出条件。当我们递归到最后一维元素的时候,用store_scalar函数来将元素放到对应的空间上
  2. 递归函数循环体。这个函数本质上是一个深度优先遍历算法。
void recursive_store(
    char* data,
    IntArrayRef sizes,
    IntArrayRef strides,
    int64_t dim,
    ScalarType scalarType,
    int elementSize,
    PyObject* obj) {
  int64_t ndim = sizes.size();
  if (dim == ndim) {
    torch::utils::store_scalar(data, scalarType, obj);
    return;
  }

  auto n = sizes[dim];
  auto seq = THPObjectPtr(PySequence_Fast(obj, "not a sequence"));
  ...
  PyObject** items = PySequence_Fast_ITEMS(seq.get());
  for (const auto i : c10::irange(n)) {
    recursive_store(
        data, sizes, strides, dim + 1, scalarType, elementSize, items[i]);
    data += strides[dim] * elementSize;
  }
}

Tensor的函数与方法

Tensor创建之后我们就需要通过函数或方法来对Tensor进行操作。前面也说到在Python中,类外的函数称之为函数,类内的函数称之为方法。我们通过THPVariable_initModule 函数来将Tensor相关的类和函数与Python对象联系起来。从下面code可以看出,Tensor的方法主要是通过torch::autograd::variable_methods extra_methods 两个对象来初始化的。

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;
}

而Tensor的函数则是通过函数initTorchFunctions 来初始化。在initTorchFunctions 中我们调用gatherTorchFunctions 来初始化函数,其里面主要分两种函数

  1. torch_functions_manual 数组。这个数组在代码中手动加入了一些函数的声明,我们前面一直讨论的torch.tensor 函数就是在这里注册的。
  2. gatherTorchFunctions_n 函数。这是利用PyTorch的代码自动生成脚本生成的函数,脚本会根据模板文件将aten/src/ATen/native/native_functions.yaml 中声明的函数生成到对应的函数中。
// File: torch/csrc/autograd/python_torch_functions_manual.cpp
void initTorchFunctions(PyObject* module) {
  static std::vector torch_functions;
  gatherTorchFunctions(torch_functions);
  THPVariableFunctions.tp_methods = torch_functions.data();
  ...
}

void gatherTorchFunctions(std::vector& torch_functions) {
  constexpr size_t num_functions =
      sizeof(torch_functions_manual) / sizeof(torch_functions_manual[0]);
  torch_functions.assign(
      torch_functions_manual, torch_functions_manual + num_functions);
  // NOTE: Must be synced with num_shards in
  // tools/autograd/gen_python_functions.py
  gatherTorchFunctions_0(torch_functions);
  gatherTorchFunctions_1(torch_functions);
  gatherTorchFunctions_2(torch_functions);
  ...
}

你可能感兴趣的:(PyTorch源码学习,机器学习,pytorch,深度学习)