Pytorch合集1

偷懒了 ai相关也都合集了~~~

一、pytorch开发基础相关

首先 PyTorch 的安装可以根据官方文档进行操作：（根据自己cuda版本不同 安装版本也不太一样啊 自己注意）

https://pytorch.org/

pip install torch torchvision

 whaosoft aiot http://143ai.com

1. PyTorch 基础

PyTorch 是数值计算方面其中一个最流行的库，同时也是机器学习研究方面最广泛使用的框架。在很多方面，它和 NumPy 都非常相似，但是它可以在不需要代码做多大改变的情况下，在 CPUs，GPUs，TPUs 上实现计算，以及非常容易实现分布式计算的操作。PyTorch 的其中一个最重要的特征就是自动微分。它可以让需要采用梯度下降算法进行训练的机器学习算法的实现更加方便，可以更高效的自动计算函数的梯度。我们的目标是提供更好的 PyTorch 介绍以及讨论使用 PyTorch 的一些最佳实践。

对于 PyTorch 第一个需要学习的就是张量（Tensors）的概念，张量就是多维数组，它和 numpy 的数组非常相似，但多了一些函数功能。

一个张量可以存储一个标量数值、一个数组、一个矩阵：

import torch# 标量数值a = torch.tensor(3)print(a)  # tensor(3)# 数组b = torch.tensor([1, 2])print(b)  # tensor([1, 2])# 矩阵c = torch.zeros([2, 2])print(c)  # tensor([[0., 0.], [0., 0.]])# 任意维度的张量d = torch.rand([2, 2, 2])

张量还可以高效的执行代数的运算。机器学习应用中最常见的运算就是矩阵乘法。例如希望将两个随机矩阵进行相乘，维度分别是  和 ，这个运算可以通过矩阵相乘运算实现（@）：

import torch
x = torch.randn([3, 5])y = torch.randn([5, 4])z = x @ y
print(z)

对于向量相加，如下所示：

z = x + y

将张量转换为 numpy 数组，可以调用 numpy() 方法：

print(z.numpy())

当然，反过来 numpy 数组转换为张量是可以的：

x = torch.tensor(np.random.normal([3, 5]))

自动微分

PyTorch 中相比 numpy 最大优点就是可以实现自动微分，这对于优化神经网络参数的应用非常有帮助。下面通过一个例子来帮助理解这个优点。

假设现在有一个复合函数：g(u(x)) ，为了计算 g 对 x 的导数，这里可以采用链式法则，即

而 PyTorch 可以自动实现这个求导的过程。

为了在 PyTorch 中计算导数，首先要创建一个张量，并设置其 requires_grad = True ，然后利用张量运算来定义函数，这里假设 u 是一个二次方的函数，而 g 是一个简单的线性函数，代码如下所示：

x = torch.tensor(1.0, requires_grad=True)
def u(x):  return x * x
def g(u):  return -u

在这个例子中，复合函数就是 ，所以导数是 ，如果 x=1 ，那么可以得到 -2 。

在 PyTorch 中调用梯度函数：​​​​​​​

dgdx = torch.autograd.grad(g(u(x)), x)[0]print(dgdx)  # tensor(-2.)

拟合曲线

为了展示自动微分有多么强大，这里介绍另一个例子。

首先假设我们有一些服从一个曲线（也就是函数 ）的样本，然后希望基于这些样本来评估这个函数 f(x) 。我们先定义一个带参数的函数:

函数的输入是 x，然后 w 是参数，目标是找到合适的参数使得下列式子成立：

实现的一个方法可以是通过优化下面的损失函数来实现：

尽管这个问题里有一个正式的函数（即 f(x) 是一个具体的函数），但这里我们还是采用一个更加通用的方法，可以应用到任何一个可微分的函数，并采用随机梯度下降法，即通过计算 L(w) 对于每个参数 w 的梯度的平均值，然后不断从相反反向移动。

利用 PyTorch 实现的代码如下所示：​​​​​​​

import numpy as npimport torch
# Assuming we know that the desired function is a polynomial of 2nd degree, we# allocate a vector of size 3 to hold the coefficients and initialize it with# random noise.w = torch.tensor(torch.randn([3, 1]), requires_grad=True)
# We use the Adam optimizer with learning rate set to 0.1 to minimize the loss.opt = torch.optim.Adam([w], 0.1)
def model(x):    # We define yhat to be our estimate of y.    f = torch.stack([x * x, x, torch.ones_like(x)], 1)    yhat = torch.squeeze(f @ w, 1)    return yhat
def compute_loss(y, yhat):    # The loss is defined to be the mean squared error distance between our    # estimate of y and its true value.    loss = torch.nn.functional.mse_loss(yhat, y)    return loss
def generate_data():    # Generate some training data based on the true function    x = torch.rand(100) * 20 - 10    y = 5 * x * x + 3    return x, y
def train_step():    x, y = generate_data()
    yhat = model(x)    loss = compute_loss(y, yhat)
    opt.zero_grad()    loss.backward()    opt.step()
for _ in range(1000):    train_step()
print(w.detach().numpy())

运行上述代码，可以得到和下面相近的结果：

[4.9924135, 0.00040895029, 3.4504161]

这和我们的参数非常接近。

上述只是 PyTorch 可以做的事情的冰山一角。很多问题，比如优化一个带有上百万参数的神经网络，都可以用 PyTorch 高效的用几行代码实现，PyTorch 可以跨多个设备和线程进行拓展，并且支持多个平台。

---

2. 将模型封装为模块

在之前的例子中，我们构建模型的方式是直接实现张量间的运算操作。但为了让代码看起来更加有组织，推荐采用 PyTorch 的 modules 模块。一个模块实际上是一个包含参数和压缩模型运算的容器。

比如，如果想实现一个线性模型 ，那么实现的代码可以如下所示：​​​​​​​

import torch
class Net(torch.nn.Module):
  def __init__(self):    super().__init__()    self.a = torch.nn.Parameter(torch.rand(1))    self.b = torch.nn.Parameter(torch.rand(1))
  def forward(self, x):    yhat = self.a * x + self.b    return yhat

使用的例子如下所示，需要实例化声明的模型，并且像调用函数一样使用它：​​​​​​​

x = torch.arange(100, dtype=torch.float32)
net = Net()y = net(x)

参数都是设置 requires_grad 为 true 的张量。

通过模型的 parameters() 方法可以很方便的访问和使用参数，如下所示：
​​

for p in net.parameters():    print(p)

现在，假设是一个未知的函数 y=5x+3+n ，注意这里的 n 是表示噪音，然后希望优化模型参数来拟合这个函数，首先可以简单从这个函数进行采样，得到一些样本数据：​​​​​​​

x = torch.arange(100, dtype=torch.float32) / 100y = 5 * x + 3 + torch.rand(100) * 0.3

和上一个例子类似，需要定义一个损失函数并优化模型的参数，如下所示：
​​​​​​

criterion = torch.nn.MSELoss()optimizer = torch.optim.SGD(net.parameters(), lr=0.01)
for i in range(10000):  net.zero_grad()  yhat = net(x)  loss = criterion(yhat, y)  loss.backward()  optimizer.step()
print(net.a, net.b) # Should be close to 5 and 3

在 PyTorch 中已经实现了很多预定义好的模块。比如 torch.nn.Linear 就是一个类似上述例子中定义的一个更加通用的线性函数，所以我们可以采用这个函数来重写我们的模型代码，如下所示：​​​​​​​

class Net(torch.nn.Module):
  def __init__(self):    super().__init__()    self.linear = torch.nn.Linear(1, 1)
  def forward(self, x):    yhat = self.linear(x.unsqueeze(1)).squeeze(1)    return yhat

这里用到了两个函数，squeeze 和 unsqueeze ，主要是torch.nn.Linear 会对一批向量而不是数值进行操作。

同样，默认调用 parameters() 会返回其所有子模块的参数：​​​​​​​

net = Net()for p in net.parameters():    print(p)

当然也有一些预定义的模块是作为包容其他模块的容器，最常用的就是 torch.nn.Sequential ，它的名字就暗示了它主要用于堆叠多个模块（或者网络层），例如堆叠两个线性网络层，中间是一个非线性函数 ReLU ，如下所示：​​​​​​​

model = torch.nn.Sequential(    torch.nn.Linear(64, 32),    torch.nn.ReLU(),    torch.nn.Linear(32, 10),)

---

3. 广播机制的优缺点

优点

PyTorch 支持广播的元素积运算。正常情况下，当想执行类似加法和乘法操作的时候，你需要确认操作数的形状是匹配的，比如无法进行一个 [3, 2] 大小的张量和 [3, 4] 大小的张量的加法操作。

但是存在一种特殊的情况：只有单一维度的时候，PyTorch 会隐式的根据另一个操作数的维度来拓展只有单一维度的操作数张量。因此，实现 [3,2] 大小的张量和 [3,1] 大小的张量相加的操作是合法的。

如下代码展示了一个加法的例子：​​​​​​​

import torch
a = torch.tensor([[1., 2.], [3., 4.]])b = torch.tensor([[1.], [2.]])# c = a + b.repeat([1, 2])c = a + b
print(c)

广播机制可以实现隐式的维度复制操作（repeat 操作），并且代码更短，内存使用上也更加高效，因为不需要存储复制的数据的结果。这个机制非常适合用于结合多个维度不同的特征的时候。

为了拼接不同维度的特征，通常的做法是先对输入张量进行维度上的复制，然后拼接后使用非线性激活函数。整个过程的代码实现如下所示：​​​​​​​

a = torch.rand([5, 3, 5])b = torch.rand([5, 1, 6])
linear = torch.nn.Linear(11, 10)
# concat a and b and apply nonlinearitytiled_b = b.repeat([1, 3, 1]) # b shape:  [5, 3, 6]c = torch.cat([a, tiled_b], 2) # c shape: [5, 3, 11]d = torch.nn.functional.relu(linear(c))
print(d.shape)  # torch.Size([5, 3, 10])

但实际上通过广播机制可以实现得更加高效，即 f(m(x+y)) 是等同于 f(mx+my) 的，也就是我们可以先分别做线性操作，然后通过广播机制来做隐式的拼接操作，如下所示：​​​​​​​

a = torch.rand([5, 3, 5])b = torch.rand([5, 1, 6])
linear1 = torch.nn.Linear(5, 10)linear2 = torch.nn.Linear(6, 10)
pa = linear1(a) # pa shape: [5, 3, 10]pb = linear2(b) # pb shape: [5, 1, 10]d = torch.nn.functional.relu(pa + pb)
print(d.shape)  # torch.Size([5, 3, 10])

实际上这段代码非常通用，可以用于任意维度大小的张量，只要它们之间是可以实现广播机制的，如下所示：​​​​​​​

class Merge(torch.nn.Module):    def __init__(self, in_features1, in_features2, out_features, activation=None):        super().__init__()        self.linear1 = torch.nn.Linear(in_features1, out_features)        self.linear2 = torch.nn.Linear(in_features2, out_features)        self.activation = activation
    def forward(self, a, b):        pa = self.linear1(a)        pb = self.linear2(b)        c = pa + pb        if self.activation is not None:            c = self.activation(c)        return c

缺点

到目前为止，我们讨论的都是广播机制的优点。但它的缺点是什么呢？原因也是出现在隐式的操作，这种做法非常不利于进行代码的调试。

这里给出一个代码例子：​​​​​​​

a = torch.tensor([[1.], [2.]])b = torch.tensor([1., 2.])c = torch.sum(a + b)
print(c)

所以上述代码的输出结果 c 是什么呢？你可能觉得是 6，但这是错的，正确答案是 12 。这是因为当两个张量的维度不匹配的时候，PyTorch 会自动将维度低的张量的第一个维度进行拓展，然后在进行元素之间的运算，所以这里会将b 先拓展为 [[1, 2], [1, 2]]，然后 a+b 的结果应该是 [[2,3], [3, 4]] ，然后sum 操作是将所有元素求和得到结果 12。

那么避免这种结果的方法就是显式的操作，比如在这个例子中就需要指定好想要求和的维度，这样进行代码调试会更简单，代码修改后如下所示：​​​​​​​

a = torch.tensor([[1.], [2.]])b = torch.tensor([1., 2.])c = torch.sum(a + b, 0)
print(c)

这里得到的 c 的结果是 [5, 7]，而我们基于结果的维度可以知道出现了错误。

这有个通用的做法，就是在做累加（ reduction ）操作或者使用 torch.squeeze 的时候总是指定好维度。

---

4. 使用好重载的运算符

和 NumPy 一样，PyTorch 会重载 python 的一些运算符来让 PyTorch 代码更简短和更有可读性。

例如，切片操作就是其中一个重载的运算符，可以更容易的对张量进行索引操作，如下所示：

z = x[begin:end]  # z = torch.narrow(0, begin, end-begin)

但需要谨慎使用这个运算符，它和其他运算符一样，也有一些副作用。正因为它是一个非常常用的运算操作，如果过度使用可以导致代码变得低效。

这里给出一个例子来展示它是如何导致代码变得低效的。这个例子中我们希望对一个矩阵手动实现行之间的累加操作：​​​​​​​

import torchimport time
x = torch.rand([500, 10])
z = torch.zeros([10])
start = time.time()for i in range(500):    z += x[i]print("Took %f seconds." % (time.time() - start))

上述代码的运行速度会非常慢，因为总共调用了 500 次的切片操作，这就是过度使用了。一个更好的做法是采用 torch.unbind 运算符在每次循环中将矩阵切片为一个向量的列表，如下所示：

z = torch.zeros([10])for x_i in torch.unbind(x):    z += x_i

这个改进会提高一些速度（在作者的机器上是提高了大约30%）。

但正确的做法应该是采用 torch.sum 来一步实现累加的操作：

z = torch.sum(x, dim=0)

这种实现速度就非常的快（在作者的机器上提高了100%的速度）。

其他重载的算数和逻辑运算符分别是：

z = -x  # z = torch.neg(x)z = x + y  # z = torch.add(x, y)z = x - yz = x * y  # z = torch.mul(x, y)z = x / y  # z = torch.div(x, y)z = x // yz = x % yz = x ** y  # z = torch.pow(x, y)z = x @ y  # z = torch.matmul(x, y)z = x > yz = x >= yz = x < yz = x <= yz = abs(x)  # z = torch.abs(x)z = x & yz = x | yz = x ^ y  # z = torch.logical_xor(x, y)z = ~x  # z = torch.logical_not(x)z = x == y  # z = torch.eq(x, y)z = x != y  # z = torch.ne(x, y)

还可以使用这些运算符的递增版本，比如 x += y 和 x **=2 都是合法的。

另外，Python 并不允许重载 and 、or 和 not 三个关键词。

---

5. 采用 TorchScript 优化运行时间

PyTorch 优化了维度很大的张量的运算操作。在 PyTorch 中对小张量进行太多的运算操作是非常低效的。所以有可能的话，将计算操作都重写为批次（batch）的形式，可以减少消耗和提高性能。而如果没办法自己手动实现批次的运算操作，那么可以采用 TorchScript 来提升代码的性能。

TorchScript 是一个 Python 函数的子集，但经过了 PyTorch 的验证，PyTorch 可以通过其 just in time(jtt) 编译器来自动优化 TorchScript 代码，提高性能。

下面给出一个具体的例子。在机器学习应用中非常常见的操作就是 batch gather ，也就是 output[i] = input[i, index[i]]。其代码实现如下所示：

​​​​​​​

import torchdef batch_gather(tensor, indices):    output = []    for i in range(tensor.size(0)):        output += [tensor[i][indices[i]]]    return torch.stack(output)

通过 torch.jit.script 装饰器来使用 TorchScript 的代码：

@torch.jit.scriptdef batch_gather_jit(tensor, indices):    output = []    for i in range(tensor.size(0)):        output += [tensor[i][indices[i]]]    return torch.stack(output)

这个做法可以提高 10% 的运算速度。

但更好的做法还是手动实现批次的运算操作，下面是一个向量化实现的代码例子，提高了 100 倍的速度：

def batch_gather_vec(tensor, indices):    shape = list(tensor.shape)    flat_first = torch.reshape(        tensor, [shape[0] * shape[1]] + shape[2:])    offset = torch.reshape(        torch.arange(shape[0]).cuda() * shape[1],        [shape[0]] + [1] * (len(indices.shape) - 1))    output = flat_first[indices + offset]    return output

---

6. 构建高效的自定义数据加载类

上一节介绍了如何写出更加高效的 PyTorch 的代码，但为了让你的代码运行更快，将数据更加高效加载到内存中也是非常重要的。幸运的是 PyTorch 提供了一个很容易加载数据的工具，即 DataLoader 。一个 DataLoader 会采用多个 workers 来同时将数据从 Dataset 类中加载，并且可以选择使用 Sampler 类来对采样数据和组成 batch 形式的数据。

如果你可以随时访问你的数据，那么使用 DataLoader 会非常简单：只需要继承 Dataset 类别并实现 __getitem__ (读取每个数据)和 __len__（返回数据集的样本数量）这两个方法。下面给出一个代码例子，如何从给定的文件夹中加载图片数据：

import globimport osimport randomimport cv2import torch
class ImageDirectoryDataset(torch.utils.data.Dataset):    def __init__(path, pattern):        self.paths = list(glob.glob(os.path.join(path, pattern)))
    def __len__(self):        return len(self.paths)
    def __item__(self):        path = random.choice(paths)        return cv2.imread(path, 1)

比如想将文件夹内所有的 jpeg 图片都加载，代码实现如下所示：​​​​​​​

dataloader = torch.utils.data.DataLoader(ImageDirectoryDataset("/data/imagenet/*.jpg"), num_workers=8)for data in dataloader:    # do something with data

这里采用了 8 个 workers 来并行的从硬盘中读取数据。这个数量可以根据实际使用机器来进行调试，得到一个最佳的数量。

当你的数据都很大或者你的硬盘读写速度很快，采用DataLoader进行随机读取数据是可行的。但也可能存在一种情况，就是使用的是一个很慢的连接速度的网络文件系统，请求单个文件的速度都非常的慢，而这可能就是整个训练过程中的瓶颈。

一个更好的做法就是将数据保存为一个可以连续读取的连续文件格式。例如，当你有非常大量的图片数据，可以采用 tar 命令将其压缩为一个文件，然后用 python 来从这个压缩文件中连续的读取图片。要实现这个操作，需要用到 PyTorch 的 IterableDataset。创建一个 IterableDataset 类，只需要实现 __iter__ 方法即可。

下面给出代码实现的例子：​​​​​​​

import tarfileimport torch
def tar_image_iterator(path):    tar = tarfile.open(self.path, "r")    for tar_info in tar:        file = tar.extractfile(tar_info)        content = file.read()        yield cv2.imdecode(content, 1)        file.close()        tar.members = []    tar.close()
class TarImageDataset(torch.utils.data.IterableDataset):    def __init__(self, path):        super().__init__()        self.path = path
    def __iter__(self):        yield from tar_image_iterator(self.path)

不过这个方法有一个问题，当使用 DataLoader 以及多个 workers 读取这个数据集的时候，会得到很多重复的数据：​​​​​​​

dataloader = torch.utils.data.DataLoader(TarImageDataset("/data/imagenet.tar"), num_workers=8)for data in dataloader:    # data contains duplicated items

这个问题主要是因为每个 worker 都会创建一个单独的数据集的实例，并且都是从数据集的起始位置开始读取数据。一种避免这个问题的办法就是不是压缩为一个tar 文件，而是将数据划分成 num_workers 个单独的 tar 文件，然后每个 worker 分别加载一个，代码实现如下所示：
​​​​​

class TarImageDataset(torch.utils.data.IterableDataset):    def __init__(self, paths):        super().__init__()        self.paths = paths
    def __iter__(self):        worker_info = torch.utils.data.get_worker_info()        # For simplicity we assume num_workers is equal to number of tar files        if worker_info is None or worker_info.num_workers != len(self.paths):            raise ValueError("Number of workers doesn't match number of files.")        yield from tar_image_iterator(self.paths[worker_info.worker_id])

所以使用例子如下所示：​​​​​​​

dataloader = torch.utils.data.DataLoader(    TarImageDataset(["/data/imagenet_part1.tar", "/data/imagenet_part2.tar"]), num_workers=2)for data in dataloader:    # do something with data

这是一种简单的避免重复数据的问题。而 tfrecord 则用了比较复杂的办法来共享数据，具体可以查看：

https://github.com/vahidk/tfrecord

---

7. PyTorch 的数值稳定性

当使用任意一个数值计算库，比如 NumPy 或者 PyTorch ，都需要知道一点，编写数学上正确的代码不一定会得到正确的结果，你需要确保这个计算是稳定的。

首先以一个简单的例子开始。从数学上来说，对任意的非零 x ，都可以知道式子  是成立的。但看看具体实现的时候，是不是总是正确的：​​​​​​​

import numpy as np
x = np.float32(1)
y = np.float32(1e-50)  # y would be stored as zeroz = x * y / y
print(z)  # prints nan

代码的运行结果是打印 nan ，原因是 y 的数值对于 float32 类型来说非常的小，这导致它的实际数值是 0 而不是 1e-50。

另一种极端情况就是 y 非常的大：​​​​​​​

y = np.float32(1e39)  # y would be stored as infz = x * y / y
print(z)  # prints nan

输出结果依然是 nan ，因为 y 太大而被存储为 inf 的情况，对于 float32 类型来说，其范围是 1.4013e-45 ~ 3.40282e+38，当超过这个范围，就会被置为 0 或者 inf。

下面是如何查看一种数据类型的数值范围：

print(np.nextafter(np.float32(0), np.float32(1)))  # prints 1.4013e-45
print(np.finfo(np.float32).max)  # print 3.40282e+38

为了让计算变得稳定，需要避免过大或者过小的数值。这看起来很容易，但这类问题是很难进行调试，特别是在 PyTorch 中进行梯度下降的时候。这不仅因为需要确保在前向传播过程中的所有数值都在使用的数据类型的取值范围内，还要保证在反向传播中也做到这一点。

下面给出一个代码例子，计算一个输出向量的 softmax，一种不好的代码实现如下所示：​​​​​​​

import torch
def unstable_softmax(logits):    exp = torch.exp(logits)    return exp / torch.sum(exp)
print(unstable_softmax(torch.tensor([1000., 0.])).numpy())  # prints [ nan, 0.]

这里计算 logits 的指数数值可能会得到超出 float32 类型的取值范围，即过大或过小的数值，这里最大的 logits 数值是 ln(3.40282e+38) = 88.7，超过这个数值都会导致 nan 。

那么应该如何避免这种情况，做法很简单。因为有 ，也就是我们可以对 logits 减去一个常量，但结果保持不变，所以我们选择logits 的最大值作为这个常数，这种做法，指数函数的取值范围就会限制为 [-inf, 0] ，然后最终的结果就是 [0.0, 1.0] 的范围，代码实现如下所示：​​​​​​​

import torch
def softmax(logits):    exp = torch.exp(logits - torch.reduce_max(logits))    return exp / torch.sum(exp)
print(softmax(torch.tensor([1000., 0.])).numpy())  # prints [ 1., 0.]

接下来是一个更复杂点的例子。

假设现在有一个分类问题。我们采用 softmax 函数对输出值 logits 计算概率。接着定义采用预测值和标签的交叉熵作为损失函数。对于一个类别分布的交叉熵可以简单定义为 ：

所以有一个不好的实现交叉熵的代码实现为：​​​​​​​

def unstable_softmax_cross_entropy(labels, logits):    logits = torch.log(softmax(logits))    return -torch.sum(labels * logits)
labels = torch.tensor([0.5, 0.5])logits = torch.tensor([1000., 0.])
xe = unstable_softmax_cross_entropy(labels, logits)
print(xe.numpy())  # prints inf

在上述代码实现中，当 softmax 结果趋向于 0，其 log 输出会趋向于无穷，这就导致计算结果的不稳定性。所以可以对其进行重写，将 softmax 维度拓展并做一些归一化的操作：​​​​​​​

def softmax_cross_entropy(labels, logits, dim=-1):    scaled_logits = logits - torch.max(logits)    normalized_logits = scaled_logits - torch.logsumexp(scaled_logits, dim)    return -torch.sum(labels * normalized_logits)
labels = torch.tensor([0.5, 0.5])logits = torch.tensor([1000., 0.])
xe = softmax_cross_entropy(labels, logits)
print(xe.numpy())  # prints 500.0

可以验证计算的梯度也是正确的：​​​​​​​

logits.requires_grad_(True)xe = softmax_cross_entropy(labels, logits)g = torch.autograd.grad(xe, logits)[0]print(g.numpy())  # prints [0.5, -0.5]

这里需要再次提醒，进行梯度下降操作的时候需要额外的小心谨慎，需要确保每个网络层的函数和梯度的范围都在合法的范围内，指数函数和对数函数在不正确使用的时候都可能导致很大的问题，它们都能将非常小的数值转换为非常大的数值，或者从很大变为很小的数值。

二、预处理相关

如何使用 Taichi Kernel 来实现 PyTorch 程序中特殊的数据预处理和自定义的算子，告别手写 CUDA，用轻巧便捷的方式提升机器学习模型算法的开发效率和灵活性。

边缘填充（Padding）是机器学习中常用的预处理方法。如在对图像执行卷积操作时，用户需要对图像边缘进行填充，以保证图像输入输出前后的尺寸不变。一般来说，填充的方法有零填充或 torch.nn.functional.pad 提供的重复填充、循环填充等其他预设模式。但有时候我们想要在边缘上填充某个特殊的纹理或者模式，却并没有一个精心优化过的 PyTorch 算子能够适配这种场景。

解决方案有两个：使用 PyTorch 或者 Python 逐个操作矩阵元素；手写 C++ 或 CUDA 代码并接入PyTorch。前者的计算效率非常低，会拖累神经网络的训练速度；后者学习曲线陡峭，实操非常麻烦，开发流程冗长。

那么，有没有更好的方案呢？接下来我们将通过一个例子，带大家体验如何用 Taichi 做一个砖墙纹理的边缘填充。

 用Taichi给PyTorch「添砖加瓦」！

第一步，我们在PyTorch中创建一个如下图所示的「砖块」。为了更好地观察填充的规律，我们给这块「砖」填充上了渐变的颜色：

填充的基本单元

第二步，我们想要在x轴上错位重复这个「砖」，也就是如下所示的效果：

由于PyTorch中没有为这样的填充提供原生的算子，为了提高运算效率，需要将padding过程改写成一系列PyTorch的原生矩阵运算：

def torch_pad(arr, tile, y):
    # image_pixel_to_coord
    arr[:, :, 0] = image_height - 1 + ph - arr[:, :, 0]
    arr[:, :, 1] -= pw
    arr1 = torch.flip(arr, (2, ))
    # map_coord
    v = torch.floor(arr1[:, :, 1] / tile_height).to(torch.int)
    u = torch.floor((arr1[:, :, 0] - v * shift_y[0]) / tile_width).to(torch.int)
    uu = torch.stack((u, u), axis=2)
    vv = torch.stack((v, v), axis=2)
    arr2 = arr1 - uu * shift_x - vv * shift_y
    # coord_to_tile_pixel 
    arr2[:, :, 1] = tile_height - 1 - arr2[:, :, 1]
    table = torch.flip(arr2, (2, ))
    table = table.view(-1, 2).to(torch.float)
    inds = table.mv(y)
    gathered = torch.index_select(tile.view(-1), 0, inds.to(torch.long))
    return gathered
   
with Timer():
    gathered = torch_pad(coords, tile, y)
    torch.cuda.synchronize(device=device)

这一系列的矩阵操作并不是特别直观，而且需要在GPU内存中保存多个中间结果矩阵。一个较为明显的缺点是显存比较小的卡上可能就跑不起来了。而如果使用Taichi，我们可以非常直接地描述这个运算：

@ti.kernel
def ti_pad(image_pixels: ti.types.ndarray(), tile: ti.types.ndarray()):
    for row, col in ti.ndrange(image_height, image_width):
        # image_pixel_to_coord
        x1, y1 = ti.math.ivec2(col - pw, image_height - 1 - row + ph) 
        # map_coord
        v: ti.i32 = ti.floor(y1 / tile_height)
        u: ti.i32 = ti.floor((x1 - v * shift_y[0]) / tile_width)
        x2, y2 = ti.math.ivec2(x1 - u * shift_x[0] - v * shift_y[0],
                 y1 - u * shift_x[1] - v * shift_y[1])
        # coord_to_tile_pixel
        x, y = ti.math.ivec2(tile_height - 1 - y2, x2) 
        image_pixels[row, col] = tile[x, y]
with Timer():
    ti_pad(image_pixels, tile）
    ti.sync()

这段代码逻辑非常简单：遍历输出图片的每个像素，计算当前像素对应到输入的「砖块」图片中的位置，最后复制该位置的颜色到这个像素。虽然看起来是在逐个写入每个像素，但Taichi会将kernel的顶层for-loop编译成高度并行的GPU代码。同时，上一段代码中我们直接把两个PyTorch的Tensor传给了 Taichi 函数ti_pad ，Taichi会直接使用PyTorch分配好的内存，不会因为两个框架间的数据交互而产生额外开销。

最后，实际的运算性能是：在RTX3090 GPU上运行时，PyTorch (v1.12.1)耗费了30.392 ms[1]，而Taichi版本的Kernel耗时仅0.267 ms[2]，Taichi相对PyTorch的加速比超过了100倍。

*加速比会因实现细节和运行硬件略有不同

事实上，上述的 PyTorch 底层实现需要启动 58 个 CUDA Kernel，而本例中 Taichi 将全部运算编译成了 1 个 CUDA Kernel。更少的 Kernel 减少了 GPU 函数启动的开销，且相比 PyTorch 实现，Taichi 节省了大量冗余的内存操作。在 GPU 上内存操作远比运算操作开销更「昂贵」，这也是非常夸张的加速比的来源。Taichi 的设计遵循了 「Megakernel」的设计准则：使用单个大的 Kernel 去完成尽可能多的运算逻辑，这与机器学习系统设计中常见的 「算子融合优化」是一样的道理。

在数据预处理问题上，一方面 Taichi 拥有更精细的操作颗粒度，能灵活适配研究人员不同的需求，另一方面 Taichi 能达到更高的计算性能，显著提升预处理部分的运行速度。当然，预处理仅仅是机器学习训练和推理过程中的一小步，对于机器学习领域的研究人员来说，有大量时间花费在模型前向和反向的计算算子中。那么对于定制高性能 ML 算子，Taichi 有什么好办法？

案例 2：定制高性能 ML 算子

和预处理遇到的问题一样，很多时候研究员用到的算子非常新或者干脆是自己发明的，在 PyTorch 中找不到良好的支持。考虑到机器学习训练和推理计算量大、成本高昂，很多研究员不得不去学习 CUDA 并尽力调优，以提升计算效率。但 CUDA 代码编写难度大，调试困难，会拖慢模型迭代速度。

有一篇知乎文章[3]讲述了一个精彩的例子：作者开发了 RWKV 语言模型，使用了一个类似一维的深度卷积（depthwise convolution）的自定义算子。这个算子本身计算量不大，但是因为 PyTorch 中没有原生支持，跑得特别慢。为了解决计算性能的问题，作者编写 CUDA 代码并且采用循环合并、Shared Memory 等多种技巧来优化，最终性能达到了 PyTorch 实现的 20 倍性能。

参考这篇文章和发布的 CUDA 代码，我们也使用相同的优化手段实现了对应的 Taichi 版本。那么 Taichi 在这个例子中性能如何呢？请看下图：

RTX3080 上的 RWKV 运算时间，单位是毫秒，越低越好。Baseline 代表代码直接实现算法，不做任何优化。v1-v3 代表不同的优化版本。CUDA 实现代码见[4], Taichi 实现代码见[5]。

我们可以看到，在使用同样的优化技术的前提下， Taichi 版本达到了非常接近 CUDA 的性能，甚至某些情况下还略快一点。这样的性能水平是如何用 Taichi 实现的呢？会有多么简单呢？

接下来我们就以 Baseline 版本为例，体验如何用 Taichi 轻松实现深度卷积算子！算子本身的运算过程很简单：遍历两个输入 Tensor w 和 k, 把它们对应位置的元素乘起来，通过一个累加循环计算出 s 并存进输出 Tensor out。

 Python 实现（很慢很好懂）

def run_formula_very_slow(w, k, B, C, T, eps):
    out = torch.empty((B, C, T), device='cpu')
    for b in range(B):
        for c in range(C):
            for t in range(T):
                s = eps
                for u in range(t-T+1, t+1):
                    s += w[c][0][(T-1)-(t-u)] * k[b][c][u+T-1]
                out[b][c][t] = s
    return out

这段代码非常直观好懂，但它运行速度如此之慢，以至于测试出来的数据都没办法把它放进上面那张图里...

PyTorch 实现（一般慢不好懂）

out = eps + F.conv1d(nn.ZeroPad2d((T-1, 0, 0, 0))(k), w.unsqueeze(1), groups=C)

从上面的 Python 代码写出 PyTorch 的这一行还是非常有难度的，要对 PyTorch 的这几个算子底层的运算逻辑很熟悉才能写得出来。

 Taichi 实现（很快很好懂）

@ti.kernel
def taichi_forward_v0(
        out: ti.types.ndarray(field_dim=3),
        w: ti.types.ndarray(field_dim=3),
        k: ti.types.ndarray(field_dim=3),
        eps: ti.f32):
        
    for b, c, t in out:
        s = eps
        for u in range(t-T+1, t+1):
            s += w[c, 0, (T-1)-(t-u)] * k[b, c, u+T-1]
        out[b, c, t] = s

Taichi 代码和 Python 代码几乎完全一致，而且不用考虑并行、指针偏移计算等等各种编程细节，就可以达到和 CUDA 接近的性能，在开发效率上具有很大的优势。作为对比，我们也把对应的 CUDA 版本放在后面，有兴趣的读者可以看一下。CUDA 版本的可读性差了很多。它的外层循环是隐含在线程并行的逻辑里。另外，它的指针的偏移计算比较复杂，每个元素在矩阵中的位置没办法很直观地看出来，需要做一些推演才能完全理解这段代码，当算法再复杂一些的时候就很容易写错。

__global__ void kernel_forward(const float* w, const float* k, float* x,
                               const float eps, const int B, const int C, const int T)
{
    const int i = blockIdx.y;
    const int t = threadIdx.x;
    float s = eps;
    const float* www = w + (i % C) * T + (T - 1) - t;
    const float* kk = k + i * T;
    for (int u = 0; u <= t; u++){
        s += www[u] * kk[u];
    }
    x[i * T + t] = s;
}

更重要的是，CUDA 代码需要编译环境才能运行。如果提前编译成动态库，又需要对齐 CUDA 运行时环境。环境配置、Python 接口封装等等都需要耗费精力去做。而 Taichi 代码本身就是一小段 Python 代码，可以通过 pip 安装管理，与 PyTorch 完全一致，简单了很多，其良好的可复现性，也便于机器学习开发者开源、分享代码。更好的性能、更敏捷的开发效率、更便捷的分享方式，共同构成了使用 Taichi 开发自定义 ML 算子的显著优势。

总结 

虽然 PyTorch 可以高效完成机器学习中大部分的运算任务，但仍有许多算子没有实现或者运算效率无法满足需求。作为嵌在 Python 中的高性能编程语言，Taichi 易于编写、内存消耗小，计算性能接近手写 CUDA。本文展示的两个例子，正是结合 Taichi 和 PyTorch 之所长，解决了预处理算子和新算法中的算子的高性能编程问题，同时 Taichi 和 Pytorch Tensor 零开销交互的特性也省去了编写「脚手架」代码的时间，极大地提升了开发效率。希望 Taichi 可以将机器学习研究人员从繁复晦涩的高性能代码编写、验证、调优中解放出来，专注于算法本身，创造出更多有趣的东西 :-)