Pytorch源代码分析

1.Tensor

张量(Tensor)是现代机器学习的基础。它的核心是一个数据容器。

张量.png

Tensor在pytorch中的实现

Tensor由存储和解释组成。存储即底层数据的存储空间，和内存管理相关；解释赋予了存储空间含义，例如存储空间的数据类型是什么，数据的维度是多少等等。

例1.png

Tensor类在aten/src/ATen/templates/TensorBody.h中

/// Returns a `Tensor`'s dtype (`TypeMeta`). Defined in TensorMethods.h
caffe2::TypeMeta dtype() const noexcept;

/// Returns a `Tensor`'s device.
Device device() const;

int64_t numel() const ;
  // Length of one array element in bytes.  This is the traditional
  // Numpy naming.
size_t itemsize() const;
...
//TensorImpl指针
c10::intrusive_ptr<TensorImpl, UndefinedTensorImpl> impl_;

Tensor类提供了众多成员方法查询Tensor的相关信息。在pytorch里真正意义上的Tensor实现类是TensorImpl，Tensor类中包含impl_这个TensorImpl类变量。

TensorImpl类在c10/core/TensorImpl.h中

//存储类
Storage storage_;
//自动梯度元数据
std::unique_ptr<c10::AutogradMetaInterface> autograd_meta_ = nullptr;

TensorImpl维护一个存储类（Storage）的对象（storage_），即Storage表示了Tensor存储空间。

Storage类在c10/core/Storage.h中

c10::intrusive_ptr<StorageImpl> storage_impl_;

和Tensor、TensorImpl类比，Storage类中包含StorageImpl类变量，pytorch真正意义上的存储类是StorageImpl

Storage类在c10/core/StorageImpl.h中

private:
  DataPtr data_ptr_;//数据指针
  size_t size_bytes_;//所占的空间大小
  bool resizable_;
  // Identifies that Storage was received from another process and doesn't have
  // local to process cuda memory allocation
  bool received_cuda_;
  Allocator* allocator_;//分配器

其中，最重要的是分配器(Allocator)和数据指针(DataPtr)。

DataPtr类和Allocator类的定义在c10/core/Allocator.h中

// A DataPtr is a unique pointer (with an attached deleter and some
// context for the deleter) to some memory, which also records what
// device is for its data.
Dataptr指向一段内存，并且记录了数据所在设备和内存的析构函数

class DataPtr
private:
  c10::detail::UniqueVoidPtr ptr_;//指向内存的指针并带有析构函数
  Device device_;

class UniqueVoidPtr 
private:
  // Lifetime tied to ctx_
  void* data_;
//当data_指针生命周期到，自动执行DeleteFnPtr函数
  std::unique_ptr<void, DeleterFnPtr> ctx_;

因此数据指针DataPtr三个重要部分是void* data指针、设备类型、deleter。

struct C10_API Allocator {
  virtual ~Allocator() = default;
//分配指定大小的内存，返回DataPtr，用来Storage使用
  virtual DataPtr allocate(size_t n) const = 0;
  virtual DeleterFnPtr raw_deleter() const {
    return nullptr;
  }
//分配指定大小的内存，返回指针
  void* raw_allocate(size_t n) {
    auto dptr = allocate(n);
    AT_ASSERT(dptr.get() == dptr.get_context());
    return dptr.release_context();
  }
//释放指针指向空间
  void raw_deallocate(void* ptr) {
    auto d = raw_deleter();
    AT_ASSERT(d);
    d(ptr);
  }
};

Allocator类作为分配类的基类，有三个重要的函数:分配返回原始指针、分配返回DataPtr、释放空间。其中virtual DataPtr allocate(size_t n) const = 0;虚拟方法供不同的继承类实现。pytorch实现了不同内存的分配类。

DefaultCPUAllocator类在c10/core/CPUAllocator.cpp中
分配CPU内存

at::DataPtr allocate(size_t nbytes) const override {
    void* data = alloc_cpu(nbytes);
    profiledCPUMemoryReporter().New(data, nbytes);
    return {data, data, &ReportAndDelete, at::Device(at::DeviceType::CPU)};
  }

其中分配函数alloc_cpu调用 posix_memalign或者_aligned_malloc或者memalign

void* alloc_cpu(size_t nbytes) {
... 
void* data;
#ifdef __ANDROID__
  data = memalign(gAlignment, nbytes);
#elif defined(_MSC_VER)
  data = _aligned_malloc(nbytes, gAlignment);
#else
  int err = posix_memalign(&data, gAlignment, nbytes);
 ...
#endif
...
  return data;
}

THCCachingHostAllocator类在aten/src/THC/THCCachingHostAllocator.cpp中
分配pinned memory
CudaCachingAllocator类在c10/cuda/CUDACachingAllocator.cpp中
分配GPU内存

2.数据读入

数据读入深度学习程序运行的第一个操作。数据读入也就是把磁盘中的数据读入内存，又可以分为两个部分：第一步，在内存中分配空间，写入分配的内存中。
在pytorch，数据都是以Tensor形式存在的。因此，数据的输入以创建Tensor再填充Tensor的形式。填充Tensor也就是对Tensor的数据内存上赋值，通过数据指针访问内存。

创建Tensor

创建Tensor主要是分配Tensor所需的空间
pytorch提供empty函数创建一个Tensor
在cpu上的实现代码在aten/src/ATen/native/TensorFactories.cpp中
在cuda上的实现代码aten/src/ATen/native/cuda/TensorFactories.cu中
empty的主要流程可以分为4个部分。

1.分配器
c10::Allocator* allocator;
2.计算分配空间
int64_t size_bytes = nelements * dtype.itemsize();
3.构建storageImpl对象
auto storage_impl = c10::make_intrusive<StorageImpl>(
      c10::StorageImpl::use_byte_size_t(),
      size_bytes,
      allocator->allocate(size_bytes),//返回DataPtr对象
      allocator,
      /*resizeable=*/true);
4.构建Tensor
auto tensor = detail::make_tensor<TensorImpl>(
      std::move(storage_impl), at::DispatchKey::CPU, dtype);

CIFAR10数据集输入

CIFAR10二进制文件

一个CIFAR10二进制文件包括10000条记录，每一条记录包括label和数据，数据是一张33232的图片的像素值，二进制文件以这样的格式存储。因此文件大小是10000*3073个字节

<1×标签> <3072×像素>
....
<1×标签> <3072×像素>

CIFAR10数据输入

1.创建10000*3*32*32的图像Tensor和10000的labelTensor
images_ = torch::empty({10000, 3,32, 32}, torch::kByte);
targets_ =  torch::empty(10000, torch::kByte);
2.根据目录创建输入流
std::ifstream reader(path,std::ios::binary);
for(int i=0;i<10000;i++){
      3.输入流读取指定的字节到Tensor的数据指针指向的内存
      reader.read(reinterpret_cast<char*>(targets_.data_ptr())+i/**偏移值**/,1);
      reader.read(reinterpret_cast<char*>(images_.data_ptr())+i*3072/**偏移值**/,3072);
}

代码创建的Images_,targets_是一个二进制文件的输入Tensor，包含10000张图片。之后如果设置batch size，可以根据偏移取输入Tensor的子Tensor作为一个iteration的输入。

3.计算

深度学习的核心部分就是计算操作，深度学习的训练过程包括了一系列的数学运算。pytorch里实现成百上千个计算子，包括简单的加减数学运算和复杂的卷积操作等。

1.前向计算函数
static inline Tensor mul(const Tensor & self, const Tensor & other);
static inline Tensor & mul_out(Tensor & out, const Tensor & self, const Tensor & other);
static inline Tensor mul(const Tensor & self, Scalar other);
static inline Tensor mv(const Tensor & self, const Tensor & vec);
2后向计算函数
variable_list MulBackward0::apply(variable_list&& grads) ;
variable_list MvBackward::apply(variable_list&& grads);

在pytorch中，最简单的计算过程代码是

y = net(input);
y.backward();

这两句分别代表了前向计算和后向计算。

网络结构

网络由layer组成，layer封装了计算子和参数。

例

net是由三层全连接层组成的。

struct net : torch::nn::Module{
    torch::Tensor forward(Torch::Tensor x){
        torch::Tensor x1 = fc1->forward(x);
        torch::Tensor x2 = fc2->forward(x);
        torch::Tensor x3 = fc3->forward(x);
    }
    torch::nn::Linear fc1;
    torch::nn::Linear fc2;
    torch::nn::Linear fc2;
};

前向计算

net的前向计算依次经过三层全连接层的forward

net.png

那么具体全连接层的forward执行什么样的操作，先了解全连接层的类定义

class TORCH_API LinearImpl : public Cloneable<LinearImpl> {
 public:
  ...
  Tensor forward(const Tensor& input);
  LinearOptions options;
  Tensor weight;
  Tensor bias;
  ...
};

看代码可以知道LinearImpl类中含有参数weight和bias，由此可见，网络的参数其实是由每一层的参数组成。再看一下forward函数的具体实现

return torch::addmm(bias, input, weight.t());

追踪代码，发现foward函数最终调用的addmm计算操作，输入为linear层的输入，参数weight和bias。
那么这个网络的前向计算也可以展开为

net展开.png

类似的，复杂的网络前向计算过程也可以看成参数、输入、中间数据经过计算子的计算图

后向计算

后向计算的作用是更新参数的梯度，更新梯度需要用到中间数据，因此中间数据暂存在内存中。x1虽然在执行第三次addmm操作时没有被访问，但是需要暂存在内存等待后向传播时使用。
前向传播的数学公式为
$x1 = w1x+b1$ $x2 = w2x1+b2$ $x3 = w3x2+b3$
后向传播的数学公式为
$\frac{\partial x3 }{\partial x2}=w3$ $\frac{\partial x3 }{\partial w3}=x2$ $\frac{\partial x3 }{\partial b3}=1$
$\frac{\partial x3}{\partial x1}= \frac{\partial x3 }{\partial x2}.w2$ $\frac{\partial x3}{\partial w2}= \frac{\partial x3 }{\partial x2}.x1$ $\frac{\partial x3}{\partial b2}= \frac{\partial x3 }{\partial x2}.1$
$\frac{\partial x3}{\partial x}= \frac{\partial x3 }{\partial x1}.w1$ $\frac{\partial x3}{\partial w1}= \frac{\partial x3 }{\partial x1}.x$ $\frac{\partial x3}{\partial b1}= \frac{\partial x3 }{\partial x1}.1$
从公式中可以发现，后向传播的计算遵从链式法则，梯度从后向前传播。
根据图和公式，
addmm前向计算的输入是w1、x、b1，输出是x1
后向传播的输入是w1、x、b1和 $\frac{\partial x3}{\partial x1}$ （目标值对前向输出x1的梯度），输出是目标值对w1、x、b1的梯度
那么pytorch是怎么实现这样的计算呢
与操作前向相对应的，pytorch也实现了操作的后向计算，例如addmm的后向函数是AddmmBackward
但是，与前向不同的时，后向传播搭建显示的静态计算图。构成静态图的节点是后向操作，通过边连接后向传播的输入输出。例子的后向过程就如下图所示

后向展开.png

后向函数的输入除了前一个后向函数的输出，还有前向函数的输入，也就是说前向函数的输入也是后向函数的输入。
那么后向是怎么通过后向操作构建计算流程的呢，后向操作又是怎么访问前向操作的输入呢？
Node是一个关键的类，定义在torch/csrc/autograd/function.h中
各种后向操作都继承了这个类，AddmmBackward也是继承了Node类。
在Tensor类定义中,有grad_fn成员函数
const std::shared_ptr<torch::autograd::Node>& grad_fn() const;
grad_fn理解为梯度函数，这个函数的意思是，tensor生成操作的后向函数。例如x1由前向传播图中第一个addmm生成，x1.grad_fn()就是后向传播图上的第一个AddmmBackward。
x1的梯度函数是什么时候创建的呢，是在进行前向计算的过程中。

1. 创建后向函数
std::shared_ptr<AddmmBackward> grad_fn;
  if (compute_requires_grad( self, mat1, mat2 )) {
    grad_fn = std::shared_ptr<AddmmBackward>(new AddmmBackward(), deleteNode);
    grad_fn->set_next_edges(collect_next_edges( self, mat1, mat2 ));
    2. 保存后向函数需要的输入，这里的mat1，mat2也是前向函数的输入
    grad_fn->mat1_ = SavedVariable(mat1, false);
    grad_fn->mat2_ = SavedVariable(mat2, false);
    grad_fn->alpha = alpha;
    grad_fn->mat2_sizes = mat2.sizes().vec();
    grad_fn->beta = beta;
  }
3 进行前向计算，result即为前向函数的输出Tensor
auto tmp = ([&]() {
    at::AutoNonVariableTypeMode non_var_type_mode(true);
    return at::addmm(self_, mat1_, mat2_, beta, alpha);
  })();
  auto result = std::move(tmp);
4 设置result的梯度函数
if (grad_fn) {
      set_history(flatten_tensor_args( result ), grad_fn);
  }

经过4个步骤，在进行前向传播的同时，对每一个输出Tensor都构建后向操作，同时把输入的数据放入后向操作变量中，对这些后向操作节点构建图（计算流程）。
x3.backward()这一句简单的后向传播也就可以分解为，

通过x3.grad_fn()获得梯度函数，执行这个函数，输出三个梯度。
x3.grad_fn()的下一个节点是x2.grad_fn(),并且其中x3对x2的梯度作为x2.grad_fn()的一个输入。最后执行x1.grad_fn()。
最终，参数w1、w2、w3、b1、b2、b3都获得了梯度。

附
后向传播的代码主要在torch/csrc/autograd中