优化之前每轮训练大概需要 3’10”，优化之后大概只需要 35‘’
问题提出
我有一个图像异常检测的模型，训练过程很慢，且在训练过程中GPU的使用率剧烈波动，一会100%一会又降到0。我使用的是GTX3090 24G显卡，在MVTec数据集上训练，该数据集有15个类别，单卡全部训练完大约需要4天。
寻找原因
python或pytorch本身提供了一些性能分析的工具，但我担心并行环境下，一些方法执行耗时并不能反应真实的代码性能，所以我采用一种较为简单粗暴的方法：控制变量法。不断注释或修改某些函数，或者在死循环中执行某些语句，查看执行的速度和GPU占用。
一旦找到出问题的语句就能对该代码进行优化，手段包括但不限于：避免对tensor元素的for循环、tensor直接创建到GPU上、修改pytorch中Dataloader的一些优化参数等
下面是一个具体的优化过程：
我发现代码在loss函数上执行的特别慢，为了验证，我在loss计算和backward这里使用了一个死循环
注意：pytorch的backward计算是基于计算图的，默认情况下，执行backward之后就会删掉计算图，如果尝试再次backward就会报错，解决办法就是如图所示，加上 retain_graph=True 参数即可保留特征图
代码执行过程中发现，GPU利用率剧烈上下波动，可以肯定loss函数是出了问题的，追踪代码进入到loss函数的实现，可以发现它是由三个loss相加得到的，故还是一个一个注释看效果，前两个loss计算时，GPU基本能处于满载，但是到第三个loss时，GPU使用率又开始上下翻动，再进去查看实现，发现代码并不多
依然是给可能出问题的语句使用死循环执行，发现第455行和457行执行得极慢，且当代码执行到这时，原本满载的GPU很快就降到0了，与此同时cpu使用率开始增大。问题一定出在这。
torch.randperm(num)函数的作用是随机打乱一个序列，例如这里为了增加模型的鲁棒性，会随机打乱norm和anom两个矩阵中元素的索引。通过pytorch文档查询到该函数会返回一个随机序列的tensor，问题就出在这，如果没有指定 device 参数，该tensor将在cpu上生成，由于我这里要打乱的序列长度较大（亿级别），故cpu执行时，需要1~2s，而GPU则可以瞬间完成，故指定device为目标训练设备即可（这里是GPU）
结合一些其他的代码优化，模型训练过程中可以看到GPU基本一直处于满载状态，训练速度大大提升。
优化建议
tensor直接创建到GPU
pytorch提供的很多创建tensor的方法都会有一个 device 参数用于指定tensor创建位置，可以直接指定创建到GPU。或者在cpu上创建tensor之后使用 .to(device) 函数复制到GPU上，显然直接指定device会好很多，免去了复制过程且避免使用cpu做计算，例如上面那个例子。
避免计算图中出现流程控制语句
模型forward或backward过程应始终保持矩阵运算的思想，尽量避免使用for、if等流程控制语句，它使函数变得不连续且无法发挥并行运算的优势。
例如，在本次优化过程中，我发现了一处瓶颈，在某个子模型的forward中存在对某些tensor一些维度上的遍历，代码如下：
这段代码实际上在做一个量阶的事，即给你一个数组，你需要将这个数组中的值分成n个等级，你需要计算每个等级的上下边界，其实就是连续变离散的过程。但上述代码对一个tensor进行了 2*n 次沿第一维进行的遍历，很显然这样并不会并行运算，极大浪费了GPU的计算资源，以下为我修改后的代码：
事实上，上述代码应该还能进一步优化，不过目前已经可以做到并行运算了。
缺点当然也很明显，它降低了代码的可读性，当我忘记原来的代码时干嘛的时候，我大概看看代码就能知道意思，但修改后的代码，我估计过段时间再看就会一脸懵逼，这也是我特地加上注释的原因。
尽量使用pytorch提供的api
如果需要选取矩阵中的部分元素，则可以考虑 torch.gather()、torch.meshgrid()等方法
若需要对矩阵进行运算，优先考虑pytorch内置的一些运算方法，对于比较复杂的计算可以考虑使用 torch.einsum()，它使用爱因斯坦求和约定，能够进行较为复杂的运算且支持并行运算。
给Dataloader选取合适的优化参数
具体的，比如：batch_size、worker_num、prefetch_factor、pin_memory、non_blocking、persistent_workers
batch_size 越大，相对来说数据的噪声就越小，模型更容易欠拟合（大批量的数据导致批与批之间的数据更加相似，缺乏多样性数据），当然占用显存肯定越大，训练速度越快。但是我记得我之前看过一些论文说它还会影响模型的一些其他什么性能。
pin_memory 设置为True，则会在显存中固定一片区域来存放batch数据，这样就省去了每次获取数据的寻址过程，不过在我这个实验中作用不明显。
此外，pytorch中.to(device)函数还有个 non_blocking 参数，用于异步将tensor复制到指定设备上。原理是已经传输的数据就能开始参与计算，还没传输完成的部分继续传输。但是在我本次实验中，效果不明显。
worker_num 是一个相对较为重要的参数，它指定使用多少条额外线程加载数据（在load data过程中，可能会伴随一些transformer），该参数默认为0，表示模型的训练和数据加载都使用同一条线程，那么就加载一次数据，训练一次。gpu每次都需要等待数据加载，及其耗费资源。
当把该参数值设定为大于0时，就会在加载数据时同时启用多条线程加载数据（加载数量为 prefetch_factor * worker_num），当加载数据的速度大于gpu处理数据的速度时，gpu就能处于满载状态而不需要等待数据加载了。本次实验时，若将workers_num设置为4时，可以明显看到GPU使用率会时不时降下去，而训练过程也会偶尔变慢，这就是数据加载速度小于GPU处理速度导致的，若将该值设定为8则没有出现这种情况了，但两个epoch之间的时间间隔明显增大了（上一个epoch结束到下一个epoch开始），这是因为worker_num越多，训练开始时需要加载的数据批次越多（需要所有的worker都加载好数据），训练准备工作的时间也就越长。当讲该值设定为16时，两个epoch的间隔更长了，但训练速度没有提升，也就是数据加载速度过剩。所以这三个数中，8是最合适的。
prefetch_factor 参数指定每个worker预加载批次的数量，例如当该参数指定为2时，则每个worker在加载完一个批次的数据后，还需要再继续加载一个批次数据缓存起来。每个worker始终保持有两个批次数据待处理。这样可以提升数据加载速度，但我本次实验效果并不明显。我想，在资源有限的情况下可以少点worker多点prefetch，然而资源较为充裕的情况下，还是多点worker更好。
persistent_workers 默认为false，表示当dataset消耗完一次之后要不要销毁worker，在我本次实验中，我开启10个worker，两个epoch之间的时间数据加载时间明显减少（11s减少到7s）
cuDNN加速
cuDNN是一个专门为深度学习开发的库，参见 深度学习GPU环境准备
默认情况下会使用这个库，但这个库在使用过程中会出现一些随机性，所以如果希望每次训练结果相同则可以将其关闭。
我本地有一个模型，每次训练时间大约8s左右，我记得以前这个模型训练速度应该没这么慢才是，经过反复排查，问题就出现在这个库上，我之前为了保证结果复现把这个关掉了。开启之后，每次训练时间骤减到3s（因为是图像模型，大量使用了卷积操作，benchmark = True在这里也发挥了大约1s的作用）。
# 开启cudnn加速
torch.backends.cudnn.enabled = True
# 开启后cudnn会自动根据你的模型选择一些最优算法
torch.backends.cudnn.benchmark = True
# 将这个设置为true就不会选择最优算法了，使用确定性算法保证结果复现
torch.backends.cudnn.deterministic = Falsetorch.einsum()
einsum()函数使用的是爱因斯坦求和约定的标记，什么是爱因斯坦求和约定？在我的理解看就是一种对符号简化的符号。
当我要计算
\[ [a_0, a_1, a_2] · [b_0, b_1, b_2]^T \]
时，我可以写成：
\[ a_0·b_0 + a_1·b_1 + a_2·b_2 = \sum_{i=0}^2{a_ib_i} = a_ib_i \]
上面三个式子都表示一个意思，只不过后一个是前一个的简写形式。而最右边的式子就是爱因斯坦求和约定的写法，当两个元素角标相同时（例如上面中间式子中，a和b都有相同的角标i）就能省略求和符号。
以上这个式子若体现在pytorch或numpy中就应该是：
其中，i表示a和b两个数组相同的下标，而箭头后面留空则表示求和，你可能会想，求的是什么的和？事实上，箭头左边的式子如果是一样的下标表示，则表示这两个下标位置相乘，故上式就理解为：a和b对应位置相乘并求和。可以想到，如果不计算就和那就是一个向量了：
einsum表达式几乎能做任何矩阵的计算，下面是对einsum函数中的表达式进行解释，对于这个表达式，你只用知道三点就够了：
1. 箭头左边的输入中，两个矩阵（逗号区分）中的相同字母表示，这两个矩阵对应的这个维度要相乘
例如表达式 ij,jk->ijk 表示第一个矩阵的第二维与第二个矩阵的第一维相乘（即合并两个相同的维度），所以它们的输出就是一个三维的矩阵（各自两个维度，但是合并了一个维度）
2. 箭头右边的输出中，如果忽略了左边的某个或多个字母，则意味着计算结果矩阵中在忽略的这一维或多维求和，如果忽略所有的字母就是对整个矩阵求和
例如表达式 ij,ji->ik，它和上面的表达式相比，右边少了一个j，则表示在合并后的三维矩阵中，沿着合并的这个轴累加（即降维，讲原来的一个向量元素累加成一个标量）。
3. 你可以在箭头右边任意调整输出轴的位置
例如你想转置某个三维数组，讲其前两个轴进行交换，你当然可以使用permute()方法
你也可以使用einsum来做：
当然，你也可以边运算边转置，例如表达式 ij,jk->ki，它还是同上文中的表达式，只不过讲输出的两个字母调换顺序，意思就是计算的最终结果还需要转置成这种新的形式。
注：更详细的解释见：https://ajcr.net/Basic-guide-to-einsum/
我个人的简单理解：
表达式中的字母代表着矩阵的各个维度，若你希望两个矩阵某个相同维度合并（合并就是将两个相同维度中对应的数字相乘得到一个新的维度）时，你就在左边将它们这一维用相同的字母表示。
合并之后可以得到一个新的大矩阵，如果你希望新矩阵某些维度要降维（向量元素累加成一个标量），则在表达式右边省略代表该维度的字母。
我觉得这个不太好理解，比如这个表达式 abcde,fgbij->abcdefhi，这个表达式比较简单，只不过矩阵维度较多，可以看到，输入端分别是第二维和第三维用相同字母表示，则意味着它们这一维是一样大的且要计算它们这一维对应位置的乘积，问题是，第一个输入第二维度上是一个c*d*e的三维矩阵，而第二个输入第三维度上是一个i*j的二维数组，它们怎么相乘的？
我的理解就是，你可以将上面表达式看成，先将第一个输入的第二维转置成最后一维，那它就只是一个数字标量了，它就可以和第二个输入的第三维度中的矩阵相乘了，乘得的结果再次转置到之前的位置。当然，实际计算结果我实验也确实是这个效果。
这样虽然对理解表达式有帮助，但对使用einsum却没什么帮助，主要是我对爱因斯坦求和约定理解不深，只能记下结果，后面再找机会学习一下这个。