PyTorch源码浅析

最新推荐文章于 2024-06-02 15:51:08 发布

血_影

最新推荐文章于 2024-06-02 15:51:08 发布

阅读量567

点赞数

CPU上的张量（多维数组）库

TH库的实现使用了用C语言的宏产生的泛型，并且通过命名规则来产生类似面向对象的效果。这部分我们在这一章后面介绍。

TH负责实现CPU上的张量（Tensor）运算，存储，内存分配。张量的内存管理和运算类型通过THTensor和THStorage两个C泛型来进行建模。张量这个数学对象被TH分解为THTensor和THStorage，THTensor提供一种查看THStorage的方法，THStorage负责管理张量的存储方式。

数据存储

存储的数据结构声明如下

typedef struct THStorage
{
 real *data;
 ptrdiff_t size;
 int refcount;
 char flag;
 THAllocator *allocator;
 void *allocatorContext;
 struct THStorage *view;
} THStorage;

所有在CPU上的张量实际上都是内存中的一个一维C数组（C指针）data来存储，并且使用引用计数（reference count）来管理内存。

构造函数

所有构造新THStorage的函数都以new开头，后面跟具有相关含义的后缀名。

TH_API THStorage* THStorage_(new)(void);

TH_API THStorage* THStorage_(newWithSize)(ptrdiff_t size);

TH_API THStorage* THStorage_(newWithSize1)(num);

TH_API THStorage* THStorage_(newWithSize2)(num, num);

TH_API THStorage* THStorage_(newWithSize3)(num, num, num);

TH_API THStorage* THStorage_(newWithSize4)(num, num, num, num);

TH_API THStorage* THStorage_(newWithMapping)(const char *filename, ptrdiff_t size, int flags);


/* takes ownership of data */

TH_API THStorage* THStorage_(newWithData)(num *data, ptrdiff_t size);


TH_API THStorage* THStorage_(newWithAllocator)(ptrdiff_t size,

                                           THAllocator* allocator,

                                           void *allocatorContext);

TH_API THStorage* THStorage_(newWithDataAndAllocator)(

num* data, ptrdiff_t size, THAllocator* allocator, void *allocatorContext);

析构函数都以`free`开头（实际上只有一个名为`free`的函数）

张量

张量在TH中是一种查看存储（Storage）的方法。它包括以下内容：

- long *size：一个用来存储张量各个维度大小的一维数组

- long *stride：一个用来存储张量各个角标偏移量的数组

- int nDimension：维度

- THStorage *storage：存储的指针（作者在这里也注明了，张量大小实际上是小于等于存储大小的）

- ptrdiff_t: 存储偏移

- refcount：引用计数

- char flag：（暂时还没完全看懂`flag`有啥用）

它的具体声明如下

typedef struct THTensor
{
 long *size;
 long *stride;
 int nDimension;

 // 注意: storage->size 可能比张量的大小要大
 THStorage *storage;
 ptrdiff_t storageOffset;
 int refcount;

 char flag;

} THTensor;

我们接下来具体解释几个可能不太容易理解的地方。首先是`stride`，说道`stride`我们要先简单介绍诸如**NumPy**，**Eigen**等提供了BLAS（基本线性代数运算）是如何存储一个矩阵的。首先矩阵在内存中实际上都作为一个内存块进行存储，在C语言看来它是一个一维数组或者说是由`malloc`或者`calloc`分配的某个给定大小的内存块，例如下表是一个有20个浮点类型（双精度）的内存块，它可能存储了一个4x5矩阵的值，也有可能存储了一个2x5x2的三阶张量的值。

内存中的具体存储信息

向系统申请这个内存块，在不再使用之后删除所分配的内存，将内存块固定到硬盘等存储中，以及访问指定地址等任务实际上就可以单独交给THStorage来完成，因为我们并不需要知道其对应张量的大小。甚至有可能几个元素数目不同但是总数相同的张量（比如4x4,2x2x2x2,1x16的不同大小张量）可以通过用不同的THTensor共享一块内存（共用一个THStorage，但此THStorage的引用计数将会大于等于3）。但当我们需要完成张量的一些运算，例如对于矩阵，他们的乘积（matrix product），点积（dot product）等运算会需要使用维度的信息（各个维度的大小）并且这个时候我们将按照维度来访问不同位置的元素，这使得我们首先需要存储各个维度的大小long *size，但是这还不够，我们实际上在访问一块连续内存的时候实际上使用的是不同维度上的间隔，例如第一个维度上的间隔一般是0，第二个维度上的间隔是第一个维度的大小size[0]，依次类推，但也有可能由于是由某个较大的张量分割来的，并不满足上述间隔分配方式，所以我们有必要再用一个数组存储各个维度的间隔大小long *stride，同时再加上内存的偏移量storageOffset。这样在访问某个角标ijk对应的内存地址时就可以用

storageOffset + i * stride[0] + j * stride[1] + k * stride[2]

来获得其真实内存地址了。而类似于存储，一个张量也有可能被不用的变量所使用，这也需要一个引用计数refcount来管理内存。

张量构造

张量的构造相比存储对象的构造就麻烦多了，但很多时候这些操作的共性就是对每一个或者部分张量元素使用某一个函数，在一些语言或者框架中，这被称为map函数。在TH中，使用了宏函数来做到一个高性能的map函数，我们首先介绍一下TH是如何使用宏函数做到高性能的map的。

TensorApply宏

`TensorApply`系列的宏函数是TH实现各种张量元素操作最重要的操作，它们负责把一个针对某些标量的操作应用到多个张量元素上去。在GPU部分是相当于一个map的操作。大致方法是优先去操作内存连续部分，然后再操作不连续的部分，以增加CPU cache命中率。详细内容留到下一篇文章讲。

使用C语言实现面向对象以及泛型

在PyTorch/Torch中，后端的库都使用了宏来进行泛型等功能的实现。下面我们用一个例子介绍这一部分。面向对象这一点可以通过命名规范来完成，例如我们的向量结构体如果是`Vector`，那么属于这个结构体的方法就是`Vector_xxx`。下面主要介绍泛型。

需求

现在我们需要在C语言中实现对两个向量的加法`add`。并且向量的类型可以是：`float`, `double`。

实现一

很容易想到的一个方法就是针对不同类型编写按照规则命名的`add`函数以及向量结构体`Vector`，例如我们可以分别实现如下的`Vector`类型：`Float_Vector`, `Double_Vector`。同时实现其各自对应的加法函数（假设函数输入类型必须一致）：`Float_Vector_add`, `Double_Vector_add`。

实现二

上述的实现方法实际上重复写了很多代码，我们知道两个向量的加法就是各个元素对应相加。以上所有类型所需的算法是完全相同的。假如在实现了泛型语言中去做这件事情是非常简单的，比如在C++中我们可以使用模板函数

// 这里的Vector是某个自己实现的类型

template<typename T>
void add(Vector<T> &c, Vector<T> &a, Vector<T> &b)
{
 for(int i=0; i<a.size(); i++)
 {
c.data[i] = a.data[i] + b.data[i]
 }
}

或者对于一些有自动类型匹配的语言，比如Julia，直接将变量指定为这些类型的抽象类型即可

function add!{T<:Number}(c::Vector{T}, a::Vector{T}, b::Vector{T})
 for i=1:size(c)
     c[i] = a[i] + b[i]
 end
end

而C并不具备这样的功能。但不难从实现一中发现，不同类型的命名方式是固定的，这使得我们可以通过借助文本替换的方式来完成自动命名，也就间接实现了泛型。而文本替换可以借助外部程序来完成例如一些模板语言（template language）,也可以自己来写。好在我们现在的后端是用C语言而不是Fortran95，C自身提供了宏来实现类似的功能。而对于Fortran95，就只能使用像Jinja这样的模板语言来完成泛型的支持了。

PyTorch选择了两种方案，在后端代码中利用宏来完成泛型的支持，而在中间的胶水代码部分，使用了一个用Python实现的，通过一种YAML标记语言的变体生成泛型胶水代码的生成器。不过这一部分我们着重关注第一种实现方案。下面我们继续。

回顾一下C语言的宏

关于C语言的宏，可以说是C语言中最有趣的一部分。下面关于C语言宏预处理器的介绍来自于GNU的宏命令在线文档我们只是简单的回顾，如果有疑问请详细阅读这份在线文档。

指令（Directive）

#define MACRO_NAME VALUE

定义一个宏（Macro），其名称为MACRO_NAME，值（将被展开的形式）VALUE

#line digit "finename"

改变编译器存储的当前行号`digit`和文件名`finename`为指定的行号和文件名。

#include "filepath"
#include <filepath>

预读取指定文件`filepath`，对于双引号中的文件，将在本地目录查找。对尖括号中的内容将在环境目录中查找。

宏变量

宏变量是最简单的宏，例如

#define BUFFER_SIZE 1024

在预处理器工作的时候，当后面的代码出现了`BUFFER_SIZE`，就会将其替换为`1024`，例如下面的代码

BUFFER_SIZE + 2

就会被替换为

1024 + 2

宏的定义支持续行符`\`，当一个宏命令过长时，我们可以通过使用续行符来整理你的代码。这个我们会在后面遇到。

所以也正式因为它只是简单的**文本替换**使用它也是很危险的，如果定义不当，程序可能会出现作者没有预料的行为。所以一定要小心。

有时候，我恰好和需要再次使用相同宏变量的名字，这个时候需要取消定义

#undef BUFFER_SIZE

这样在此之后预处理器就不会将`BUFFER_SIZE`替换为宏后面的内容了

宏函数

宏也可以具有参数，其行为类似于函数，但实际上很不一样。例如

#define MIN(X, Y) X < Y? X : Y

这个宏函数实现了比较其输入变量大小的功能，例如执行

// 获得最小的数字
MIN(2, 3); //注意一定要有分号

将会得到`2`，这是因为预处理器将宏`MIN`替换成了

2 < 3? 2 : 3;

这个表达式将返回`2`。可见实际上宏函数也是某种文本替换，但是不当的声明是很危险的，例如上面的这个宏，若我们

#define G 1 + 2

MIN(G, 2);

预处理器将替换为

1 + 2 < 2? 2 : 2;

这是不符合我们原本的意图的。所以我们要修改原来的定义来防止不必要的意外发生。

#define MIN(X, Y) ((X) < (Y)? (X): (Y))

还有就是一定不要在宏的最后使用分号，这是为了保证代码样式的统一。例如

#define UglyMIN(X, Y) ((X) < (Y) ? (X): (Y));

会使得在使用时没有分号，看起来和正常的语句不同。

将宏名称转换为字符串

如果我们使用宏来产生泛型，那么在抛出错误等场景可能会需要输出是哪个类型错了在宏内部可以使用`#`来产生字符串，例如

#define WARN(EXP) printf(#EXP)

会将输入的变量变为字符串再替换

WARN(test);

被替换为

printf("test");

组合名字

当我们使用不同的宏产生名字时，我们最终需要将它们组合起来。

#define CONCAT(A, B, C) A ## B ## C

例如这个宏可以用来产生`Double_Matrix_add`这个变量名

Double_Matrix CONCAT(Double, Matrix, add)(Double_Matrix *A, Double_Matrix *B);

一些预定义的宏

C语言的预处理器有一些预定义的宏

- `__FILE__` 当前输入文件名称，是一个C字符串常量，这个变量会展开为完整路径，而不是像是在`#include`中使用的简写。

- `__LINE__` 当前输入行号，是一个整数常量，这个宏的值会随着预处理器读入的行的改变而改变，所以其行为与其它预定义宏略有不同。

构建你的C泛型

首先假设我们已经有了泛型`num`，接下来我们试着按照实现一中的命名规则写出利用这个泛型构造的向量类型和`add`函数

struct NumVector
{
 num *data;
 int n;
}

// C = A + B
void NumVector_add(NumVector *C, NumVector *A, NumVector *B)
{
// check size
 if(!((C->n == A->n) && (C->n == B->n)))
 {
     exit(1); // 稍后再说产生异常的问题，先这么退出
 }

 int i,j, n;
 n = C->n;
 for(i=0; i<n; i++)
 {
     C->data[i] = A->data[i] + B->data[i];
 }
}

现在考虑如何将类似于`Num_add`的形式特例化为`FloatVector_add`等类型名称。这个可以用宏函数实现

#define Vector_(NAME) Num ## Vector_ ## NAME
#define Vector Num ## Vector

#define num float
#define Num Float

struct Vector
{
 num *data;
 int n;
};

void Vector_(add)(Vector *C, Vector *A, Vector *B)
{
//codes
}

我们期望这些宏将把以上函数和结构体替换为

struct FloatVector
{
 float *data;
 int n;
};

void FloatVector_add(FloatVector *C, FloatVector *A, FloatVector *B)
{
//codes
}

但是实际上以上代码只能产生`NumVector`的名字，这是因为C的宏定义在出现`#`和`##`时不会展开宏名，我们需要使用一个中间宏来让编译器先展开宏名，然后再组合它们。修改后如下

#define CONCAT_2_EXPAND(A, B) A ## B
#define CONCAT_2(A, B) CONCAT_2_EXPAND(A, B)
#define CONCAT_3_EXPAND(A, B, C) A ## B ## C
#define CONCAT_3(A, B, C) CONCAT_3_EXPAND(A, B, C)

#define Vector_(NAME) CONCAT_3(Num, Vector_, NAME)
#define Vector CONCAT_2(Num, Vector)

#define num float
#define Num Float

struct Vector
{
 num *data;
 int n;
};

void Vector_(add)(Vector *C, Vector *A, Vector *B)
{
//codes
}

但是这只能产生一种类型对应的函数，如果要产生多种类型的函数就需要有如下的结构

// add.c
#define CONCAT_2_EXPAND(A, B) A ## B
#define CONCAT_2(A, B) CONCAT_2_EXPAND(A, B)
#define CONCAT_3_EXPAND(A, B, C) A ## B ## C
#define CONCAT_3(A, B, C) CONCAT_3_EXPAND(A, B, C)

#define Vector_(NAME) CONCAT_3(Num, Vector_, NAME)
#define Vector CONCAT_2(Num, Vector)

#define num float
#define Num Float
struct Vector
{
 num *data;
int n;
};

void Vector_(add)(Vector *C, Vector *A, Vector *B)
{
//codes
}
#undef num
#undef Num

#define num double
#define Num Double
struct Vector
{
 num *data;
 int n;
};

void Vector_(add)(Vector *C, Vector *A, Vector *B)
{
//codes
}
#undef num
#undef Num

// etc.

这样不断复制粘贴之前的带宏命令的代码肯定是不现实的。但如果这部分泛型代码在另外一个文件里的话，那么岂不是每次从这个文件开始读取不就好了？我们现在将这部分代码分离出去，放在`generic/`文件夹下(这样就可以取相同的名字，方便记忆)，现在工程目录如下

.
├── add.c # 用来展开generic/add.c
├── add.h # 用来展开generic/add.h
├── general.h # 用来包含其它头文件
└── generic
 ├── add.c # 泛型add函数定义
 └── add.h # 泛型Vector类型的定义

现在`add.h`和`add.c`里变成了这样

// add.h
#include "general.h"

#define Vector_(NAME) CONCAT_3(Num, Vector_, NAME)
#define Vector CONCAT_2(Num, Vector)

#define num float
#define Num Float
#include "generic/add.h"
#undef num
#undef Num

#define num double
#define Num Double
#include "generic/add.h"
#undef num
#undef Num

// add.c
#include "add.h"

#define num float
#define Num Float
#include "generic/add.c"
#undef num
#undef Num

#define num double
#define Num Double
#include "generic/add.c"
#undef num
#undef Num

用`nm`命令查看一下链接库里的函数名

>>> nm *.a
add.c.o:
000000000000007e T DoubleVector_add
0000000000000000 T FloatVector_add

成功了，现在写一个测试文件来看看是否正确

#include "general.h"
#include "add.h"

int main(int argc, char const *argv[])
{
int i, n;
FloatVector *A, *B, *C;
A = (FloatVector *)malloc(sizeof(FloatVector));
B = (FloatVector *)malloc(sizeof(FloatVector));
C = (FloatVector *)malloc(sizeof(FloatVector));

n = 10;
A->data = (float *)calloc(n, sizeof(float));
B->data = (float *)calloc(n, sizeof(float));
C->data = (float *)calloc(n, sizeof(float));

A->n = n;
B->n = n;
C->n = n;

for(i=0;i<n;i++)
{
 A->data[i] = i;
 B->data[i] = 2 * i;
 C->data[i] = 0;
}

FloatVector_add(C, A, B);

for(i=0;i<n;i++)
{
printf("%f\n", C->data[i]);
}

free(A);
free(B);
free(C);
return 0;
}

正确无误！

唔，但是我们总不能每次都写一遍`num`这泛型的宏定义，我们现在把它打包到一个头文件`GenerateFloat.h`里去，然后用一个宏`GENERIC_FILE`来存储要进行特例化的文件名。首先判断是否定义了这个宏

#ifndef GENERIC_FILE
#error "You must define GENERIC_FILE before including GenerateFloat.h"
#endif

然后把刚才的特例化宏代码挪进来，加入`#line`使得编译器每次加载`GENERIC_FILE`的时候`__LINE__`都是从1开始，就好像是重新读入一样。

// GenerateFloat.h
#ifndef GENERIC_FILE
#error "You must define GENERIC_FILE before including GenerateFloat.h"
#endif

#define num float
#define Num Float
#line 1 GENERIC_FILE
#include GENERIC_FILE
#undef num
#undef Num

#define num double
#define Num Double
#line 1 GENERIC_FILE
#include GENERIC_FILE
#undef num
#undef Num

现在再修改`generic/add.h`和`generic/add.c`定义`GENERIC_FILE`这个宏

// generic/add.h
#ifndef GENERIC_FILE
#define GENERIC_FILE "generic/add.h"
#else

typedef struct Vector
{
num *data;
int n;
} Vector;

extern void Vector_(add)(Vector *C, Vector *A, Vector *B);

#endif

// generic/add.c
#ifndef GENERIC_FILE
#define GENERIC_FILE "generic/add.c"
#else

void Vector_(add)(Vector *C, Vector *A, Vector *B)
{
int i, n;
n = C->n;
for(i=0;i<n;i++)
{
C->data[i] = A->data[i] + B->data[i];
}
}

#endif

完成！

下一篇具体讲TensorApply，CPU的THNN库部分。然后讲THC也就是CUDA部分。

备注：转自 PyTorch源码浅析 PyTorch框架学习

血_影

关注

0
点赞
踩
1

收藏

觉得还不错? 一键收藏
0
评论
PyTorch源码浅析

CPU上的张量（多维数组）库TH库的实现使用了用C语言的宏产生的泛型，并且通过命名规则来产生类似面向对象的效果。这部分我们在这一章后面介绍。TH负责实现CPU上的张量（Tensor）运算，存储，内存分配。张量的内存管理和运算类型通过THTensor和THStorage两个C泛型来进行建模。张量这个数学对象被TH分解为THTensor和THStorage，THTensor提供一种查看T...
复制链接

扫一扫