简介
多年来，我们一直在使用 PyTorch 来构建和训练深度学习模型。尽管对其语法和规则已经了如指掌，但内心深处总有一些疑问：这些操作背后究竟发生了什么？它们是如何运作的？如果询问你如何在 PyTorch 中创建和训练一个模型，你可能会想到类似以下的代码：
import torch
import torch.nn as nn
import torch.optim as optim
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.fc1 = nn.Linear(1, 10)
self.sigmoid = nn.Sigmoid()
self.fc2 = nn.Linear(10, 1)
def forward(self, x):
out = self.fc1(x)
out = self.sigmoid(out)
out = self.fc2(out)
return out
...
model = MyModel().to(device)
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)
for epoch in range(epochs):
for x, y in ...
x = x.to(device)
y = y.to(device)
outputs = model(x)
loss = criterion(outputs, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
如果询问你这个回溯步骤是如何运作的，你会如何回答？再比如，当你改变一个张量的形状时，背后发生了什么？数据是否在内部被重新组织？这个过程是如何进行的？PyTorch为何能够运行得如此迅速？它是如何处理GPU运算的？这些问题一直让我充满好奇，我相信它们同样也引起了你的兴趣。因此，为了更深入地掌握这些概念，有什么方法能比自己动手从零开始构建一个张量库更有效呢？这就是你在本文中将要学习的内容！
Tensor
要创建一个张量库，你首先必须掌握的概念当然是：张量是什么？
你可能已经有一个直观的理解，即张量是一个包含数值的多维数据结构的数学概念。但在这儿，我们得从计算的角度来理解如何构建这种数据结构。我们可以将张量想象成由两部分组成：一部分是数据本身，另一部分是描述张量特征的元数据，比如它的维度形状，或者它存储在哪种设备上（比如CPU内存或GPU内存）。
还有一个你可能不太熟悉的元数据概念，称为步长（stride）。理解这个概念对于深入掌握张量数据重排的内部工作原理至关重要，因此我们有必要对其进行更深入的探讨。
设想一个形状为 [4, 8] 的二维张量，如下图所示。
张量的数据（即浮点数）实际上在内存中存储为一维数组：
所以，为了将这个一维数组表示为多维张量，我们利用了步长（strides）的概念。简单来说，思路是这样的：
我们有一个 4 行 8 列的矩阵。假设所有元素都是按行顺序存储在一维数组中，如果我们想访问位于 [2, 3] 的元素值，我们需要跳过 2 行（每行有 8 个元素）再加上 3 个位置。用数学语言描述，我们需要在一维数组中跳过 3 + 2 * 8 个元素来访问它：
这个“8”代表的是第二维度的步长（stride）。简单来说，它告诉我们在数组中需要跳过多少个元素才能“移动”到第二维度上的其他位置。因此，要访问形状为 [shape_0, shape_1] 的二维张量中的元素 [i, j]，我们实际上需要定位到 j + i * shape_1 的位置。
接下来，让我们设想一个三维张量的情况：
你可以将这个三维张量视作矩阵的序列。比如，你可以把这个形状为 [5, 4, 8] 的张量看作是 5 个形状为 [4, 8] 的矩阵。现在，要获取位于 [1, 2, 7] 的元素，你需要跳过 1 个完整的 [4,8] 形状的矩阵，2 行 [8] 形状的行，以及 7 列 [1] 形状的列。所以，你需要在一维数组中跳过 (1 4 8) + (2 8) + (7 1) 个位置。
因此，要访问 1-D 数据数组上具有 [shape_0, shape_1, shape_2] 的 3-D 张量的元素 [i][j][k]，您可以执行以下操作：
shape_1 * shape_2 是第一维度的步幅，shape_2 是第二维度的步幅，1 是第三维度的步幅。然后，为了概括：
其中每个维度的步幅可以使用下一维度张量形状的乘积来计算：
然后我们设置 stride[n-1] = 1。在形状为 [5, 4, 8] 的张量示例中，步长 = [4*8, 8, 1] = [32, 8, 1] 您可以自行测试：
import torch
torch.rand([5, 4, 8]).stride()
#(32, 8, 1)
好的，但我们为什么要使用形状和步长呢？这个概念不仅可以用于访问存储为一维数组形式的N维张量的元素，还可以非常方便地用来调整张量的布局。
比如，当你想要改变一个张量的形状时，只需指定新的形状，并据此计算出新的步长即可！（这样做是因为新的形状确保了元素的总数保持不变）
import torch
t = torch.rand([5, 4, 8])
print(t.shape)
# [5, 4, 8]
print(t.stride())
# [32, 8, 1]
new_t = t.reshape([4, 5, 2, 2, 2])
print(new_t.shape)
# [4, 5, 2, 2, 2]
print(new_t.stride())
# [40, 8, 4, 2, 1]
在内部，张量仍然以一维数组的形式存储。重塑操作并没有改变数组中元素的排列顺序！这真是令人惊叹，对吧？你可以通过下面的函数来亲自验证这一点，该函数能够访问 PyTorch 内部的一维数组：
import ctypes
def print_internal(t: torch.Tensor):
print(
torch.frombuffer(
ctypes.string_at(t.data_ptr(), t.storage().nbytes()), dtype=t.dtype
)
)
print_internal(t)
# [0.0752, 0.5898, 0.3930, 0.9577, 0.2276, 0.9786, 0.1009, 0.138, ...
print_internal(new_t)
# [0.0752, 0.5898, 0.3930, 0.9577, 0.2276, 0.9786, 0.1009, 0.138, ...
或者例如，您想要转置两个轴。在内部，您只需交换各自的步幅即可！
t = torch.arange(0, 24).reshape(2, 3, 4)
print(t)
# [[[ 0, 1, 2, 3],
# [ 4, 5, 6, 7],
# [ 8, 9, 10, 11]],
# [[12, 13, 14, 15],
# [16, 17, 18, 19],
# [20, 21, 22, 23]]]
print(t.shape)
# [2, 3, 4]
print(t.stride())
# [12, 4, 1]
new_t = t.transpose(0, 1)
print(new_t)
# [[[ 0, 1, 2, 3],
# [12, 13, 14, 15]],
# [[ 4, 5, 6, 7],
# [16, 17, 18, 19]],
# [[ 8, 9, 10, 11],
# [20, 21, 22, 23]]]
print(new_t.shape)
# [3, 2, 4]
print(new_t.stride())
# [4, 12, 1]
如果打印内部数组，两者具有相同的值：
print_internal(t)
# [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]
print_internal(new_t)
# [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]
但是，新张量 new_t 的步长现在与我之前展示的公式不一致。这是因为张量现在不再连续。也就是说，尽管内部数组本身没有变化，它在内存中的存储顺序与张量的实际排列顺序并不一致。
t.is_contiguous()
# True
new_t.is_contiguous()
# False
这意味着，如果我们要按顺序访问那些不连续的元素，效率会比较低（因为这些元素在内存中的存储并不是连续的）。为了提高效率，我们可以通过以下方式来解决：
new_t_contiguous = new_t.contiguous()
print(new_t_contiguous.is_contiguous())
# True
如果我们分析内部数组，它的顺序现在与实际张量顺序匹配，这可以提供更好的内存访问效率：
print(new_t)
# [[[ 0, 1, 2, 3],
# [12, 13, 14, 15]],
# [[ 4, 5, 6, 7],
# [16, 17, 18, 19]],
# [[ 8, 9, 10, 11],
# [20, 21, 22, 23]]]
print_internal(new_t)
# [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]
print_internal(new_t_contiguous)
# [ 0, 1, 2, 3, 12, 13, 14, 15, 4, 5, 6, 7, 16, 17, 18, 19, 8, 9, 10, 11, 20, 21, 22, 23]
构建自己的张量库（Norch）
现在我们已经明白了张量是如何构建的，接下来就让我们着手开发我们的库。我打算给它命名为Norch，这个名字既表示“非PyTorch”，也隐含了我自己的姓氏Nogueira 😁首先需要了解的是，尽管PyTorch是通过Python接口使用的，但其核心实际上是用C/C++编写的。因此，我们首先需要开发我们自己的C/C++底层功能。我们可以首先定义一个结构体来存储张量的数据和元数据，并编写一个函数来创建这个结构体的实例。
//norch/csrc/tensor.cpp
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
typedef struct {
float* data;
int* strides;
int* shape;
int ndim;
int size;
char* device;
} Tensor;
Tensor* create_tensor(float* data, int* shape, int ndim) {
Tensor* tensor = (Tensor*)malloc(sizeof(Tensor));
if (tensor == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
tensor->data = data;
tensor->shape = shape;
tensor->ndim = ndim;
tensor->size = 1;
for (int i = 0; i < ndim; i++) {
tensor->size *= shape[i];
}
tensor->strides = (int*)malloc(ndim * sizeof(int));
if (tensor->strides == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
int stride = 1;
for (int i = ndim - 1; i >= 0; i--) {
tensor->strides[i] = stride;
stride *= shape[i];
}
return tensor;
}
为了访问某些元素，我们可以利用步幅，正如我们之前学到的：
//norch/csrc/tensor.cpp
float get_item(Tensor* tensor, int* indices) {
int index = 0;
for (int i = 0; i < tensor->ndim; i++) {
index += indices[i] * tensor->strides[i];
}
float result;
result = tensor->data[index];
return result;
}
现在，我们可以创建张量运算。我将展示一些示例，您可以在本文末尾链接的存储库中找到完整版本。
//norch/csrc/cpu.cpp
void add_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
for (int i = 0; i < tensor1->size; i++) {
result_data[i] = tensor1->data[i] + tensor2->data[i];
}
}
void sub_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
for (int i = 0; i < tensor1->size; i++) {
result_data[i] = tensor1->data[i] - tensor2->data[i];
}
}
void elementwise_mul_tensor_cpu(Tensor* tensor1, Tensor* tensor2, float* result_data) {
for (int i = 0; i < tensor1->size; i++) {
result_data[i] = tensor1->data[i] * tensor2->data[i];
}
}
void assign_tensor_cpu(Tensor* tensor, float* result_data) {
for (int i = 0; i < tensor->size; i++) {
result_data[i] = tensor->data[i];
}
}
...
之后我们就可以创建其他张量函数来调用这些操作：
//norch/csrc/tensor.cpp
Tensor* add_tensor(Tensor* tensor1, Tensor* tensor2) {
if (tensor1->ndim != tensor2->ndim) {
fprintf(stderr, "Tensors must have the same number of dimensions %d and %d for addition\n", tensor1->ndim, tensor2->ndim);
exit(1);
}
int ndim = tensor1->ndim;
int* shape = (int*)malloc(ndim * sizeof(int));
if (shape == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
for (int i = 0; i < ndim; i++) {
if (tensor1->shape[i] != tensor2->shape[i]) {
fprintf(stderr, "Tensors must have the same shape %d and %d at index %d for addition\n", tensor1->shape[i], tensor2->shape[i], i);
exit(1);
}
shape[i] = tensor1->shape[i];
}
float* result_data = (float*)malloc(tensor1->size * sizeof(float));
if (result_data == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
add_tensor_cpu(tensor1, tensor2, result_data);
return create_tensor(result_data, shape, ndim, device);
}
如前所述，张量重塑不会修改内部数据数组：
//norch/csrc/tensor.cpp
Tensor* reshape_tensor(Tensor* tensor, int* new_shape, int new_ndim) {
int ndim = new_ndim;
int* shape = (int*)malloc(ndim * sizeof(int));
if (shape == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
for (int i = 0; i < ndim; i++) {
shape[i] = new_shape[i];
}
// Calculate the total number of elements in the new shape
int size = 1;
for (int i = 0; i < new_ndim; i++) {
size *= shape[i];
}
// Check if the total number of elements matches the current tensor's size
if (size != tensor->size) {
fprintf(stderr, "Cannot reshape tensor. Total number of elements in new shape does not match the current size of the tensor.\n");
exit(1);
}
float* result_data = (float*)malloc(tensor->size * sizeof(float));
if (result_data == NULL) {
fprintf(stderr, "Memory allocation failed\n");
exit(1);
}
assign_tensor_cpu(tensor, result_data);
return create_tensor(result_data, shape, ndim, device);
}
虽然我们已经能够执行一些张量操作，但不应该有人需要直接使用C/C++来运行这些操作，对吗？现在，让我们着手开发Python接口！在Python中执行C/C++代码有很多选择，例如Pybind11和Cython。在我们的示例中，我选择使用ctypes。下面是ctypes的基本结构：
//C code
#include <stdio.h>
float add_floats(float a, float b) {
return a + b;
}
# Compile
gcc -shared -o add_floats.so -fPIC add_floats.c
# Python code
import ctypes
# Load the shared library
lib = ctypes.CDLL('./add_floats.so')
# Define the argument and return types for the function
lib.add_floats.argtypes = [ctypes.c_float, ctypes.c_float]
lib.add_floats.restype = ctypes.c_float
# Convert python float to c_float type
a = ctypes.c_float(3.5)
b = ctypes.c_float(2.2)
# Call the C function
result = lib.add_floats(a, b)
print(result)
# 5.7
正如您所看到的，它非常直观。编译 C/C++ 代码后，您可以非常轻松地通过 ctypes 在 Python 上使用它。您只需要定义函数的参数和返回 c_types，将变量转换为其各自的 c_types 并调用该函数。对于更复杂的类型，例如数组（浮点列表），您可以使用指针。
data = [1.0, 2.0, 3.0]
data_ctype = (ctypes.c_float * len(data))(*data)
lib.some_array_func.argstypes = [ctypes.POINTER(ctypes.c_float)]
...
lib.some_array_func(data)
对于结构类型，我们可以创建自己的 c_type：
class CustomType(ctypes.Structure):
_fields_ = [
('field1', ctypes.POINTER(ctypes.c_float)),
('field2', ctypes.POINTER(ctypes.c_int)),
('field3', ctypes.c_int),
]
# Can be used as ctypes.POINTER(CustomType)
经过简短的解释后，让我们为张量 C/C++ 库构建 Python 包装器！
# norch/tensor.py
import ctypes
class CTensor(ctypes.Structure):
_fields_ = [
('data', ctypes.POINTER(ctypes.c_float)),
('strides', ctypes.POINTER(ctypes.c_int)),
('shape', ctypes.POINTER(ctypes.c_int)),
('ndim', ctypes.c_int),
('size', ctypes.c_int),
]
class Tensor:
os.path.abspath(os.curdir)
_C = ctypes.CDLL("COMPILED_LIB.so"))
def __init__(self):
data, shape = self.flatten(data)
self.data_ctype = (ctypes.c_float * len(data))(*data)
self.shape_ctype = (ctypes.c_int * len(shape))(*shape)
self.ndim_ctype = ctypes.c_int(len(shape))
self.shape = shape
self.ndim = len(shape)
Tensor._C.create_tensor.argtypes = [ctypes.POINTER(ctypes.c_float), ctypes.POINTER(ctypes.c_int), ctypes.c_int]
Tensor._C.create_tensor.restype = ctypes.POINTER(CTensor)
self.tensor = Tensor._C.create_tensor(
self.data_ctype,
self.shape_ctype,
self.ndim_ctype,
)
def flatten(self, nested_list):
"""
This method simply convert a list type tensor to a flatten tensor with its shape
Example:
Arguments:
nested_list: [[1, 2, 3], [-5, 2, 0]]
Return:
flat_data: [1, 2, 3, -5, 2, 0]
shape: [2, 3]
"""
def flatten_recursively(nested_list):
flat_data = []
shape = []
if isinstance(nested_list, list):
for sublist in nested_list:
inner_data, inner_shape = flatten_recursively(sublist)
flat_data.extend(inner_data)
shape.append(len(nested_list))
shape.extend(inner_shape)
else:
flat_data.append(nested_list)
return flat_data, shape
flat_data, shape = flatten_recursively(nested_list)
return flat_data, shape
现在我们包含 Python 张量操作来调用 C/C++ 操作。
# norch/tensor.py
def __getitem__(self, indices):
"""
Access tensor by index tensor[i, j, k...]
"""
if len(indices) != self.ndim:
raise ValueError("Number of indices must match the number of dimensions")
Tensor._C.get_item.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(ctypes.c_int)]
Tensor._C.get_item.restype = ctypes.c_float
indices = (ctypes.c_int * len(indices))(*indices)
value = Tensor._C.get_item(self.tensor, indices)
return value
def reshape(self, new_shape):
"""
Reshape tensor
result = tensor.reshape([1,2])
"""
new_shape_ctype = (ctypes.c_int * len(new_shape))(*new_shape)
new_ndim_ctype = ctypes.c_int(len(new_shape))
Tensor._C.reshape_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(ctypes.c_int), ctypes.c_int]
Tensor._C.reshape_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.reshape_tensor(self.tensor, new_shape_ctype, new_ndim_ctype)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = new_shape.copy()
result_data.ndim = len(new_shape)
result_data.device = self.device
return result_data
def __add__(self, other):
"""
Add tensors
result = tensor1 + tensor2
"""
if self.shape != other.shape:
raise ValueError("Tensors must have the same shape for addition")
Tensor._C.add_tensor.argtypes = [ctypes.POINTER(CTensor), ctypes.POINTER(CTensor)]
Tensor._C.add_tensor.restype = ctypes.POINTER(CTensor)
result_tensor_ptr = Tensor._C.add_tensor(self.tensor, other.tensor)
result_data = Tensor()
result_data.tensor = result_tensor_ptr
result_data.shape = self.shape.copy()
result_data.ndim = self.ndim
result_data.device = self.device
return result_data
# Include the other operations:
# __str__
# __sub__ (-)
# __mul__ (*)
# __matmul__ (@)
# __pow__ (**)
# __truediv__ (/)
# log
# ...
您现在就可以运行代码并开始执行一些张量运算！
import norch
tensor1 = norch.Tensor([[1, 2, 3], [3, 2, 1]])
tensor2 = norch.Tensor([[3, 2, 1], [1, 2, 3]])
result = tensor1 + tensor2
print(result[0, 0])
# 4
