一、pytorch开发基础相关
首先 PyTorch 的安装可以根据官方文档进行操作：（根据自己cuda版本不同 安装版本也不太一样啊 自己注意）
pip install torch torchvision1. 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 np
import 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
trueparameters()for p in net.parameters():
print(p)现在，假设是一个未知的函数 y=5x+3+n ，注意这里的 n 是表示噪音，然后希望优化模型参数来拟合这个函数，首先可以简单从这个函数进行采样，得到一些样本数据：
x = torch.arange(100, dtype=torch.float32) / 100
y = 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 nonlinearity
tiled_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 torch
import 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 - y
z = x * y # z = torch.mul(x, y)
z = x / y # z = torch.div(x, y)
z = x // y
z = x % y
z = x ** y # z = torch.pow(x, y)
z = x @ y # z = torch.matmul(x, y)
z = x > y
z = x >= y
z = x < y
z = x <= y
z = abs(x) # z = torch.abs(x)
z = x & y
z = x | y
z = 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 torch
def batch_gather(tensor, indices):
output = []
for i in range(tensor.size(0)):
output += [tensor[i][indices[i]]]
return torch.stack(output)torch.jit.script@torch.jit.script
def 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 output6. 构建高效的自定义数据加载类
上一节介绍了如何写出更加高效的 PyTorch 的代码，但为了让你的代码运行更快，将数据更加高效加载到内存中也是非常重要的。幸运的是 PyTorch 提供了一个很容易加载数据的工具，即 DataLoader 。一个 DataLoader 会采用多个 workers 来同时将数据从 Dataset 类中加载，并且可以选择使用 Sampler 类来对采样数据和组成 batch 形式的数据。
如果你可以随时访问你的数据，那么使用 DataLoader 会非常简单：只需要继承 Dataset 类别并实现 __getitem__ (读取每个数据)和 __len__（返回数据集的样本数量）这两个方法。下面给出一个代码例子，如何从给定的文件夹中加载图片数据：
import glob
import os
import random
import cv2
import 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 dataworkers当你的数据都很大或者你的硬盘读写速度很快，采用DataLoader进行随机读取数据是可行的。但也可能存在一种情况，就是使用的是一个很慢的连接速度的网络文件系统，请求单个文件的速度都非常的慢，而这可能就是整个训练过程中的瓶颈。
一个更好的做法就是将数据保存为一个可以连续读取的连续文件格式。例如，当你有非常大量的图片数据，可以采用 tar 命令将其压缩为一个文件，然后用 python 来从这个压缩文件中连续的读取图片。要实现这个操作，需要用到 PyTorch 的 IterableDataset。创建一个 IterableDataset 类，只需要实现 __iter__ 方法即可。
下面给出代码实现的例子：
import tarfile
import 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 itemsworker
tar
num_workers
tar
workerclass 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 datatfrecordhttps://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 zero
z = x * y / y
print(z) # prints nannan
y
float32另一种极端情况就是 y 非常的大：
y = np.float32(1e39) # y would be stored as inf
z = x * y / y
print(z) # prints nannan
y
inf
float32
1.4013e-45 ~ 3.40282e+38下面是如何查看一种数据类型的数值范围：
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] = sTaichi 代码和 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 可以将机器学习研究人员从繁复晦涩的高性能代码编写、验证、调优中解放出来，专注于算法本身，创造出更多有趣的东西 :-)
二、在PyTorch中创建和使用Python自定义操作符
关于如何在PyTorch中创建和使用Python自定义操作符（Custom Operators）的教程
在vllm里面看到flash attention包了一层@torch.library.custom_op装饰器（https://github.com/vllm-project/vllm/pull/7536），查阅了一下资料，发现这个是torch 2.4之后的新feature，防止打算torch compile的graph，翻译一下官方教程稍微了解一下这个用法。来源：https://pytorch.org/tutorials/advanced/python_custom_ops.html
Python Custom Operators 教程
这个教程介绍了Python自定义运算符的主题。它列出了我们将从这一教程中学习到的内容,包括如何将用Python编写的自定义运算符与PyTorch集成,以及如何使用torch.library.opcheck来测试自定义运算符。所需的先决条件是安装了PyTorch 2.4或更高版本。
PyTorch提供了大量可以在Tensor上运行的运算符(例如torch.add、torch.sum等)。但是,您可能希望在PyTorch中使用一个新的自定义运算符,可能是由第三方库编写的。本教程展示了如何封装Python函数,使它们的行为类似于PyTorch原生运算符。创建PyTorch中的自定义运算符的原因可能包括:
将任意Python函数视为不透明的可调用对象,与torch.compile相对应(即防止torch.compile跟踪进入函数)。
为任意Python函数添加训练支持。
请注意,如果您的操作可以表示为现有PyTorch运算符的组合,那么通常就不需要使用自定义运算符(例如,支持自动微分的运算应该可以直接工作)。
例子：将PIL库的crop功能封装为一个自定义运算符
假设我们在使用PIL的crop操作
import torch
from torchvision.transforms.functional import to_pil_image, pil_to_tensor
import PIL
import IPython
import matplotlib.pyplot as plt
def crop(pic, box):
img = to_pil_image(pic.cpu())
cropped_img = img.crop(box)
return pil_to_tensor(cropped_img).to(pic.device) / 255.
def display(img):
plt.imshow(img.numpy().transpose((1, 2, 0)))
img = torch.ones(3, 64, 64)
img *= torch.linspace(0, 1, steps=64) * torch.linspace(0, 1, steps=64).unsqueeze(-1)
display(img)
cropped_img = crop(img, (10, 10, 50, 50))
display(cropped_img)
crop功能无法被torch.compile有效地开箱即用处理:torch.compile在无法处理的函数上会引发"图中断"(https://pytorch.org/docs/stable/torch.compiler_faq.html#graph-breaks),而图中断会导致性能下降。以下代码通过引发错误来演示这一点(如果发生图中断,torch.compile(with fullgraph=True)会引发错误)。
@torch.compile(fullgraph=True)
def f(img):
return crop(img, (10, 10, 50, 50))
# The following raises an error. Uncomment the line to see it.
# cropped_img = f(img)为了能在torch.compile中使用crop作为黑盒操作,我们需要做两件事:
- 将该函数封装为一个PyTorch自定义运算符。
- 为该运算符添加"FakeTensor kernel"(又称"meta kernel")。给定输入Tensor的元数据(例如形状),此函数说明如何计算输出Tensor的元数据。
from typing import Sequence
# Use torch.library.custom_op to define a new custom operator.
# If your operator mutates any input Tensors, their names must be specified
# in the ``mutates_args`` argument.
@torch.library.custom_op("mylib::crop", mutates_args=())
def crop(pic: torch.Tensor, box: Sequence[int]) -> torch.Tensor:
img = to_pil_image(pic.cpu())
cropped_img = img.crop(box)
return (pil_to_tensor(cropped_img) / 255.).to(pic.device, pic.dtype)
# Use register_fake to add a ``FakeTensor`` kernel for the operator
@crop.register_fake
def _(pic, box):
channels = pic.shape[0]
x0, y0, x1, y1 = box
return pic.new_empty(channels, y1 - y0, x1 - x0)做了上述操作之后,crop现在可以在不产生图中断的情况下正常工作了。
@torch.compile(fullgraph=True)
def f(img):
return crop(img, (10, 10, 50, 50))
cropped_img = f(img)
display(img)
display(cropped_img)
为crop添加训练支持
使用torch.library.register_autograd为运算符添加训练支持。相比直接使用torch.autograd.Function,优先使用这种方式;因为autograd.Function与PyTorch运算符注册API组合使用时,可能会在与torch.compile组合时导致无声的不正确性。crop的梯度公式本质上是PIL.paste(我们把推导留作读者练习)。让我们首先将paste封装为一个自定义运算符:
@torch.library.custom_op("mylib::paste", mutates_args=())
def paste(im1: torch.Tensor, im2: torch.Tensor, coord: Sequence[int]) -> torch.Tensor:
assert im1.device == im2.device
assert im1.dtype == im2.dtype
im1_pil = to_pil_image(im1.cpu())
im2_pil = to_pil_image(im2.cpu())
PIL.Image.Image.paste(im1_pil, im2_pil, coord)
return (pil_to_tensor(im1_pil) / 255.).to(im1.device, im1.dtype)
@paste.register_fake
def _(im1, im2, coord):
assert im1.device == im2.device
assert im1.dtype == im2.dtype
return torch.empty_like(im1)现在让我们使用register_autograd来为crop指定梯度公式:
def backward(ctx, grad_output):
grad_input = grad_output.new_zeros(ctx.pic_shape)
grad_input = paste(grad_input, grad_output, ctx.coords)
return grad_input, None
def setup_context(ctx, inputs, output):
pic, box = inputs
ctx.coords = box[:2]
ctx.pic_shape = pic.shape
crop.register_autograd(backward, setup_cnotallow=setup_context)注意,backward必须是由PyTorch可理解的运算符组成,这也是我们将paste封装为自定义运算符而不直接使用PIL的paste的原因。
img = img.requires_grad_()
result = crop(img, (10, 10, 50, 50))
result.sum().backward()
display(img.grad)
这是正确的梯度,在裁剪区域内是1(白色),在未使用的区域内是0(黑色)。
测试Python自定义运算符
使用torch.library.opcheck来测试自定义运算符是否正确注册。这不会测试梯度是否在数学上正确,请单独编写测试(手动测试或使用torch.autograd.gradcheck)。要使用opcheck,请传入一组示例输入用于测试。如果你的运算符支持训练,那么示例应该包括需要计算梯度的Tensor。如果你的运算符支持多个设备,那么示例应该包括来自每个设备的Tensor。
examples = [
[torch.randn(3, 64, 64), [0, 0, 10, 10]],
[torch.randn(3, 91, 91, requires_grad=True), [10, 0, 20, 10]],
[torch.randn(3, 60, 60, dtype=torch.double), [3, 4, 32, 20]],
[torch.randn(3, 512, 512, requires_grad=True, dtype=torch.double), [3, 4, 32, 45]],
]
for example in examples:
torch.library.opcheck(crop, example)可变的Python自定义运算符
你也可以将一个会修改其输入的Python函数封装为自定义运算符。修改输入的函数很常见,因为这是许多low-level kernel编写的方式;例如,计算sin的kernel可能会修改输入,并将输出张量赋值为input.sin()。我们将使用numpy.sin来演示一个可变的Python自定义运算符的示例。
import numpy as np
@torch.library.custom_op("mylib::numpy_sin", mutates_args={"output"}, device_types="cpu")
def numpy_sin(input: torch.Tensor, output: torch.Tensor) -> None:
assert input.device == output.device
assert input.device.type == "cpu"
input_np = input.numpy()
output_np = output.numpy()
np.sin(input_np, out=output_np)由于该运算符没有返回值,因此不需要注册FakeTensor kernel(元kernel)就可以在torch.compile中正常工作。
@torch.compile(fullgraph=True)
def f(x):
out = torch.empty(3)
numpy_sin(x, out)
return out
x = torch.randn(3)
y = f(x)
assert torch.allclose(y, x.sin())这里是一次opcheck运行的结果,告诉我们确实正确注册了该kernel。如果我们忘记添加输出到mutates_args,例如,opcheck将会报错。
总结
在本教程中,我们学习了如何使用torch.library.custom_op创建一个与PyTorch子系统如torch.compile和autograd协同工作的Python自定义运算符。
本教程提供了自定义运算符的基本介绍。更多详细信息,请参阅:
- torch.library文档: https://pytorch.org/docs/stable/library.html
- 自定义运算符手册: https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html#the-custom-operators-manual
三、PyTorch 2.0
先说一个小内存的部署 然后在说一下刚出的pytorch2.0 100%向后兼容，一行代码将训练提速76%
前段时间，PyTorch 团队在官方博客宣布 Pytorch 1.13 发布，包含 BetterTransformer 稳定版等多项更新。在体验新特性的同时，不少人也在期待下一个版本的推出。
出乎意料的是，这个新版本很快就来了，而且是跨越式的 2.0 版！
新版本的重要进步体现在速度和可用性，而且完全向后兼容。PyTorch 团队表示，PyTorch 2.0 是他们向 2.x 系列迈出的第一步，其稳定版预计在 2023 年 3 月初发布。
首先，PyTorch 2.0 引入了 torch.compile，这是一种编译模式，可以在不更改模型代码的情况下加速模型。在 163 个涵盖视觉、NLP 和其他领域的开源模型中，该团队发现使用 2.0 可以将训练速度提高 38-76%。
其次，PyTorch 2.0 是 100% 向后兼容的：代码库一样，API 一样，写模型的方式也一样。团队之所以称它为 2.0，是因为它有一些标志性的新特性，包括：
- TorchDynamo 可以从字节码分析生成 FX 图；
- AOTAutograd 可以以 ahead-of-time 的方式生成反向图；
- PrimTorch 引入了一个小型算子集，使后端更容易；
- TorchInductor：一个由 OpenAI Triton 支持的 DL 编译器。
PyTorch 2.0 将延续 PyTorch 一贯的优势，包括 Python 集成、命令式风格、API 简单等等。此外，PyTorch 2.0 提供了相同的 eager-mode 开发和用户体验，同时从根本上改变和增强了 PyTorch 在编译器级别的运行方式。该版本能够为「Dynamic Shapes」和分布式运行提供更快的性能和更好的支持。
在官方博客中，PyTorch团队还公布了他们对于整个2.0系列的展望：
以下是详细内容。
PyTorch 2.X：速度更快、更加地 Python 化、一如既往地 dynamic
PyTorch 2.0 官宣了一个重要特性——torch.compile，这一特性将 PyTorch 的性能推向了新的高度，并将 PyTorch 的部分内容从 C++ 移回 Python。torch.compile 是一个完全附加的（可选的）特性，因此 PyTorch 2.0 是 100% 向后兼容的。
支撑 torch.compile 的技术包括研发团队新推出的 TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor。
- TorchDynamo 使用 Python Frame Evaluation Hooks 安全地捕获 PyTorch 程序，这是一项重大创新，是研究团队对快速可靠地获取图进行 5 年研发的结果；
- AOTAutograd 重载 PyTorch 的 autograd 引擎作为一个跟踪 autodiff，用于生成 ahead-of-time 向后跟踪；
- PrimTorch 将约 2000 多个 PyTorch 算子规范化为一组约 250 个原始算子的闭集，开发人员可以将其作为构建完整 PyTorch 后端的目标。这大大降低了编写 PyTorch 特性或后端的障碍；
- TorchInductor 是一种深度学习编译器，可为多个加速器和后端生成快速代码。对于 NVIDIA GPU，它使用 OpenAI Triton 作为关键构建块。
TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor 是用 Python 编写的，并支持 dynamic shapes（即能够发送不同大小的张量而无需重新编译），这使得它们具备灵活、易于破解的特性，降低了开发人员和供应商的使用门槛。
为了验证这些技术，研发团队在各种机器学习领域测试了 163 个开源模型。实验精心构建了测试基准，包括各种 CV 任务（图像分类、目标检测、图像生成等）、NLP 任务（语言建模、问答、序列分类、推荐系统等）和强化学习任务，测试模型主要有 3 个来源：
- 46 个来自 HuggingFace Transformers 的模型；
- 来自 TIMM 的 61 个模型：一系列 SOTA PyTorch 图像模型；
- 来自 TorchBench 的 56 个模型：包含来自 github 的精选流行代码库。
然后研究者测量加速性能并验证这些模型的准确性。加速可能取决于数据类型，研究团队选择测量 float32 和自动混合精度 (AMP) 的加速。
在 163 个开源模型中，torch.compile 在 93% 的情况下都有效，模型在 NVIDIA A100 GPU 上的训练速度提高了 43%。在 float32 精度下，它的平均运行速度提高了 21%，而在 AMP 精度下，它的运行速度平均提高了 51%。
目前，torch.compile 还处于早期开发阶段，预计 2023 年 3 月上旬将发布第一个稳定的 2.0 版本。
TorchDynamo：快速可靠地获取图
TorchDynamo 是一种使用 Frame Evaluation API （PEP-0523 中引入的一种 CPython 特性）的新方法。研发团队采用数据驱动的方法来验证其在 Graph Capture 上的有效性，并使用 7000 多个用 PyTorch 编写的 Github 项目作为验证集。TorchScript 等方法大约在 50% 的时间里都难以获取图，而且通常开销很大；而 TorchDynamo 在 99% 的时间里都能获取图，方法正确、安全且开销可忽略不计（无需对原始代码进行任何更改）。这说明 TorchDynamo 突破了多年来模型权衡灵活性和速度的瓶颈。
TorchInductor：使用 define-by-run IR 快速生成代码
对于 PyTorch 2.0 的新编译器后端，研发团队从用户编写高性能自定义内核的方式中汲取灵感：越来越多地使用 Triton 语言。此外，研究者还想要一个编译器后端——使用与 PyTorch eager 类似的抽象，并且具有足够的通用性以支持 PyTorch 中广泛的功能。
TorchInductor 使用 pythonic define-by-run loop level IR 自动将 PyTorch 模型映射到 GPU 上生成的 Triton 代码和 CPU 上的 C++/OpenMP。TorchInductor 的 core loop level IR 仅包含约 50 个算子，并且是用 Python 实现的，易于破解和扩展。
AOTAutograd：将 Autograd 重用于 ahead-of-time 图
PyTorch 2.0 的主要特性之一是加速训练，因此 PyTorch 2.0 不仅要捕获用户级代码，还要捕获反向传播。此外，研发团队还想要复用现有的经过实践检验的 PyTorch autograd 系统。AOTAutograd 利用 PyTorch 的 torch_dispatch 可扩展机制来跟踪 Autograd 引擎，使其能够「ahead-of-time」捕获反向传递（backwards pass）。这使 TorchInductor 能够加速前向和反向传递。
PrimTorch：稳定的原始算子
为 PyTorch 编写后端具有挑战性。PyTorch 有 1200 多个算子，如果考虑每个算子的各种重载，则有 2000 多个。
在 PrimTorch 项目中，研发团队致力于定义更小且稳定的算子集，将 PyTorch 程序缩减到这样较小的算子集。目标是定义两个算子集：
- Prim ops：约有 250 个相当低级的算子。这些算子适用于编译器，需要将它们重新融合在一起以获得良好的性能；
- ATen ops：约有 750 个规范算子。这些算子适用于已经在 ATen 级别集成的后端或没有编译功能的后端（无法从较低级别的算子集（如 Prim ops）恢复性能）。
用户体验
PyTorch 2.0 引入了一个简单的函数 torch.compile，它会返回一个编译后的模型。
compiled_model *=* torch.compile(model)compiled_model 保存对模型的引用，并将 forward 函数编译为一个更优化的版本。在编译模型时，PyTorch 2.0 给了几项设置来调整它：
def torch.compile(model: Callable,
*,
mode: Optional[str] = "default",
dynamic: bool = False,
fullgraph:bool = False,
backend: Union[str, Callable] = "inductor",
# advanced backend options go here as kwargs
**kwargs
) -> torch._dynamo.NNOptimizedModule- 「mode」指定编译器在编译时应该优化的内容。
- default 是一种预设模式，它试图在不花费太长时间或使用额外内存的情况下高效编译。
- 其他模式，如 reduce-overhead，可以大大降低框架开销，但要消耗少量额外内存。max-autotune 编译很长时间，试图为你提供它所能生成的最快的代码。
- 「dynamic」模式指定是否为 Dynamic Shapes 启用代码路径。某些编译器优化不能应用于动态形状的程序。明确你想要一个带有动态形状还是静态形状的编译程序，将有助于编译器提供更好的优化代码。
- 「fullgraph」类似于 Numba 的 nopython。它将整个程序编译成一个图，或者给出一个错误提示，解释为什么它不能这样做。大多数用户不需要使用这种模式。如果你非常注重性能，那么你可以尝试使用它。
- 「backend 」指定使用哪个编译器后端。默认情况下使用 TorchInductor，但还有其他一些可用的工具。
编译体验想要在默认模式中提供最大的好处和最大的灵活性。
常见问答
1、什么是 PT 2.0？
2.0 是最新的 PyTorch 版本。PyTorch 2.0 提供相同的 eager 模式开发体验，同时通过 torch.compile 添加编译模式。这种编译模式有可能在训练和推理期间加速模型。
2、为什么是 2.0 而不是 1.14？
PyTorch 2.0 就是计划中的 1.14 。我们发布了大量新特性，我们相信这些新特性会改变用户使用 PyTorch 的方式，因此称其为 PyTorch 2.0。
3、如何安装 2.0？有什么额外要求吗？
安装最新的 nightlies：
CUDA 11.7
pip3 install numpy --pre torch[dynamo] torchvision torchaudio --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu117CUDA 11.6
pip3 install numpy --pre torch[dynamo] torchvision torchaudio --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu116CPU
pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cpu4、2.0 代码是否向后兼容 1.X？
是的，使用 2.0 不需要修改 PyTorch 工作流程。一行代码 model = torch.compile(model) 就可以优化模型，以使用 2.0 堆栈，并与其余 PyTorch 代码一起顺利运行。这是一个可选择项，不需要使用新的编译器。
5、默认启用 2.0 吗？
不，必须通过使用单个函数调用优化模型从而在 PyTorch 代码中显式启用 2.0。
6、如何将 PT1.X 代码迁移到 PT2.0？
代码应该按原样工作，无需任何迁移。如果想使用 2.0 中引入的新编译模式特性，那么可以从优化模型开始：
model = torch.compile(model)虽然加速主要是在训练期间观察到的，但如果你的模型运行速度比 eager 模式快，也可以将它用于推理。
import torch
def train(model, dataloader):
model = torch.compile(model)
for batch in dataloader:
run_epoch(model, batch)
def infer(model, input):
model = torch.compile(model)
return model(\*\*input)7、是否有任何不该使用 PT 2.0 的应用程序？
当前版本的 PT 2.0 仍处于实验阶段，并且处于 nightlies 版本。其中的动态形状支持还处于早期阶段，所以可以等到 2023 年 3 月稳定版发布后再使用该功能。
8、运行 PyTorch 2.0 时，代码有何不同？
开箱即用，PyTorch 2.0 与 PyTorch 1.x 相同，模型以 eager 模式运行，即 Python 的每一行都逐个执行。
在 2.0 中，如果用 model = torch.compile(model) 将模型打包，则模型在执行之前会经过 3 个步骤：
1）图获取：首先将模型重写为子图块。可由 TorchDynamo 编译的子图被「压平」，其他子图（可能包含控制流代码或其他不受支持的 Python 结构）将回退到 Eager 模式。
2）Graph lowering：所有 PyTorch 操作都被分解为特定于所选后端的组成内核。
3）图编译，内核调用其相应的低级设备专用操作。
9、2.0 目前支持哪些编译器后端？
默认和最完整的后端是 TorchInductor，但是 TorchDynamo 有一个不断增长的后端列表，可以通过调用 torchdynamo.list_backends(). 找到
10、2.0 版本的分布式训练能力如何？
Compiled 模式下的 DDP 和 FSDP ，比 FP32 中的 Eager 模式快 15%、AMP 精度快 80%。PT2.0 做了一些额外的优化，以确保 DDP 的通信 - 计算 overlap 与 Dynamo 的部分图创建良好协作。想要确保使用 static_graph=False 运行 DDP，更多细节参见：https://dev-discuss.pytorch.org/t/torchdynamo-update-9-making-ddp-work-with-torchdynamo/860
11、为什么我的代码用 2.0 的 Compiled Model 运行变慢？
性能下降最可能的原因是 graph break 太多。例如，类似模型前向 trigger 中的输出语句这样的东西会触发 graph break。详见：https://pytorch.org/docs/master/dynamo/faq.html#why-am-i-not-seeing-speedups
12、以前运行的代码在 2.0 中崩溃了，该如何调试？
参见：https ://pytorch.org/docs/master/dynamo/faq.html#why-is-my-code-crashing
参考链接：https://pytorch.org/get-started/pytorch-2.0/
四、Pytorch Debug
在使用Pytorch时你或多或少会遇到各种bug
CrossEntropyLoss和NLLLoss
最常见的错误是损失函数和输出激活函数之间的不匹配。nn.CrossEntropyLossPyTorch中的损失模块执行两个操作：nn.LogSoftmax和nn.NLLLoss。
因此nn.CrossEntropyLossPyTorch的输入应该是最后一个线性层的输出。不要在nn.CrossEntropyLossPyTorch之前应用Softmax。 否则将对Softmax输出计算log-softmax，将会降低模型精度。
如果使用nn.NLLLoss模块，则需要自己应用log-softmax。nn.NLLLoss需要对数概率，而不是普通概率。因此确保应用nn.LogSoftmaxor nn.functional.log_softmax，而不是nn.Softmax。
Softmax的计算维度
注意Softmax的计算维度。通常是输出张量的最后一个维度，例如nn.Softmax(dim=-1)。如果混淆了维度，模型最终会得到随机预测。
类别数据与嵌入操作
对于类别数据，常见的做法是进行数值编码。但对于深度学习而言，这并不是一个很好的操作，数值会带来大小关系，且会丢失很多信息。因此对于类别数据建议使用one-hot或Embedding操作，对于nn.Embedding模块，你需要设置的参数包括：
num_embeddings：数据类别的数量embedding_dim：每个类别的嵌入维度padding_idx：填充符号的索引
嵌入特征向量从随机初始化，不要用 Kaiming、Xavier初始化方法。因为标准差为1，初始化、激活函数等被设计为输入标准差为 1。nn.Embedding模块的示例用法：
import torch
import torch.nn as nn
# Create 5 embedding vectors each with 32 features
embedding = nn.Embedding(num_embeddings=5,
embedding_dim=32)
# Example integer input
input_tensor = torch.LongTensor([[0, 4], [2, 3], [0, 1]])
# Get embeddings
embed_vectors = embedding(input_tensor)
print("Input shape:", input_tensor.shape)
print("Output shape:", embed_vectors.shape)
print("Example features:\n", embed_vectors[:,:,:2])nn.LSTM 中 数据维度
默认情况下，PyTorch的nn.LSTM模块假定输入维度为[seq_len, batch_size, input_size]，所以确保不要混淆序列长度和批大小的次数。如果混淆LSTM仍然可以正常运行，但会给出错误的结果。
维度不匹配
如果Pytorch执行矩阵乘法，并两个矩阵出现维度不匹配，PyTorch会报错并抛出错误。但是也存在PyTorch不会抛出错误的情况，此时未对齐的维度具有相同的大小。建议使用多个不同的批量大小测试您的代码，以防止维度不对齐。
训练和评估模式
在PyTorch中，神经网络有两种模式：train和train。您可以使用model.eval（)和model.train()对模型时进行切换。不同的模式决定是否使用dropout，以及如何处理Batch Normalization。常见的错误是在eval后忘记将模型设置回train模式，确定模型在预测阶段为eval模式。
参数继承
PyTorch支持nn.Modules，一个模块可以包含另一个模块，另一个模块又可以包含一个模块，依此类推。
当调用.parameters()时，PyTorch会查找该模块内的所有模块，并将它们的参数添加到最高级别模块的参数中。
但是PyTorch不会检测列表、字典或类似结构中模块的参数。如果有一个模块列表，请确保将它们放入一个nn.ModuleList或nn.Sequential对象中。
参数初始化
正确初始化模型的参数非常重要。用标准正态分布初始化参数不是好的选择，推荐的方法有Kaiming或Xavier。
zero_grad()
请记住在执行loss.backward()之前调用optimizer.zero_grad()。如果在执行反向传播之前没有重置所有参数的梯度，梯度将被添加到上一批的梯度中。
指标计算逻辑
在怀疑自己或模型之前，请经常检查您的指标计算逻辑计算两次或更多次。像准确性这样的指标很容易计算，但在代码中添加错误也很容易。例如，检查您是否对批次维度进行了平均，而不是意外对类维度或任何其他维度进行平均。
设备不匹配
如果使用GPU可能会看到一个错误，例如：
Runtime Error: Input type (torch.FloatTensor) dand weigh type (torch.cuda.FloatTensor) should be on the same device.
此错误表示输入数据在CPU上，而权重在GPU上。确保所有数据都在同一设备上。这通常是GPU，因为它支持训练和测试加速。
nn.Sequential和nn.ModuleList
如果模型有很多层，推荐将它们汇总为一个nn.Sequential或nn.ModuleList对象。在前向传递中，只需要调用sequential，或者遍历模块列表。
class MLP(nn.Module):
def __init__(self, input_dims=64, hidden_dims=[128,256], output_dims=10):
super().__init__()
hidden_dims = [input_dims] + hidden_dims
layers = []
for idx in range(len(hidden_dims)-1):
layers += [
nn.Linear(hidden_dims[i], hidden_dims[i+1]),
nn.ReLU(inplace=True)
]
self.layers = nn.Sequential(*layers)
def forward(self, x):
return self.layers(x)参数重复计算
在深度神经网络中，通常会有重复添加到模型中的块。如果这些块需要比更复杂的前向函数，建议在单独的模块中实现它们。例如，一个 ResNet 由多个具有残差连接的ResNet块组成。ResNet模块应用一个小型神经网络，并将输出添加回输入。最好在单独的类中实现这种动态，以保持主模型类小而清晰。
输入相同的维度
如果您有多个具有相同输入的线性层或卷积，则可以将它们堆叠在一起以提高效率。假设我们有：
虽然可以通过两个线性层来实现它，但您可以通过将两层堆叠为一层来获得完全相同的神经网络。单层效率更高，因为这代表单个矩阵运算，而不是GPU的两个矩阵运算，因此我们可以并行化计算。
x = torch.randn(2, 10)
# Implementation of separate layers:
y1_layer = nn.Linear(10, 20)
y2_layer = nn.Linear(10, 30)
y1 = y1_layer(x)
y2 = y2_layer(x)
# Implementation of a stacked layer:
y_layer = nn.Linear(10, 50)
y = y_layer(x)
y1, y2 = y[:,:20], y[:,20:50]使用带logits的损失函数
分类损失函数（例如二元交叉熵）在PyTorch中有两个版本：nn.BCELoss和nn.BCEWithLogitsLoss，建议和推荐的做法是使用后者。这因为它在数值上更稳定，并在您的模型预测非常错误时防止出现任何不稳定性。
如果您不使用logit损失函数，则当模型预测不正确的非常高或非常低的值时，您可能会遇到问题。
五、 常用代码段
PyTorch常用代码段合集，涵盖基本配置、张量处理、模型定义与操作、数据处理、模型训练与测试等5个方面，还给出了多个值得注意的Tips，内容非常全面。
PyTorch最好的资料是官方文档。本文是PyTorch常用代码段，在参考资料[1](张皓：PyTorch Cookbook)的基础上做了一些修补，方便使用时查阅。
基本配置
导入包和版本查询
import torch
import torch.nn as nn
import torchvision
print(torch.__version__)
print(torch.version.cuda)
print(torch.backends.cudnn.version())
print(torch.cuda.get_device_name(0))可复现性
在硬件设备（CPU、GPU）不同时，完全的可复现性无法保证，即使随机种子相同。但是，在同一个设备上，应该保证可复现性。具体做法是，在程序开始的时候固定torch的随机种子，同时也把numpy的随机种子固定。
np.random.seed(0)
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False显卡设置
如果只需要一张显卡
# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')如果需要指定多张显卡，比如0，1号显卡。
import osos.environ['CUDA_VISIBLE_DEVICES'] = '0,1'也可以在命令行运行代码时设置显卡：
CUDA_VISIBLE_DEVICES=0,1 python train.py清除显存
torch.cuda.empty_cache()也可以使用在命令行重置GPU的指令
nvidia-smi --gpu-reset -i [gpu_id]张量(Tensor)处理
张量的数据类型
PyTorch有9种CPU张量类型和9种GPU张量类型。
张量基本信息
tensor = torch.randn(3,4,5)print(tensor.type()) # 数据类型print(tensor.size()) # 张量的shape，是个元组print(tensor.dim()) # 维度的数量命名张量
张量命名是一个非常有用的方法，这样可以方便地使用维度的名字来做索引或其他操作，大大提高了可读性、易用性，防止出错。
# 在PyTorch 1.3之前，需要使用注释
# Tensor[N, C, H, W]
images = torch.randn(32, 3, 56, 56)
images.sum(dim=1)
images.select(dim=1, index=0)
# PyTorch 1.3之后
NCHW = [‘N’, ‘C’, ‘H’, ‘W’]
images = torch.randn(32, 3, 56, 56, names=NCHW)
images.sum('C')
images.select('C', index=0)
# 也可以这么设置
tensor = torch.rand(3,4,1,2,names=('C', 'N', 'H', 'W'))
# 使用align_to可以对维度方便地排序
tensor = tensor.align_to('N', 'C', 'H', 'W')数据类型转换
# 设置默认类型，pytorch中的FloatTensor远远快于DoubleTensor
torch.set_default_tensor_type(torch.FloatTensor)
# 类型转换
tensor = tensor.cuda()
tensor = tensor.cpu()
tensor = tensor.float()
tensor = tensor.long()torch.Tensor与np.ndarray转换
除了CharTensor，其他所有CPU上的张量都支持转换为numpy格式然后再转换回来。
ndarray = tensor.cpu().numpy()
tensor = torch.from_numpy(ndarray).float()
tensor = torch.from_numpy(ndarray.copy()).float() # If ndarray has negative stride.Torch.tensor与PIL.Image转换
# pytorch中的张量默认采用[N, C, H, W]的顺序，并且数据范围在[0,1]，需要进行转置和规范化
# torch.Tensor -> PIL.Image
image = PIL.Image.fromarray(torch.clamp(tensor*255, min=0, max=255).byte().permute(1,2,0).cpu().numpy())
image = torchvision.transforms.functional.to_pil_image(tensor) # Equivalently way
# PIL.Image -> torch.Tensor
path = r'./figure.jpg'
tensor = torch.from_numpy(np.asarray(PIL.Image.open(path))).permute(2,0,1).float() / 255
tensor = torchvision.transforms.functional.to_tensor(PIL.Image.open(path)) # Equivalently waynp.ndarray与PIL.Image的转换
image = PIL.Image.fromarray(ndarray.astype(np.uint8))
ndarray = np.asarray(PIL.Image.open(path))从只包含一个元素的张量中提取值
value = torch.rand(1).item()张量形变
# 在将卷积层输入全连接层的情况下通常需要对张量做形变处理，
# 相比torch.view，torch.reshape可以自动处理输入张量不连续的情况
tensor = torch.rand(2,3,4)
shape = (6, 4)
tensor = torch.reshape(tensor, shape)打乱顺序
tensor = tensor[torch.randperm(tensor.size(0))] # 打乱第一个维度水平翻转
# pytorch不支持tensor[::-1]这样的负步长操作，水平翻转可以通过张量索引实现
# 假设张量的维度为[N, D, H, W].
tensor = tensor[:,:,:,torch.arange(tensor.size(3) - 1, -1, -1).long()]复制张量
# Operation | New/Shared memory | Still in computation graph |tensor.clone() # | New | Yes |tensor.detach() # | Shared | No |tensor.detach.clone()() # | New | No |张量拼接
'''
注意torch.cat和torch.stack的区别在于torch.cat沿着给定的维度拼接，
而torch.stack会新增一维。例如当参数是3个10x5的张量，torch.cat的结果是30x5的张量，
而torch.stack的结果是3x10x5的张量。
'''
tensor = torch.cat(list_of_tensors, dim=0)
tensor = torch.stack(list_of_tensors, dim=0)将整数标签转为one-hot编码
# pytorch的标记默认从0开始
tensor = torch.tensor([0, 2, 1, 3])
N = tensor.size(0)
num_classes = 4
one_hot = torch.zeros(N, num_classes).long()
one_hot.scatter_(dim=1, index=torch.unsqueeze(tensor, dim=1), src=torch.ones(N, num_classes).long())得到非零元素
torch.nonzero(tensor) # index of non-zero elements
torch.nonzero(tensor==0) # index of zero elements
torch.nonzero(tensor).size(0) # number of non-zero elements
torch.nonzero(tensor == 0).size(0) # number of zero elements判断两个张量相等
torch.allclose(tensor1, tensor2) # float tensor
torch.equal(tensor1, tensor2) # int tensor张量扩展
# Expand tensor of shape 64*512 to shape 64*512*7*7.
tensor = torch.rand(64,512)
torch.reshape(tensor, (64, 512, 1, 1)).expand(64, 512, 7, 7)矩阵乘法
# Matrix multiplcation: (m*n) * (n*p) * -> (m*p).
result = torch.mm(tensor1, tensor2)
# Batch matrix multiplication: (b*m*n) * (b*n*p) -> (b*m*p)
result = torch.bmm(tensor1, tensor2)
# Element-wise multiplication.
result = tensor1 * tensor2计算两组数据之间的两两欧式距离
利用广播机制
dist = torch.sqrt(torch.sum((X1[:,None,:] - X2) ** 2, dim=2))模型定义和操作
一个简单两层卷积网络的示例
# convolutional neural network (2 convolutional layers)
class ConvNet(nn.Module):
def __init__(self, num_classes=10):
super(ConvNet, self).__init__()
self.layer1 = nn.Sequential(
nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
nn.BatchNorm2d(16),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2))
self.layer2 = nn.Sequential(
nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2))
self.fc = nn.Linear(7*7*32, num_classes)
def forward(self, x):
out = self.layer1(x)
out = self.layer2(out)
out = out.reshape(out.size(0), -1)
out = self.fc(out)
return out
model = ConvNet(num_classes).to(device)卷积层的计算和展示可以用这个网站辅助。
双线性汇合（bilinear pooling）
X = torch.reshape(N, D, H * W) # Assume X has shape N*D*H*W
X = torch.bmm(X, torch.transpose(X, 1, 2)) / (H * W) # Bilinear pooling
assert X.size() == (N, D, D)
X = torch.reshape(X, (N, D * D))
X = torch.sign(X) * torch.sqrt(torch.abs(X) + 1e-5) # Signed-sqrt normalization
X = torch.nn.functional.normalize(X) # L2 normalization多卡同步 BN（Batch normalization）
当使用 torch.nn.DataParallel 将代码运行在多张 GPU 卡上时，PyTorch 的 BN 层默认操作是各卡上数据独立地计算均值和标准差，同步 BN 使用所有卡上的数据一起计算 BN 层的均值和标准差，缓解了当批量大小（batch size）比较小时对均值和标准差估计不准的情况，是在目标检测等任务中一个有效的提升性能的技巧。
sync_bn = torch.nn.SyncBatchNorm(num_features,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=True)将已有网络的所有BN层改为同步BN层
def convertBNtoSyncBN(module, process_group=None):
'''Recursively replace all BN layers to SyncBN layer.
Args:
module[torch.nn.Module]. Network
'''
if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
sync_bn = torch.nn.SyncBatchNorm(module.num_features, module.eps, module.momentum,
module.affine, module.track_running_stats, process_group)
sync_bn.running_mean = module.running_mean
sync_bn.running_var = module.running_var
if module.affine:
sync_bn.weight = module.weight.clone().detach()
sync_bn.bias = module.bias.clone().detach()
return sync_bn
else:
for name, child_module in module.named_children():
setattr(module, name) = convert_syncbn_model(child_module, process_group=process_group))
return module类似 BN 滑动平均
如果要实现类似 BN 滑动平均的操作，在 forward 函数中要使用原地（inplace）操作给滑动平均赋值。
class BN(torch.nn.Module)
def __init__(self):
...
self.register_buffer('running_mean', torch.zeros(num_features))
def forward(self, X):
...
self.running_mean += momentum * (current - self.running_mean)计算模型整体参数量
num_parameters = sum(torch.numel(parameter) for parameter in model.parameters())查看网络中的参数
可以通过model.state_dict()或者model.named_parameters()函数查看现在的全部可训练参数（包括通过继承得到的父类中的参数）
params = list(model.named_parameters())
(name, param) = params[28]
print(name)
print(param.grad)
print('-------------------------------------------------')
(name2, param2) = params[29]
print(name2)
print(param2.grad)
print('----------------------------------------------------')
(name1, param1) = params[30]
print(name1)
print(param1.grad)模型可视化（使用pytorchviz）
szagoruyko/pytorchvizgithub.com类似 Keras 的 model.summary() 输出模型信息，使用pytorch-summary
sksq96/pytorch-summarygithub.com模型权重初始化
注意 model.modules() 和 model.children() 的区别：model.modules() 会迭代地遍历模型的所有子层，而 model.children() 只会遍历模型下的一层。
# Common practise for initialization.
for layer in model.modules():
if isinstance(layer, torch.nn.Conv2d):
torch.nn.init.kaiming_normal_(layer.weight, mode='fan_out',
nonlinearity='relu')
if layer.bias is not None:
torch.nn.init.constant_(layer.bias, val=0.0)
elif isinstance(layer, torch.nn.BatchNorm2d):
torch.nn.init.constant_(layer.weight, val=1.0)
torch.nn.init.constant_(layer.bias, val=0.0)
elif isinstance(layer, torch.nn.Linear):
torch.nn.init.xavier_normal_(layer.weight)
if layer.bias is not None:
torch.nn.init.constant_(layer.bias, val=0.0)
# Initialization with given tensor.
layer.weight = torch.nn.Parameter(tensor)提取模型中的某一层
modules()会返回模型中所有模块的迭代器，它能够访问到最内层，比如self.layer1.conv1这个模块，还有一个与它们相对应的是name_children()属性以及named_modules(),这两个不仅会返回模块的迭代器，还会返回网络层的名字。
# 取模型中的前两层
new_model = nn.Sequential(*list(model.children())[:2]
# 如果希望提取出模型中的所有卷积层，可以像下面这样操作：
for layer in model.named_modules():
if isinstance(layer[1],nn.Conv2d):
conv_model.add_module(layer[0],layer[1])部分层使用预训练模型
注意如果保存的模型是 torch.nn.DataParallel，则当前的模型也需要是
model.load_state_dict(torch.load('model.pth'), strict=False)将在 GPU 保存的模型加载到 CPU
model.load_state_dict(torch.load('model.pth', map_location='cpu'))导入另一个模型的相同部分到新的模型
模型导入参数时，如果两个模型结构不一致，则直接导入参数会报错。用下面方法可以把另一个模型的相同的部分导入到新的模型中。
# model_new代表新的模型
# model_saved代表其他模型，比如用torch.load导入的已保存的模型
model_new_dict = model_new.state_dict()
model_common_dict = {k:v for k, v in model_saved.items() if k in model_new_dict.keys()}
model_new_dict.update(model_common_dict)
model_new.load_state_dict(model_new_dict)数据处理
计算数据集的均值和标准差
import os
import cv2
import numpy as np
from torch.utils.data import Dataset
from PIL import Image
def compute_mean_and_std(dataset):
# 输入PyTorch的dataset，输出均值和标准差
mean_r = 0
mean_g = 0
mean_b = 0
for img, _ in dataset:
img = np.asarray(img) # change PIL Image to numpy array
mean_b += np.mean(img[:, :, 0])
mean_g += np.mean(img[:, :, 1])
mean_r += np.mean(img[:, :, 2])
mean_b /= len(dataset)
mean_g /= len(dataset)
mean_r /= len(dataset)
diff_r = 0
diff_g = 0
diff_b = 0
N = 0
for img, _ in dataset:
img = np.asarray(img)
diff_b += np.sum(np.power(img[:, :, 0] - mean_b, 2))
diff_g += np.sum(np.power(img[:, :, 1] - mean_g, 2))
diff_r += np.sum(np.power(img[:, :, 2] - mean_r, 2))
N += np.prod(img[:, :, 0].shape)
std_b = np.sqrt(diff_b / N)
std_g = np.sqrt(diff_g / N)
std_r = np.sqrt(diff_r / N)
mean = (mean_b.item() / 255.0, mean_g.item() / 255.0, mean_r.item() / 255.0)
std = (std_b.item() / 255.0, std_g.item() / 255.0, std_r.item() / 255.0)
return mean, std得到视频数据基本信息
import cv2
video = cv2.VideoCapture(mp4_path)
height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
num_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
fps = int(video.get(cv2.CAP_PROP_FPS))
video.release()TSN 每段（segment）采样一帧视频
K = self._num_segments
if is_train:
if num_frames > K:
# Random index for each segment.
frame_indices = torch.randint(
high=num_frames // K, size=(K,), dtype=torch.long)
frame_indices += num_frames // K * torch.arange(K)
else:
frame_indices = torch.randint(
high=num_frames, size=(K - num_frames,), dtype=torch.long)
frame_indices = torch.sort(torch.cat((
torch.arange(num_frames), frame_indices)))[0]
else:
if num_frames > K:
# Middle index for each segment.
frame_indices = num_frames / K // 2
frame_indices += num_frames // K * torch.arange(K)
else:
frame_indices = torch.sort(torch.cat((
torch.arange(num_frames), torch.arange(K - num_frames))))[0]
assert frame_indices.size() == (K,)
return [frame_indices[i] for i in range(K)]常用训练和验证数据预处理
其中 ToTensor 操作会将 PIL.Image 或形状为 H×W×D，数值范围为 [0, 255] 的 np.ndarray 转换为形状为 D×H×W，数值范围为 [0.0, 1.0] 的 torch.Tensor。
train_transform = torchvision.transforms.Compose([
torchvision.transforms.RandomResizedCrop(size=224,
scale=(0.08, 1.0)),
torchvision.transforms.RandomHorizontalFlip(),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225)),
])
val_transform = torchvision.transforms.Compose([
torchvision.transforms.Resize(256),
torchvision.transforms.CenterCrop(224),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225)),
])模型训练和测试
分类模型训练代码
# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
# Train the model
total_step = len(train_loader)
for epoch in range(num_epochs):
for i ,(images, labels) in enumerate(train_loader):
images = images.to(device)
labels = labels.to(device)
# Forward pass
outputs = model(images)
loss = criterion(outputs, labels)
# Backward and optimizer
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (i+1) % 100 == 0:
print('Epoch: [{}/{}], Step: [{}/{}], Loss: {}'
.format(epoch+1, num_epochs, i+1, total_step, loss.item()))分类模型测试代码
# Test the model
model.eval（) # eval mode(batch norm uses moving mean/variance
#instead of mini-batch mean/variance)
with torch.no_grad():
correct = 0
total = 0
for images, labels in test_loader:
images = images.to(device)
labels = labels.to(device)
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print('Test accuracy of the model on the 10000 test images: {} %'
.format(100 * correct / total))自定义loss
继承torch.nn.Module类写自己的loss。
class MyLoss(torch.nn.Moudle):
def __init__(self):
super(MyLoss, self).__init__()
def forward(self, x, y):
loss = torch.mean((x - y) ** 2)
return loss标签平滑（label smoothing）
写一个label_smoothing.py的文件，然后在训练代码里引用，用LSR代替交叉熵损失即可。label_smoothing.py内容如下：
import torch
import torch.nn as nn
class LSR(nn.Module):
def __init__(self, e=0.1, reduction='mean'):
super().__init__()
self.log_softmax = nn.LogSoftmax(dim=1)
self.e = e
self.reduction = reduction
def _one_hot(self, labels, classes, value=1):
"""
Convert labels to one hot vectors
Args:
labels: torch tensor in format [label1, label2, label3, ...]
classes: int, number of classes
value: label value in one hot vector, default to 1
Returns:
return one hot format labels in shape [batchsize, classes]
"""
one_hot = torch.zeros(labels.size(0), classes)
#labels and value_added size must match
labels = labels.view(labels.size(0), -1)
value_added = torch.Tensor(labels.size(0), 1).fill_(value)
value_added = value_added.to(labels.device)
one_hot = one_hot.to(labels.device)
one_hot.scatter_add_(1, labels, value_added)
return one_hot
def _smooth_label(self, target, length, smooth_factor):
"""convert targets to one-hot format, and smooth
them.
Args:
target: target in form with [label1, label2, label_batchsize]
length: length of one-hot format(number of classes)
smooth_factor: smooth factor for label smooth
Returns:
smoothed labels in one hot format
"""
one_hot = self._one_hot(target, length, value=1 - smooth_factor)
one_hot += smooth_factor / (length - 1)
return one_hot.to(target.device)
def forward(self, x, target):
if x.size(0) != target.size(0):
raise ValueError('Expected input batchsize ({}) to match target batch_size({})'
.format(x.size(0), target.size(0)))
if x.dim() < 2:
raise ValueError('Expected input tensor to have least 2 dimensions(got {})'
.format(x.size(0)))
if x.dim() != 2:
raise ValueError('Only 2 dimension tensor are implemented, (got {})'
.format(x.size()))
smoothed_target = self._smooth_label(target, x.size(1), self.e)
x = self.log_softmax(x)
loss = torch.sum(- x * smoothed_target, dim=1)
if self.reduction == 'none':
return loss
elif self.reduction == 'sum':
return torch.sum(loss)
elif self.reduction == 'mean':
return torch.mean(loss)
else:
raise ValueError('unrecognized option, expect reduction to be one of none, mean, sum')或者直接在训练文件里做label smoothing
for images, labels in train_loader:
images, labels = images.cuda(), labels.cuda()
N = labels.size(0)
# C is the number of classes.
smoothed_labels = torch.full(size=(N, C), fill_value=0.1 / (C - 1)).cuda()
smoothed_labels.scatter_(dim=1, index=torch.unsqueeze(labels, dim=1), value=0.9)
score = model(images)
log_prob = torch.nn.functional.log_softmax(score, dim=1)
loss = -torch.sum(log_prob * smoothed_labels) / N
optimizer.zero_grad()
loss.backward()
optimizer.step()Mixup训练
beta_distribution = torch.distributions.beta.Beta(alpha, alpha)
for images, labels in train_loader:
images, labels = images.cuda(), labels.cuda()
# Mixup images and labels.
lambda_ = beta_distribution.sample([]).item()
index = torch.randperm(images.size(0)).cuda()
mixed_images = lambda_ * images + (1 - lambda_) * images[index, :]
label_a, label_b = labels, labels[index]
# Mixup loss.
scores = model(mixed_images)
loss = (lambda_ * loss_function(scores, label_a)
+ (1 - lambda_) * loss_function(scores, label_b))
optimizer.zero_grad()
loss.backward()
optimizer.step()L1 正则化
l1_regularization = torch.nn.L1Loss(reduction='sum')
loss = ... # Standard cross-entropy loss
for param in model.parameters():
loss += torch.sum(torch.abs(param))
loss.backward()不对偏置项进行权重衰减（weight decay)
pytorch里的weight decay相当于l2正则
bias_list = (param for name, param in model.named_parameters() if name[-4:] == 'bias')
others_list = (param for name, param in model.named_parameters() if name[-4:] != 'bias')
parameters = [{'parameters': bias_list, 'weight_decay': 0},
{'parameters': others_list}]
optimizer = torch.optim.SGD(parameters, lr=1e-2, momentum=0.9, weight_decay=1e-4)梯度裁剪（gradient clipping）
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=20)得到当前学习率
# If there is one global learning rate (which is the common case).
lr = next(iter(optimizer.param_groups))['lr']
# If there are multiple learning rates for different layers.
all_lr = []
for param_group in optimizer.param_groups:
all_lr.append(param_group['lr'])另一种方法，在一个batch训练代码里，当前的lr是optimizer.param_groups[0]['lr']
学习率衰减
# Reduce learning rate when validation accuarcy plateau.
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', patience=5, verbose=True)
for t in range(0, 80):
train(...)
val(...)
scheduler.step(val_acc)
# Cosine annealing learning rate.
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=80)
# Reduce learning rate by 10 at given epochs.
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[50, 70], gamma=0.1)
for t in range(0, 80):
scheduler.step()
train(...)
val(...)
# Learning rate warmup by 10 epochs.
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda t: t / 10)
for t in range(0, 10):
scheduler.step()
train(...)
val(...)优化器链式更新
从1.4版本开始，torch.optim.lr_scheduler 支持链式更新（chaining），即用户可以定义两个 schedulers，并交替在训练中使用。
import torch
from torch.optim import SGD
from torch.optim.lr_scheduler import ExponentialLR, StepLR
model = [torch.nn.Parameter(torch.randn(2, 2, requires_grad=True))]
optimizer = SGD(model, 0.1)
scheduler1 = ExponentialLR(optimizer, gamma=0.9)
scheduler2 = StepLR(optimizer, step_size=3, gamma=0.1)
for epoch in range(4):
print(epoch, scheduler2.get_last_lr()[0])
optimizer.step()
scheduler1.step()
scheduler2.step()模型训练可视化
PyTorch可以使用tensorboard来可视化训练过程。安装和运行TensorBoard。
pip install tensorboard
tensorboard --logdir=runs使用SummaryWriter类来收集和可视化相应的数据，放了方便查看，可以使用不同的文件夹，比如'Loss/train'和'Loss/test'。
from torch.utils.tensorboard import SummaryWriter
import numpy as np
writer = SummaryWriter()
for n_iter in range(100):
writer.add_scalar('Loss/train', np.random.random(), n_iter)
writer.add_scalar('Loss/test', np.random.random(), n_iter)
writer.add_scalar('Accuracy/train', np.random.random(), n_iter)
writer.add_scalar('Accuracy/test', np.random.random(), n_iter)保存与加载断点
注意为了能够恢复训练，我们需要同时保存模型和优化器的状态，以及当前的训练轮数。
start_epoch = 0
# Load checkpoint.
if resume: # resume为参数，第一次训练时设为0，中断再训练时设为1
model_path = os.path.join('model', 'best_checkpoint.pth.tar')
assert os.path.isfile(model_path)
checkpoint = torch.load(model_path)
best_acc = checkpoint['best_acc']
start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
print('Load checkpoint at epoch {}.'.format(start_epoch))
print('Best accuracy so far {}.'.format(best_acc))
# Train the model
for epoch in range(start_epoch, num_epochs):
...
# Test the model
...
# save checkpoint
is_best = current_acc > best_acc
best_acc = max(current_acc, best_acc)
checkpoint = {
'best_acc': best_acc,
'epoch': epoch + 1,
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
}
model_path = os.path.join('model', 'checkpoint.pth.tar')
best_model_path = os.path.join('model', 'best_checkpoint.pth.tar')
torch.save(checkpoint, model_path)
if is_best:
shutil.copy(model_path, best_model_path)提取 ImageNet 预训练模型某层的卷积特征
# VGG-16 relu5-3 feature.
model = torchvision.models.vgg16(pretrained=True).features[:-1]
# VGG-16 pool5 feature.
model = torchvision.models.vgg16(pretrained=True).features
# VGG-16 fc7 feature.
model = torchvision.models.vgg16(pretrained=True)
model.classifier = torch.nn.Sequential(*list(model.classifier.children())[:-3])
# ResNet GAP feature.
model = torchvision.models.resnet18(pretrained=True)
model = torch.nn.Sequential(collections.OrderedDict(
list(model.named_children())[:-1]))
with torch.no_grad():
model.eval（)
conv_representation = model(image)提取 ImageNet 预训练模型多层的卷积特征
class FeatureExtractor(torch.nn.Module):
"""Helper class to extract several convolution features from the given
pre-trained model.
Attributes:
_model, torch.nn.Module.
_layers_to_extract, list<str> or set<str>
Example:
>>> model = torchvision.models.resnet152(pretrained=True)
>>> model = torch.nn.Sequential(collections.OrderedDict(
list(model.named_children())[:-1]))
>>> conv_representation = FeatureExtractor(
pretrained_model=model,
layers_to_extract={'layer1', 'layer2', 'layer3', 'layer4'})(image)
"""
def __init__(self, pretrained_model, layers_to_extract):
torch.nn.Module.__init__(self)
self._model = pretrained_model
self._model.eval（)
self._layers_to_extract = set(layers_to_extract)
def forward(self, x):
with torch.no_grad():
conv_representation = []
for name, layer in self._model.named_children():
x = layer(x)
if name in self._layers_to_extract:
conv_representation.append(x)
return conv_representation微调全连接层
model = torchvision.models.resnet18(pretrained=True)
for param in model.parameters():
param.requires_grad = False
model.fc = nn.Linear(512, 100) # Replace the last fc layer
optimizer = torch.optim.SGD(model.fc.parameters(), lr=1e-2, momentum=0.9, weight_decay=1e-4)以较大学习率微调全连接层，较小学习率微调卷积层
model = torchvision.models.resnet18(pretrained=True)
finetuned_parameters = list(map(id, model.fc.parameters()))
conv_parameters = (p for p in model.parameters() if id(p) not in finetuned_parameters)
parameters = [{'params': conv_parameters, 'lr': 1e-3},
{'params': model.fc.parameters()}]
optimizer = torch.optim.SGD(parameters, lr=1e-2, momentum=0.9, weight_decay=1e-4)其他注意事项
- 不要使用太大的线性层。因为nn.Linear(m,n)使用的是的内存，线性层太大很容易超出现有显存。
- 不要在太长的序列上使用RNN。因为RNN反向传播使用的是BPTT算法，其需要的内存和输入序列的长度呈线性关系。
- model(x) 前用 model.train() 和 model.eval（) 切换网络状态。
- 不需要计算梯度的代码块用 with torch.no_grad() 包含起来。
- model.eval（) 和 torch.no_grad() 的区别在于，model.eval（) 是将网络切换为测试状态，例如 BN 和dropout在训练和测试阶段使用不同的计算方法。torch.no_grad() 是关闭 PyTorch 张量的自动求导机制，以减少存储使用和加速计算，得到的结果无法进行 loss.backward()。
- model.zero_grad()会把整个模型的参数的梯度都归零, 而optimizer.zero_grad()只会把传入其中的参数的梯度归零.
- torch.nn.CrossEntropyLoss 的输入不需要经过 Softmax。torch.nn.CrossEntropyLoss 等价于 torch.nn.functional.log_softmax + torch.nn.NLLLoss。
- loss.backward() 前用 optimizer.zero_grad() 清除累积梯度。
- torch.utils.data.DataLoader 中尽量设置 pin_memory=True，对特别小的数据集如 MNIST 设置 pin_memory=False 反而更快一些。num_workers 的设置需要在实验中找到最快的取值。
- 用 del 及时删除不用的中间变量，节约 GPU 存储。使用 inplace 操作可节约 GPU 存储，如：
x = torch.nn.functional.relu(x, inplace=True)减少 CPU 和 GPU 之间的数据传输。例如如果你想知道一个 epoch 中每个 mini-batch 的 loss 和准确率，先将它们累积在 GPU 中等一个 epoch 结束之后一起传输回 CPU 会比每个 mini-batch 都进行一次 GPU 到 CPU 的传输更快。使用半精度浮点数 half() 会有一定的速度提升，具体效率依赖于 GPU 型号。需要小心数值精度过低带来的稳定性问题。时常使用 assert tensor.size() == (N, D, H, W) 作为调试手段，确保张量维度和你设想中一致。除了标记 y 外，尽量少使用一维张量，使用 n*1 的二维张量代替，可以避免一些意想不到的一维张量计算结果。统计代码各部分耗时：
with torch.autograd.profiler.profile(enabled=True, use_cuda=False) as profile:
...print(profile)# 或者在命令行运行python -m torch.utils.bottleneck main.py使用TorchSnooper来调试PyTorch代码，程序在执行的时候，就会自动 print 出来每一行的执行结果的 tensor 的形状、数据类型、设备、是否需要梯度的信息。
# pip install torchsnooper
import torchsnooper# 对于函数，使用修饰器@torchsnooper.snoop()
# 如果不是函数，使用 with 语句来激活 TorchSnooper，把训练的那个循环装进 with 语句中去。
with torchsnooper.snoop():
原本的代码https://github.com/zasdfgbnm/TorchSnoopergithub.com模型可解释性，使用captum库：https://captum.ai/captum.ai
参考资料
- 张皓：PyTorch Cookbook，https://zhuanlan.zhihu.com/p/59205847?
- PyTorch官方文档和示例
- https://pytorch.org/docs/stable/notes/faq.html
- https://github.com/szagoruyko/pytorchviz
- https://github.com/sksq96/pytorch-summary等
