51c~Pytorch~合集6
我自己的原文哦~ https://blog.51cto.com/whaosoft/14269530
一、Pytorch~mm
矩阵成真！Pytorch最新工具mm，3D可视化矩阵乘法、Transformer注意力,能够将矩阵乘法模拟世界还原。
矩阵中的模拟世界，真的来了。
矩阵乘法（matmul），是机器学习中非常重要的运算，特别是在神经网络中扮演着关键角色。
Pytorch团队最新的一篇文章中，介绍了「mm」，一个用于matmuls和matmuls组合的可视化工具。
因为使用了三个空间维度，mm能够有助于建立直觉、激发想法，尤其适合（但不仅限于）视觉/空间思维者。
英伟达高级科学家Jim Fan表示，进入神经网络「矩阵」。
这是一个非常酷的可视化工具，用于矩阵、注意力、并行等等。最好的教育来自最直观的传递。这是一个具有数据并行分割功能的多层感知器。
有了三个维度来组成矩阵乘法，再加上加载训练过权重的能力，就可以用im来可视化像注意力头这样的大型复合表达式，并观察它们的实际表现。
mm工具能够交互，可在浏览器或笔记本notebook iframe中运行，并在URL中保留其完整状态，共享对话链接。
地址：https://bhosmer.github.io/mm/ref.html
博客原文：https://pytorch.org/blog/inside-the-matrix
下文中，Pytorch提供的参考指南中，介绍了mm所有可用的功能。
研究团队将首先介绍可视化方法，通过可视化一些简单的矩阵乘法、和表达式来建立直觉，然后深入研究一些更多的示例。
为什么这种可视化方式更好？
mm的可视化方法基于这样一个前提，即矩阵乘法从根本上说，是一种三维运算。
换句话说：
是一张纸，用mm打开后变成如下的样子：
当我们以这种方式，将矩阵乘法包在一个立方体周围时，参数形状、结果形状和共享维度之间的正确关系全部都会建立起来。
现在，计算就有了几何意义：
结果矩阵中的每个位置 i, j 锚定了立方体内部沿深度维度 k 运行的向量，其中从 L 中的第 i 行延伸出水平面和从 R 中的第 j 列延伸的垂直平面相交。沿着这个向量，来自左右2个参数的 (i, k) (k, j) 元素对相遇并相乘，得到的乘积沿着 k 相加，存入在结果的位置 i, j 。
这就是矩阵乘法的直观含义：
- 将两个正交矩阵投影到立方体内部
- 将每个交叉点上的一对数值相乘，形成一个乘积网格
- 沿第三个正交维度求和，得出结果矩阵
为了确定方向，mm工具会在立方体内部显示一个指向结果矩阵的箭头，蓝色指标来自左侧参数，红色指标来自右侧参数。
该工具还会显示白色指引线，以指示每个矩阵的行轴，不过在这张截图中这些指引线很模糊。
对于方向，该工具在多维数据集内部显示一个指向结果矩阵的箭头，蓝色叶片来自左参数，红色叶片来自右参数。该工具还显示白色指南来指示每个矩阵的行轴，尽管它们在这个屏幕截图中很模糊。
当然，布局限制简单明了：
- 左参数和结果必须沿着它们共享的高度 (i) 维度相邻
- 右参数和结果必须沿其共享的宽度 (j) 维度相邻
- 左参数和右参数必须沿着它们共享的（左宽/右高）维度相邻，这就是矩阵乘法的深度 (k) 维度
这个几何图形，为我们提供了可视化所有标准矩阵乘法分解的坚实基础，以及探索矩阵乘法的非难复杂组合的直观依据。
下面，我们就会看到真正的矩阵世界。
规范矩阵乘法分解动作
在深入研究一些更复杂的示例之前，Pytorch团队将介绍一些直觉构建器，以了解事物在这种可视化风格中的外观和感觉。
点
首先是标准算法。通过对相应的左行和右列进行点乘计算每个结果元素。
我们在动画中看到的是乘法值矢量在立方体内部的扫描，每个矢量都会在相应位置产生一个求和结果。
这里， L 的行块填充为1（蓝色）或-1（红色）；R 的列块填充类似。k 在这里是24，因此结果矩阵（ L @ R ）的蓝色值为24，红色值为-24。
矩阵-向量乘积
分解为矩阵向量乘积的matmul，看起来像一个垂直平面（左参数与右参数每列的乘积），当它水平扫过立方体内部时，将列绘制到结果上。
即使在简单的例子中，观察分解的中间值也会非常有趣。
例如，当我们使用随机初始化参数时，请注意中间矩阵-向量乘积中突出的垂直模式。这反映出每个中间值都是左参数的列缩放复制品：
向量-矩阵乘积
分解为向量-矩阵乘积的矩阵乘法在穿过立方体内部时，看起来就像在结果上绘制行的水平面：
切换到随机初始化参数时，我们会看到与矩阵-向量乘积类似的模式，只不过这次的模式是水平的，因为每个中间向量-矩阵乘积都是右参数的行缩放复制品。
在思考矩阵乘法如何表达其参数的秩和结构时，不妨设想一下在计算中同时出现这两种模式的情况：
这里还有一个使用向量矩阵乘积的直觉构建器，显示单位矩阵如何像镜子一样，以45度角设置其反参数和结果：
求和外积
第三个平面分解沿k轴进行，通过向量外积的点和计算出矩阵乘法结果。
在这里，我们看到外积平面「从后向前」扫过立方体，累积成结果：
使用随机初始化的矩阵进行这种分解，我们可以看到，随着每个秩-1 外积的增加，结果中不仅有数值，还有秩的累积。
除其他外，这也让我们明白了为什么「低秩因式分解」，即通过构建深度维度参数很小的矩阵乘法来逼近矩阵，在被逼近的矩阵是低秩矩阵时效果最好。
LoRA 将在后面介绍：
表达式
如何将这种可视化方法扩展到矩阵乘法的组合？
到目前为止，示例可视化了某个矩阵 L 和 R 的单一矩阵 L @ R ，如果 L 和/或 R 本身就是矩阵，并以此类推呢？
事实证明，我们可以很好地将该方法扩展到复合表达式。
关键规则很简单：子表达式（子）矩阵乘法是另一个立方体，受与父表达式相同的布局约束，子表达式的结果面同时是父表达式的相应参数面，就像共价键共享的电子一样。
在这些限制条件下，我们可以随意排列子matmul的面。
在这里，研究人员使用了工具的默认方案，即交替生成凸面和凹面的立方体，这种布局在实践中非常有效，可以最大限度地利用空间并减少遮挡。
在本节中，Pytorch将对ML模型中的一些关键构件进行可视化处理，以熟练掌握可视化习惯用语，并了解即使是简单的示例也能给我们带来哪些直观感受。
左关联表达式
先来看两个(A @ B) @ C形式的表达式，每个表达式都有自己独特的形状和特征。
首先，我们将赋予 A @ B @ C以FFN的特征形状，其中「隐藏维度」比「输入」或「输出」维度更宽。(在本例中，这意味着B的宽度大于A或C的宽度）。
与单个matmul例子一样，浮动箭头指向结果矩阵，蓝色来自左参数，红色来自右参数：
接下来，将可视化 A @ B @ C ， B的宽度比A或C窄，使其呈现一个瓶颈或「自动编码器」形状：
这种凸块和凹块交替出现的模式可以扩展到任意长度的链：例如这个多层瓶颈：
右关联表达式
接下来，将可视化一个右关联表达式A @ (B @ C) 。
有时，我们会看到 MLP 采用右侧关联方式，即输入层在右侧，权重层从右至左。
使用上图中的双层FFN例子中的矩阵--经过适当换位--如下所示，C现在扮演输入的角色，B是第一层，A是第二层：
二进制表达式
可视化工具要想超越简单的教学示例之外发挥作用，就必须在表达式变得越来越复杂时，保持可读性。
在现实世界的使用案例中，二进制表达式是一个关键的结构组件，即左右两边都有子表达式的矩阵。
在这里，将可视化最简单的表达式形状，(A @ B) @ (C @ D) ：
分割和并行性
下面，可以通过两个快速示例来了解这种可视化风格是如何通过简单的几何分割，使复合表达式的并行化推理变得非常直观的。
在第一个例子中，把规范的「数据并行」分割应用到，上述左关联多层瓶颈的例子中。
沿着i进行分割，分割左初始参数（批）和所有中间结果（激活），但不分割后续参数（权重）。
通过几何图形，我们可以清楚地看到表达式中哪些参与者被分割，哪些参与者保持完整：
第二个例子，展示了如何通过沿着其 j 轴划分左子表达式、沿着其 i 轴划分右子表达式以及沿着其 k 轴划分父表达式来并行化二进制表达式：
注意力头内部
让我们看一个GPT2注意力头——特别是来自NanoGPT的「gpt2」（小）配置（层=12，头=12，嵌入=768）的第5层，通过HuggingFace使用OpenAI权重。
输入激活来自256个token的OpenWebText训练样本的前向传递。
研究人员之所以选择它，主要是因为它计算了一种相当常见的注意力模式，而且位于模型的中间位置，这里的激活已经变得结构化，并显示出一些有趣的纹理。
结构
整个注意力头被可视化为一个单一的复合表达式，以输入开始，以投影输出结束。(注：为了保持自成一体，研究人员按照Megatron-LM中的描述对每个头进行输出投影）。
计算包含六个矩阵：
Q = input @ wQ // 1
K_t = wK_t @ input_t // 2
V = input @ wV // 3
attn = sdpa(Q @ K_t) // 4
head_out = attn @ V // 5
out = head_out @ wO // 6我们正在查看的内容的缩略图描述：
箭头叶片是矩阵乘法1、2、3和6：前一组是从输入到 Q、K 和 V 的内投影；后一组是从 attn @ V 回到嵌入维度的外投影。
在中心是双矩阵乘法，它首先计算注意力分数（后面的凸立方体），然后使用它们从值向量（前面的凹立方体）生成输出token。因果关系意味着注意力分数形成一个下三角。
计算和值
这是一个计算注意力的动画。具体来说：
sdpa(input @ wQ @ K_t) @ V @ wO(即上面的矩阵1、4、5 和 6，其中 K_t 和 V 是预先计算好的）的计算过程是一个融合的向量矩阵乘积链：序列中的每个项目都是一步完成，从输入到注意力再到输出的整个过程。
头的不同之处
继续下一步之前，这里还有一个演示，可以让我们简单地了解一下模型的详细工作原理。
这是GPT2的另一个注意头。
它的行为与上面的第5层第4个头截然不同，正如所预料的那样，因为它位于模型的一个非常不同的部分。
这个磁头位于第一层：第0层，头2：
并行注意力
我们将注意力头中的 6个矩阵中的4个可视化为融合的向量矩阵乘积链。
是一个融合了向量-矩阵乘积的链条，证实了从输入到输出的整个左关联链条沿着共享的 i 轴是层状的这一几何直觉，并且可以并行化。
比如沿着i分割
双分区
LoRA
最近的LoRA论文描述了一种高效的微调技术，该技术基于微调期间引入的权重增量是低秩的想法。
根据该论文，这使我们能够通过优化密集层在适应过程中变化的秩分解矩阵，间接训练神经网络中的一些密集层，同时保持预先训练的权重冻结。
基本思想
简而言之，关键的一步是训练权重矩阵的因子，而不是矩阵本身：用 I x K 张量和 K x J 张量的matmul替换 I x J 权重张量，保持 K 为某个小数字。
如果 K 足够小，那么所节省的大小将非常可观，但代价是降低 K 会降低乘积所能表达的秩。
下面是一个随机左 128 x 4 和右 4 x 128 参数的 matmul，也就是 128 x 128 矩阵的秩-4因式分解，可以快速说明节省的大小和对结果的结构影响。请注意 L @ R 的垂直和水平模式（：
参考资料：
https://twitter.com/PyTorch/status/1706384907377770884
https://twitter.com/DrJimFan/status/1706690238083744218
...
二、PyTorch2.2
PyTorch也迎来了重大更新，PyTorch 2.2集成了FlashAttention-2和AOTInductor等新特性，计算性能翻倍。
继去年十月份的PyTorch大会发布了2.1版本之后，全世界各地的521位开发者贡献了3628个提交，由此形成了最新的PyTorch 2.2版本。
新的版本集成了FlashAttention-2，使得scaled_dot_product_attention （SDPA）相较于之前的版本有了约2倍的性能提升。
PyTorch 2.2还引入了一个新的TorchInductor提前扩展，称为 AOTInductor，旨在为非python服务器端编译和部署PyTorch程序。
PyTorch中的torch.distributed支持了一个叫做device_mesh的新抽象，用于初始化和表示ProcessGroups。
另外，PyTorch 2.2提供了一个标准化的、可配置的日志记录机制，——TORCH_LOGS。
PyTorch 2.2还对torch.compile做了许多改进，包括改进了对编译优化器的支持，以及TorchInductor融合和布局优化。
最后值得注意的是，PyTorch将放弃对macOS x86的支持，PyTorch 2.2.x是支持macOS x64的最后一个版本。
PyTorch 2.2新特性
首先请注意，如果从源代码构建PyTorch 2.2，需要GCC 9.4或更高版本，PyTorch 代码库已从C++ 14迁移到C++ 17。
FlashAttention-2
FlashAttention-2通过优化GPU上不同线程块和warps之间的工作分区，来解决占用率低或不必要的共享内存读写。
FlashAttention-2调整了算法以减少非matmul的计算量，同时提升了Attention计算的并行性（即使是单个头，也可以跨不同的线程块，以增加占用率），在每个线程块中，优化warps之间的工作分配，以减少通过共享内存的通信。
PyTorch 2.2将FlashAttention内核更新到了v2版本，不过需要注意的是，之前的Flash Attention内核具有Windows实现，Windows用户可以强制使用sdp_kernel，仅启用Flash Attention的上下文管理器。
而在2.2中，如果必须使用 sdp_kernel 上下文管理器，请使用memory efficient或math内核（在Windows上）。
在FlashAttention-2的加持之下，torch.nn.functional.scaled_dot_product_attention的速度提升了大约2倍，在A100 GPU上达到了理论计算峰值的50%-73%。
AOTInductor
AOTInductor是TorchInductor的扩展，用于处理导出的PyTorch模型，对其进行优化，并生成共享库以及其他相关工件。
这些编译的工件可以部署在非Python环境中，经常用于服务器端的推理。
下面的示例演示了如何调用 aot_compile 将模型转换为共享库。
AOTInductor支持与Inductor相同的后端，包括CUDA、ROCm和CPU。
TORCH_LOGS
PyTorch 2.2提供了一个标准化的、可配置的日志记录机制，可用于分析各种子系统的状态，例如编译和分布式操作
可以通过TORCH_LOGS环境变量启用日志。比如通过在命令行中修改环境变量：
将TorchDynamo的日志级别设置为logging.ERROR，将TorchInductor的日志级别设置为logging.DEBUG。
当然也可以在代码中以API的形式使用：
torch.distributed.device_mesh
PyTorch 2.2引入了一个新的抽象，用于表示分布式并行中涉及的 ProcessGroups，称为torch.distributed.device_mesh。
为分布式训练设置分布式通信器（NCCL）是一件麻烦的事情。用户需要编写不同并行度的工作负载，并为每个并行度手动设置和管理NCCL通信器（ProcessGroup ）。
这个过程可能很复杂，容易出错。而DeviceMesh 可以简化此过程，使其更易于管理。
DeviceMesh 是管理 ProcessGroup 的更高级别的抽象。它允许用户毫不费力地创建节点间和节点内进程组，而不必担心如何为不同的子进程组正确设置等级。
例如，数组的其中一个维度可以表示FSDP中的数据并行（data parallelism），而另一个维度可以表示FSDP中的张量并行（tensor parallelism）。
用户还可以通过 DeviceMesh 轻松管理底层process_groups，以实现多维并行。
DeviceMesh在处理多维并行性（如3D并行）时很有用。如上图所示，当你的并行解决方案需要跨主机和每个主机内部进行通信时，可以创建一个2D网格，用于连接每个主机中的设备，并以同构设置将每个设备与其他主机上的对应设备连接起来。
借助 init_device_mesh() ，我们可以在短短两行内完成上面这个2D设置：
而如果不使用DeviceMesh，我们大概需要自己写下面这一堆代码：
当然，如果需要，我们仍然可以访问底层 ProcessGroup：
优化器的改进
大概有以下几点：
- 编译优化器在所有基准测试中都提高了性能：HuggingFace +18%、TorchBench +19%、TIMM +8% E2E；
- 编译的优化器增加对cudagraphs的支持；
- 对测试套件中所有模型进行平均，每个测试套件的基准测试平均编译时间增加约40秒；正在进行的优化可能会将其降低到30秒以下。
用于多张量优化器编译的inductor中缺少的主要功能是foreach算子的高效编码生成。
在调度器内部，将所有在下放过程中注册的缓冲区列表凝聚到ForeachKernelSchedulerNodes中（FusedSchedulerNode的子类）。
为了检查融合是否合法，每个内部 SchedulerNode 执行的写操作必须与消费SchedulerNode在同一列表索引处的读操作相匹配。
此外，正常的垂直融合规则必须允许在消费者和生产者SchedulerNode列表的每个索引处进行融合。
如果满足了这些条件，ForeachKernelSchedulerNode将垂直融合成一个 ForeachKernelSchedulerNode，其中每个列表上的相应点操作都将被融合。
通过实现这种融合，可以将一系列 foreach 运算融合到单个内核中，从而实现多张量优化器的完全融合。
性能改进
TorchInductor中添加了许多性能优化，包括对torch.concat的水平融合支持、改进的卷积布局优化、以及改进scaled_dot_product_attention模式匹配。
PyTorch 2.2还包括aarch64的许多性能增强，包括对mkldnn权重预打包的支持、改进的ideep基元缓存，以及通过对OneDNN的固定格式内核改进，来提高推理速度。
参考资料：
https://pytorch.org/blog/pytorch2-2/
...
三、拯救 4G 显卡： PyTorch 节省显存的策略总结
0 前言
本文涉及到的 PyTorch 节省显存的策略包括：
- 混合精度训练
- 大 batch 训练或者称为梯度累加
- gradient checkpointing 梯度检查点
1 混合精度训练
混合精度训练全称为 Automatic Mixed Precision，简称为 AMP，也就是我们常说的 FP16。在前系列解读中已经详细分析了 AMP 原理、源码实现以及 MMCV 中如何一行代码使用 AMP，具体链接见：
OpenMMLab：PyTorch 源码解读之 torch.cuda.amp: 自动混合精度详解
https://zhuanlan.zhihu.com/p/348554267
OpenMMLab：OpenMMLab 中混合精度训练 AMP 的正确打开方式
https://zhuanlan.zhihu.com/p/375224982
由于前面两篇文章已经分析的非常详细了，本文只简要描述原理和具体说明用法。
考虑到训练过程中梯度幅值大部分是非常小的，故训练默认是 FP32 格式，如果能直接以 FP16 格式精度进行训练，理论上可以减少一半的内存，达到加速训练和采用更大 batch size 的目的，但是直接以 FP16 训练会出现溢出问题，导致 NAN 或者参数更新失败问题，而 AMP 的出现就是为了解决这个问题，其核心思想是 混合精度训练+动态损失放大：
- 维护一个 FP32 数值精度模型的副本
- 在每个 iteration
- 拷贝并且转换成 FP16 模型
- 前向传播（FP16 的模型参数)，此时 weights, activations 都是 FP16
- loss 乘 scale factor s
- 反向传播（FP16 的模型参数和参数梯度)， 此时 gradients 也是 FP16
- 参数梯度乘 1/s
- 利用 FP16 的梯度更新 FP32 的模型参数
在 MMCV 中使用 AMP 分成两种情况：
- 在 OpenMMLab 上游库例如 MMDetection 中使用 MMCV 的 AMP
- 用户只想简单调用 MMCV 中的 AMP，而不依赖上游库
(1) OpenMMLab 上游库如何使用 MMCV 的 AMP
以 MMDectection 为例，用法非常简单，只需要在配置中设置：
fp16 = dict(loss_scale=512.) # 表示静态 scale # 表示动态 scale
fp16 = dict(loss_scale='dynamic') # 通过字典形式灵活开启动态 scale
fp16 = dict(loss_scale=dict(init_scale=512.,mode='dynamic'))三种不同设置在大部分模型上性能都非常接近，如果不想设置 loss_scale，则可以简单的采用 loss_scale='dynamic'
(2) 调用 MMCV 中的 AMP
直接调用 MMCV 中的 AMP，这通常意味着用户可能在其他库或者自己写的代码库中支持 AMP 功能。需要特别强调的是 PyTorch 官方仅仅在 1.6 版本及其之后版本中开始支持 AMP，而 MMCV 中的 AMP 支持 1.3 及其之后版本。如果你想在 1.3 或者 1.5 中使用 AMP，那么使用 MMCV 是个非常不错的选择。
使用 MMCV 的 AMP 功能，只需要遵循以下几个步骤即可：
- 将 auto_fp16 装饰器应用到 model 的 forward 函数上
- 设置模型的 fp16_enabled 为 True 表示开启 AMP 训练，否则不生效
- 如果开启了 AMP，需要同时配置对应的 FP16 优化器配置 Fp16OptimizerHook
- 在训练的不同时刻，调用 Fp16OptimizerHook，如果你同时使用了 MMCV 中的 Runner 模块，那么直接将第 3 步的参数输入到 Runner 中即可
- (可选) 如果对应某些 OP 希望强制运行在 FP32 上，则可以在对应位置引入 force_fp32 装饰器
# 1 作用到 forward 函数中
class ExampleModule(nn.Module):@auto_fp16()def forward(self, x, y):return x, y# 2 如果开启 AMP，则需要加入开启标志
model.fp16_enabled = True # 3 配置 Fp16OptimizerHook
optimizer_config = Fp16OptimizerHook(**cfg.optimizer_config, **fp16_cfg, distributed=distributed)# 4 传递给 runner
runner.register_training_hooks(cfg.lr_config, optimizer_config,cfg.checkpoint_config, cfg.log_config,cfg.get('momentum_config', None)) # 5 可选
class ExampleModule(nn.Module):@auto_fp16()def forward(self, x, y):features=self._forward(x, y)loss=self._loss(features,labels)return lossdef _forward(self, x, y):pass@force_fp32(apply_to=('features',))def _loss(features,labels) :pass注意 force_fp32 要生效，依然需要 fp16_enabled 为 True 才生效。
2 大 Batch 训练（梯度累加）
大 Batch 训练通常也称为梯度累加策略，通常 PyTorch 一次迭代训练流程为：
y_pred = model(xx)
loss = loss_fn(y_pred, y)
loss.backward()
optimizer.step()
optimizer.zero_grad()而梯度累加策略下常见的一次迭代训练流程为：
y_pred = model(xx)
loss = loss_fn(y_pred, y)loss = loss / cumulative_iters
loss.backward()if current_iter % cumulative_iters==0optimizer.step() optimizer.zero_grad()其核心思想就是对前几次梯度进行累加，然后再统一进行参数更新，从而变相实现大 batch size 功能。需要注意的是如果模型中包括 BN 等考虑 batch 信息的层，那么性能可能会有轻微的差距。
细节可以参考：
https://github.com/open-mmlab/mmcv/pull/1221
在 MMCV 中已经实现了梯度累加功能，其核心代码位于 mmcv/runner/hooks/optimizer.py
GradientCumulativeOptimizerHook 中，和 AMP 实现一样是采用 Hook 实现的。使用方法和 AMP 类似，只需要将第一节中的 Fp16OptimizerHook 替换为 GradientCumulativeOptimizerHook 或者 GradientCumulativeFp16OptimizerHook 即可。其核心实现如下所示：
@HOOKS.register_module()
class GradientCumulativeOptimizerHook(OptimizerHook):def __init__(self, cumulative_iters=1, **kwargs):self.cumulative_iters = cumulative_itersself.divisible_iters = 0 # 剩余的可以被 cumulative_iters 整除的训练迭代次数self.remainder_iters = 0 # 剩余累加次数self.initialized = Falsedef after_train_iter(self, runner):# 只需要运行一次即可if not self.initialized:self._init(runner)if runner.iter < self.divisible_iters:loss_factor = self.cumulative_iterselse:loss_factor = self.remainder_itersloss = runner.outputs['loss']loss = loss / loss_factorloss.backward()if (self.every_n_iters(runner, self.cumulative_iters)or self.is_last_iter(runner)):runner.optimizer.step()runner.optimizer.zero_grad() def _init(self, runner):residual_iters = runner.max_iters - runner.iterself.divisible_iters = (residual_iters // self.cumulative_iters * self.cumulative_iters)self.remainder_iters = residual_iters - self.divisible_itersself.initialized = True需要明白 divisible_iters 和 remainder_iters 的含义：
(1) 从头训练
此时在开始训练时 iter=0，一共迭代 max_iters=102 次，梯度累加次数是 4，由于 102 无法被 4 整除，也就是最后的 102-(102 // 4)*4=2 个迭代是额外需要考虑的，在最后 2 个训练迭代中 loss_factor 不能除以 4，而是 2，这样才是最合理的做法。其中 remainder_iters=2，divisible_iters=100，residual_iters=102。
(2) resume 训练
假设在梯度累加的中途退出，然后进行 resume 训练，此时 iter 不是 0，由于优化器对象需要重新初始化，为了保证剩余的不能被累加次数的训练迭代次数能够正常计算，需要重新计算 residual_iters。
3 梯度检查点
梯度检查点是一种用训练时间换取显存的办法，其核心原理是在反向传播时重新计算神经网络的中间激活值而不用在前向时存储，torch.utils.checkpoint 包中已经实现了对应功能。简要实现过程是：在前向阶段传递到 checkpoint 中的 forward 函数会以 _torch.no_grad_ 模式运行，并且仅仅保存输入参数和 forward 函数，在反向阶段重新计算其 forward 输出值。
具体用法非常简单，以 ResNet 的 BasicBlock 为例：
def forward(self, x):def _inner_forward(x):identity = xout = self.conv1(x)out = self.norm1(out)out = self.relu(out)out = self.conv2(out)out = self.norm2(out)if self.downsample is not None:identity = self.downsample(x)out += identityreturn out# x.requires_grad 这个判断很有必要if self.with_cp and x.requires_grad:out = cp.checkpoint(_inner_forward, x)else:out = _inner_forward(x)out = self.relu(out)return outself.with_cp 为 True，表示要开启梯度检查点功能。
checkpoint 在用法上面需要注意以下几点：
- 模型的第一层不能用 checkpoint 或者说 forward 输入中不能所有输入的 requires_grad 属性都是 False，因为其内部实现是依靠输入的 requires_grad 属性来判断输出返回是否需要梯度，而通常模型第一层输入是 image tensor，其 requires_grad 通常是 False。一旦你第一层用了 checkpoint，那么意味着这个 forward 函数不会有任何梯度，也就是说不会进行任何参数更新，没有任何使用的必要，具体见 https://discuss.pytorch.org/t/use-of-torch-utils-checkpoint-checkpoint-causes-simple-model-to-diverge/116271。如果第一层用了 checkpoint， PyTorch 会打印
None of the inputs have requires_grad=True. Gradients will be Non 警告 - 对于 dropout 这种 forward 存在随机性的层，需要保证 preserve_rng_state 为 True (默认就是 True，所以不用担心)，一旦标志位设置为 True，在 forward 会存储 RNG 状态，然后在反向传播的时候读取该 RNG，保证两次 forward 输出一致。如果你确定不需要保存 RNG，则可以设置 preserve_rng_state 为 False，省掉一些不必要的运行逻辑
- 其他注意事项，可以参考官方文档 https://pytorch.org/docs/stable/checkpoint.html#
其核心实现如下所示：
class CheckpointFunction(torch.autograd.Function):@staticmethoddef forward(ctx, run_function, preserve_rng_state, *args):# 检查输入参数是否需要梯度check_backward_validity(args)# 保存必要的状态ctx.run_function = run_functionctx.save_for_backward(*args)with torch.no_grad():# 以 no_grad 模型运行一遍outputs = run_function(*args)return outputs@staticmethoddef backward(ctx, *args):# 读取输入参数inputs = ctx.saved_tensors# Stash the surrounding rng state, and mimic the state that was# present at this time during forward. Restore the surrounding state# when we're done.rng_devices = []with torch.random.fork_rng(devices=rng_devices, enabled=ctx.preserve_rng_state):# detach 掉当前不需要考虑的节点detached_inputs = detach_variable(inputs)# 重新运行一遍with torch.enable_grad():outputs = ctx.run_function(*detached_inputs)if isinstance(outputs, torch.Tensor):outputs = (outputs,)# 计算该子图梯度torch.autograd.backward(outputs, args)grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else inpfor inp in detached_inputs)return (None, None) + grads4 实验验证
为了验证上述策略是否真的能够省显存，采用 mmdetection 库进行验证，基本环境如下：
显卡: GeForce GTX 1660
PyTorch: 1.7.1
CUDA Runtime 10.1
MMCV: 1.3.16
MMDetection: 2.17.0(1) base
- 数据集：pascal voc
- 算法是 retinanet，对应配置文件为 retinanet_r50_fpn_1x_voc0712.py
- 为了防止 lr 过大导致训练出现 nan，需要将 lr 设置为 0.01/8=0.00125
- bs 设置为 2
(2) 混合精度 AMP
在 base 配置基础上新增如下配置即可：
fp16 = dict(loss_scale=512.)(3) 梯度累加
在 base 配置基础上替换 optimizer_config 为如下：
# 累加2次
optimizer_config = dict(type='GradientCumulativeOptimizerHook', cumulative_iters=2)(4) 梯度检查点
在 base 配置基础上在 backbone 部分开启 with_cp 标志即可：
model = dict(backbone=dict(with_cp=True),bbox_head=dict(num_classes=20))每个实验总共迭代 1300 次，统计占用显存、训练总时长。
| 配置 | 显存占用(MB) | 训练时长 |
| base | 2900 | 7 分 45 秒 |
| 混合精度 AMP | 2243 | 36 分 |
| 梯度累加 | 3177 | 7 分 32 秒 |
| 梯度检查点 | 2590 | 8 分 37 秒 |
- 对比 base 和 AMP 可以发现，由于实验显卡是不支持 AMP 的，故只能节省显存，速度会特别慢，如果本身显卡支持 AMP 则可以实现在节省显存的同时提升训练速度
- 对比 base 和梯度累加可以发现，在相同 bs 情况下，梯度累加 2 次相当于 bs 扩大一倍，但是显存增加不多。如果将 bs 缩小一倍，则可以实现在相同 bs 情况下节省大概一倍显存
- 对比 base 和梯度检查点可以发现，可以节省一定的显存，但是训练时长会增加一些
从上面简单实验可以发现，AMP、梯度累加和梯度检查点确实可以在不同程度减少显存，而且这三个策略是正交的，可以同时使用。
延庆川北小区45孙老师 东屯 收卖废品破烂垃圾炒股 废品孙
买了一个带独显的笔记本也可以在家玩玩 gpu了
5 总结
本文简要描述了三个在 MMCV 中集成且可以通过配置一行开启的节省显存策略，这三个策略比较常用也比较成熟。随着模型规模的不断增长，也出现了很多新的策略，例如模型参数压缩、动态显存优化、使用 CPU 内存暂存策略以及分布式情况下 PyTorch 1.10 最新支持的 ZeroRedundancyOptimizer 等等。
快速链接直达 MMCV 算法库，欢迎大家 Star：
https://github.com/open-mmlab/mmcv
...
四、PyTorch~Faster RCNNの小麦麦穗检测
这里基于PyTorch框架，实现通过Faster RCNN算法检测图像中的小麦麦穗。当然，用YOLO算法也同样能够完成。本文最终实现的效果如下：
麦穗检测示例
一、数据下载
数据集名：Global Wheat Head Dataset
下载地址：www.kaggle.com/c/global-wheat-detection
更多深度学习数据集：https://www.cvmart.net/dataSets
相关论文：Global Wheat Head Detection (GWHD) Dataset: A Large and Diverse Dataset of High-Resolution RGB-Labelled Images to Develop and Benchmark Wheat Head Detection Methods
数据描述：全球麦穗数据集由来自7个国家的9个研究机构领导，东京大学、国家农业、营养和环境研究所、Arvalis、ETHZ、萨斯喀彻温大学、昆士兰大学、南京农业大学和洛桑研究所。包括全球粮食安全研究所、DigitAg、Kubota和Hiphen在内的许多机构都加入了这些机构的行列，致力于精确的小麦麦穗检测。
数据集贡献机构
数据集为室外小麦植物图像，包括来自全球各地不同平台采集的4698张RGB图像，标记了193,634个小麦麦穗，1024×1024像素，每张图像含有20~70个麦穗。2020年通过Kaggle举办了相关比赛，并在2021年更新了数据集。该数据集可以用于麦穗检测，评估穗数和大小。研究成果有助于准确估计不同品种小麦麦穗的密度和大小。
数据集示例
二、代码实战
2.1 导入所需要的包
# 导入所需要的包
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch import torch.nn as nn
import albumentations as A # pip install albumentations==1.1.0
from albumentations.pytorch import ToTensorV2
import torchvision
from torchvision import datasets,transforms
from tqdm import tqdm
import cv2
from torch.utils.data import Dataset,DataLoader
import torch.optim as optim
from PIL import Image
import os
import torch.nn.functional as F
import ast2.2 参数配置
# 定义参数
LR = 1e-4
SPLIT = 0.2
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
BATCH_SIZE = 4
EPOCHS = 2
DATAPATH = '../global-wheat-detection'2.3 读取数据
# 读取 train.csv文件
df = pd.read_csv(DATAPATH + '/train.csv')
df.bbox = df.bbox.apply(ast.literal_eval) # # 将string of list 转成list数据 # # 利用groupby 将同一个image_id的数据进行聚合，方式为list进行，并且用reset_index直接转变成dataframe
df = df.groupby("image_id")["bbox"].apply(list).reset_index(name="bboxes")2.4 划分数据
# # 划分数据集
def train_test_split(dataFrame,split): len_tot = len(dataFrame) val_len = int(split*len_tot) train_len = len_tot-val_len train_data,val_data = dataFrame.iloc[:train_len][:],dataFrame.iloc[train_len:][:] return train_data,val_data len(df) train_data_df,val_data_df = train_test_split(df,SPLIT) # 划分 train val 8:2
len(train_data_df), len(val_data_df) # 查看数据
train_data_df2.5 构建Dataset类
# 定义WheatDataset 返回 图片，标签
class WheatDataset(Dataset): def __init__(self,data,root_dir,transform=None,train=True): self.data = data self.root_dir = root_dir self.image_names = self.data.image_id.values self.bboxes = self.data.bboxes.values self.transform = transform self.isTrain = train def __len__(self): return len(self.data) def __getitem__(self,index):
# print(self.image_names)
# print(self.bboxes) img_path = os.path.join(self.root_dir,self.image_names[index]+".jpg") # 拼接路径 image = cv2.imread(img_path, cv2.IMREAD_COLOR) # 读取图片 image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32) # BGR2RGB image /= 255.0 # 归一化 bboxes = torch.tensor(self.bboxes[index],dtype=torch.float64)
# print(bboxes) """ As per the docs of torchvision we need bboxes in format (xmin,ymin,xmax,ymax) Currently we have them in format (xmin,ymin,width,height) """ bboxes[:,2] = bboxes[:,0]+bboxes[:,2] # 格式转换 (xmin,ymin,width,height)-----> (xmin,ymin,xmax,ymax) bboxes[:,3] = bboxes[:,1]+bboxes[:,3]
# print(image.size,type(image)) """ we need to return image and a target dictionary target: boxes,labels,image_id,area,iscrowd """ area = (bboxes[:,3]-bboxes[:,1])*(bboxes[:,2]-bboxes[:,0]) # 计算面积 area = torch.as_tensor(area,dtype=torch.float32) # there is only one class labels = torch.ones((len(bboxes),),dtype=torch.int64) # 标签 # suppose all instances are not crowded iscrowd = torch.zeros((len(bboxes),),dtype=torch.int64) target = {} # target是个字典 里面 包括 boxes,labels,image_id,area,iscrowd target['boxes'] = bboxes target['labels']= labels target['image_id'] = torch.tensor([index]) target["area"] = area target['iscrowd'] = iscrowd if self.transform is not None: sample = { 'image': image, 'bboxes': target['boxes'], 'labels': labels } sample = self.transform(**sample) image = sample['image'] # 沿着一个新维度对输入张量序列进行连接。 序列中所有的张量都应该为相同形状，
# 把多个2维的张量凑成一个3维的张量；多个3维的凑成一个4维的张量…以此类推，也就是在增加新的维度进行堆叠 target['boxes'] = torch.stack(tuple(map(torch.tensor, zip(*sample['bboxes'])))).permute(1, 0) return image,target2.6 数据增强
# 训练与验证数据增强，利用albumentations 随机翻转转换，随机图片处理
# 对象检测的增强与正常增强不同，因为在这里需要确保 bbox 在转换后仍然正确与对象对齐
train_transform = A.Compose([ A.Flip(0.5), ToTensorV2(p=1.0)
],bbox_params = {'format':"pascal_voc",'label_fields': ['labels']})
val_transform = A.Compose([ ToTensorV2(p=1.0)
],bbox_params = {'format':"pascal_voc","label_fields":['labels']})
`### 2.7 数据整理`"""
collate_fn默认是对数据（图片）通过torch.stack()进行简单的拼接。对于分类网络来说，默认方法是可以的（因为传入的就是数据的图片），
但是对于目标检测来说，train_dataset返回的是一个tuple，即(image, target)。
如果我们还是采用默认的合并方法，那么就会出错。
所以我们需要自定义一个方法，即collate_fn=train_dataset.collate_fn
"""
def collate_fn(batch): return tuple(zip(*batch))2.8 创建数据加载器
# 创建数据加载器 train_data = WheatDataset(train_data_df,DATAPATH+"/train",transform=train_transform)
valid_data = WheatDataset(val_data_df,DATAPATH+"/train",transform=val_transform)2.9 查看数据
# 查看一个训练集中的数据
image,target = train_data.__getitem__(0)
plt.imshow(image.numpy().transpose(1,2,0))
print(image.shape)
训练集示例
2.10 定义模型
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
num_classes = 2
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features,num_classes)2.11 定义Averager类
# 这一个类来保存对应的loss
class Averager: def __init__(self): self.current_total = 0.0 self.iterations = 0.0 def send(self, value): self.current_total += value self.iterations += 1 @property def value(self): if self.iterations == 0: return 0 else: return 1.0 * self.current_total / self.iterations def reset(self): self.current_total = 0.0 self.iterations = 0.02.12 构建训练和测试 dataloader
# 构建训练和测试 dataloader
train_dataloader = DataLoader(train_data,batch_size=BATCH_SIZE,shuffle=True,collate_fn=collate_fn)
val_dataloader = DataLoader(valid_data,batch_size=BATCH_SIZE,shuffle=False,collate_fn=collate_fn)2.13 定义模型参数
# 定义模型， 优化器，损失， 迭代，以及 学习率
train_loss = []
# val_loss = []
model = model.to(DEVICE)
params =[p for p in model.parameters() if p.requires_grad]
optimizer = optim.Adam(params,lr=LR)
loss_hist = Averager()
itr = 1
lr_scheduler=None loss_hist = Averager()
itr = 12.14 模型训练
if __name__ == '__main__': for epoch in range(EPOCHS): loss_hist.reset() for images, targets in train_dataloader: # print(images) # print(targets) # for image in images: # print(image.dtype) # torch.float32 # for t in targets: # for k, v in t.items(): # print(k ,v.dtype) images = list(image.to(DEVICE) for image in images) targets = [{k: v.to(DEVICE) for k, v in t.items()} for t in targets] loss_dict = model(images, targets) # for loss in loss_dict.values(): # print(loss.dtype) losses = sum(loss for loss in loss_dict.values()) loss_value = losses.item() loss_hist.send(loss_value) optimizer.zero_grad() losses.backward() optimizer.step() if itr % 50 == 0: print(f"Iteration #{itr} loss: {loss_value}") itr += 1 # update the learning rate if lr_scheduler is not None: lr_scheduler.step() print(f"Epoch #{epoch} loss: {loss_hist.value}")2.15 模型保存
# 模型保存
torch.save(model.state_dict(), 'fasterrcnn_resnet50_fpn.pth')
训练好的模型
2.16 加载模型进行预测
images, targets = next(iter(val_dataloader))
images = list(img.to(DEVICE) for img in images)
# print(images[0].shape)
targets = [{k: v.to(DEVICE) for k, v in t.items()} for t in targets]
boxes = targets[1]['boxes'].cpu().numpy().astype(np.int32)
sample = images[1].permute(1, 2, 0).cpu().numpy() model.eval（)
cpu_device = torch.device("cpu")
# print(images[0].shape) outputs = model(images)
outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]
# print(outputs[1]['boxes'].detach().numpy().astype(np.int32)) pred_boxes = outputs[1]['boxes'].detach().numpy().astype(np.int32) fig, ax = plt.subplots(1, 1, figsize=(16, 8)) for b, box in zip(boxes, pred_boxes): # 绘制预测边框 红色表示 cv2.rectangle(sample, (box[0], box[1]), (box[2], box[3]), (220, 0, 0), 3) # 绘制实际边框 绿色表示 cv2.rectangle(sample, (b[0], b[1]), (b[2], b[3]), (0, 220, 0), 3) ax.set_axis_off()
ax.imshow(sample)
plt.show()
检测结果
对比预测框与实际框，可以看出模型能够很好的预测出麦穗。可以尝试测试不同的麦穗图片，来进行测试查看效果。
...
五、Grad-CAM~2-Pytorch实现
Grad-CAM (Gradient-weighted Class Activation Mapping) 是一种可视化深度神经网络中哪些部分对于预测结果贡献最大的技术。它能够定位到特定的图像区域，从而使得神经网络的决策过程更加可解释和可视化。
Grad-CAM 的基本思想是，在神经网络中，最后一个卷积层的输出特征图对于分类结果的影响最大，因此我们可以通过对最后一个卷积层的梯度进行全局平均池化来计算每个通道的权重。这些权重可以用来加权特征图，生成一个 Class Activation Map (CAM)，其中每个像素都代表了该像素区域对于分类结果的重要性。
相比于传统的 CAM 方法，Grad-CAM 能够处理任意种类的神经网络，因为它不需要修改网络结构或使用特定的层结构。此外，Grad-CAM 还可以用于对特征的可视化，以及对网络中的一些特定层或单元进行分析。
在Pytorch中，我们可以使用钩子 (hook) 技术，在网络中注册前向钩子和反向钩子。前向钩子用于记录目标层的输出特征图，反向钩子用于记录目标层的梯度。在本篇文章中，我们将详细介绍如何在Pytorch中实现Grad-CAM。
加载并查看预训练的模型
为了演示Grad-CAM的实现，我将使用来自Kaggle的胸部x射线数据集和我制作的一个预训练分类器，该分类器能够将x射线分类为是否患有肺炎。
model\_path \= "your/model/path/" \# instantiate your model model \= XRayClassifier\(\) \# load your model. Here we're loading on CPU since we're not going to do \# large amounts of inference model.load\_state\_dict\(torch.load\(model\_path, map\_location\=torch.device\('cpu'\)\)\) \# put it in evaluation mode for inference model.eval\(\)首先我们看看这个模型的架构。就像前面提到的，我们需要识别最后一个卷积层，特别是它的激活函数。这一层表示模型学习到的最复杂的特征，它最有能力帮助我们理解模型的行为，下面是我们这个演示模型的代码：
import torch import torch.nn as nn import torch.nn.functional as F \# hyperparameters nc \= 3 \# number of channels nf \= 64 \# number of features to begin with dropout \= 0.2 device \= torch.device\('cuda' if torch.cuda.is\_available\(\) else 'cpu'\) \# setup a resnet block and its forward function class ResNetBlock\(nn.Module\): def \_\_init\_\_\(self, in\_channels, out\_channels, stride\=1\): super\(ResNetBlock, self\).\_\_init\_\_\(\) self.conv1 \= nn.Conv2d\(in\_channels, out\_channels, kernel\_size\=3, stride\=stride, padding\=1, bias\=False\) self.bn1 \= nn.BatchNorm2d\(out\_channels\) self.conv2 \= nn.Conv2d\(out\_channels, out\_channels, kernel\_size\=3, stride\=1, padding\=1, bias\=False\) self.bn2 \= nn.BatchNorm2d\(out\_channels\) self.shortcut \= nn.Sequential\(\) if stride \!= 1 or in\_channels \!= out\_channels: self.shortcut \= nn.Sequential\( nn.Conv2d\(in\_channels, out\_channels, kernel\_size\=1, stride\=stride, bias\=False\), nn.BatchNorm2d\(out\_channels\) \) def forward\(self, x\): out \= F.relu\(self.bn1\(self.conv1\(x\)\)\) out \= self.bn2\(self.conv2\(out\)\) out += self.shortcut\(x\) out \= F.relu\(out\) return out \# setup the final model structure class XRayClassifier\(nn.Module\): def \_\_init\_\_\(self, nc\=nc, nf\=nf, dropout\=dropout\): super\(XRayClassifier, self\).\_\_init\_\_\(\) self.resnet\_blocks \= nn.Sequential\( ResNetBlock\(nc, nf, stride\=2\), \# \(B, C, H, W\) -> \(B, NF, H/2, W/2\), i.e., \(64,64,128,128\) ResNetBlock\(nf, nf\*2, stride\=2\), \# \(64,128,64,64\) ResNetBlock\(nf\*2, nf\*4, stride\=2\), \# \(64,256,32,32\) ResNetBlock\(nf\*4, nf\*8, stride\=2\), \# \(64,512,16,16\) ResNetBlock\(nf\*8, nf\*16, stride\=2\), \# \(64,1024,8,8\) \) self.classifier \= nn.Sequential\( nn.Conv2d\(nf\*16, 1, 8, 1, 0, bias\=False\), nn.Dropout\(p\=dropout\), nn.Sigmoid\(\), \) def forward\(self, input\): output \= self.resnet\_blocks\(input.to\(device\)\) output \= self.classifier\(output\) return output模型3通道接收256x256的图片。它期望输入为[batch size, 3,256,256]。每个ResNet块以一个ReLU激活函数结束。对于我们的目标，我们需要选择最后一个ResNet块。
XRayClassifier((resnet_blocks): Sequential((0): ResNetBlock((conv1): Conv2d(3, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)(bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)(conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)(bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)(shortcut): Sequential((0): Conv2d(3, 64, kernel_size=(1, 1), stride=(2, 2), bias=False)(1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)))(1): ResNetBlock((conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)(bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)(conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)(bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)(shortcut): Sequential((0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)(1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)))(2): ResNetBlock((conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)(bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)(conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)(bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)(shortcut): Sequential((0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)(1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)))(3): ResNetBlock((conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)(bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)(conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)(bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)(shortcut): Sequential((0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)(1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)))(4): ResNetBlock((conv1): Conv2d(512, 1024, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)(bn1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)(conv2): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)(bn2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)(shortcut): Sequential((0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)(1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True))))(classifier): Sequential((0): Conv2d(1024, 1, kernel_size=(8, 8), stride=(1, 1), bias=False)(1): Dropout(p=0.2, inplace=False)(2): Sigmoid()))在Pytorch中，我们可以很容易地使用模型的属性进行选择。
model.resnet_blocks[-1]#ResNetBlock(# (conv1): Conv2d(512, 1024, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)# (bn1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)# (conv2): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)# (bn2): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)# (shortcut): Sequential(# (0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)# (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)# )#)Pytorch的钩子函数
Pytorch有许多钩子函数，这些函数可以处理在向前或后向传播期间流经模型的信息。我们可以使用它来检查中间梯度值，更改特定层的输出。在这里，我们这里将关注两个方法：
register_full_backward_hook(hook, prepend=False)
该方法在模块上注册了一个后向传播的钩子，当调用backward()方法时，钩子函数将会运行。后向钩子函数接收模块本身的输入、相对于层的输入的梯度和相对于层的输出的梯度
hook\(module, grad\_input, grad\_output\) -> tuple\(Tensor\) or None它返回一个_torch.utils.hooks.RemovableHandle_，可以使用这个返回值来删除钩子。我们在后面会讨论这个问题。**register_forward_hook(hook, *, prepend=False, with_kwargs=False)**这与前一个非常相似，它在前向传播中后运行，这个函数的参数略有不同。它可以让你访问层的输出:
hook\(module, args, output\) -> None or modified output它的返回也是_torch.utils.hooks.RemovableHandle_
向模型添加钩子函数
为了计算Grad-CAM，我们需要定义后向和前向钩子函数。这里的目标是关于最后一个卷积层的输出的梯度，需要它的激活，即层的激活函数的输出。钩子函数会在推理和向后传播期间为我们提取这些值。
\# defines two global scope variables to store our gradients and activations gradients \= None activations \= None def backward\_hook\(module, grad\_input, grad\_output\): global gradients \# refers to the variable in the global scope print\('Backward hook running...'\) gradients \= grad\_output \# In this case, we expect it to be torch.Size\(\[batch size, 1024, 8, 8\]\) print\(f'Gradients size: \{gradients\[0\].size\(\)\}'\) \# We need the 0 index because the tensor containing the gradients comes \# inside a one element tuple. def forward\_hook\(module, args, output\): global activations \# refers to the variable in the global scope print\('Forward hook running...'\) activations \= output \# In this case, we expect it to be torch.Size\(\[batch size, 1024, 8, 8\]\) print\(f'Activations size: \{activations.size\(\)\}'\)在定义了钩子函数和存储激活和梯度的变量之后，就可以在感兴趣的层中注册钩子，注册的代码如下：
backward\_hook \= model.resnet\_blocks\[\-1\].register\_full\_backward\_hook\(backward\_hook, prepend\=False\) forward\_hook \= model.resnet\_blocks\[\-1\].register\_forward\_hook\(forward\_hook, prepend\=False\)检索需要的梯度和激活
现在已经为模型设置了钩子函数，让我们加载一个图像，计算gradcam。
from PIL import Image img\_path \= "/your/image/path/" image \= Image.open\(img\_path\).convert\('RGB'\)
为了进行推理，我们还需要对其进行预处理：
from torchvision import transforms from torchvision.transforms import ToTensor image\_size \= 256 transform \= transforms.Compose\(\[ transforms.Resize\(image\_size, antialias\=True\), transforms.CenterCrop\(image\_size\), transforms.ToTensor\(\), transforms.Normalize\(\(0.5, 0.5, 0.5\), \(0.5, 0.5, 0.5\)\), \]\) img\_tensor \= transform\(image\) \# stores the tensor that represents the image现在就可以进行前向传播了：
model\(img\_tensor.unsqueeze\(0\)\).backward\(\)钩子函数的返回如下:
Forward hook running... Activations size: torch.Size\(\[1, 1024, 8, 8\]\) Backward hook running... Gradients size: torch.Size\(\[1, 1024, 8, 8\]\)得到了梯度和激活变量后就可以生成热图：
计算Grad-CAM
为了计算Grad-CAM，我们将原始论文公式进行一些简单的修改：
import torch.nn.functional as F import matplotlib.pyplot as plt \# weight the channels by corresponding gradients for i in range\(activations.size\(\)\[1\]\): activations\[:, i, :, :\] \*= pooled\_gradients\[i\] \# average the channels of the activations heatmap \= torch.mean\(activations, dim\=1\).squeeze\(\) \# relu on top of the heatmap heatmap \= F.relu\(heatmap\) \# normalize the heatmap heatmap /= torch.max\(heatmap\) \# draw the heatmap plt.matshow\(heatmap.detach\(\)\)结果如下：
得到的激活包含1024个特征映射，这些特征映射捕获输入图像的不同方面，每个方面的空间分辨率为8x8。通过钩子获得的梯度表示每个特征映射对最终预测的重要性。通过计算梯度和激活的元素积可以获得突出显示图像最相关部分的特征映射的加权和。通过计算加权特征图的全局平均值，可以得到一个单一的热图，该热图表明图像中对模型预测最重要的区域。这就是Grad-CAM，它提供了模型决策过程的可视化解释，可以帮助我们解释和调试模型的行为。
但是这个图能代表什么呢？我们将他与图片进行整合就能更加清晰的可视化了。
结合原始图像和热图
下面的代码将原始图像和我们生成的热图进行整合显示：
from torchvision.transforms.functional import to_pil_imagefrom matplotlib import colormapsimport numpy as npimport PIL# Create a figure and plot the first imagefig, ax = plt.subplots()ax.axis('off') # removes the axis markers# First plot the original imageax.imshow(to_pil_image(img_tensor, mode='RGB'))# Resize the heatmap to the same size as the input image and defines# a resample algorithm for increasing image resolution# we need heatmap.detach() because it can't be converted to numpy array while# requiring gradientsoverlay = to_pil_image(heatmap.detach(), mode='F').resize((256,256), resample=PIL.Image.BICUBIC)# Apply any colormap you wantcmap = colormaps['jet']overlay = (255 * cmap(np.asarray(overlay) ** 2)[:, :, :3]).astype(np.uint8)# Plot the heatmap on the same axes,# but with alpha < 1 (this defines the transparency of the heatmap)ax.imshow(overlay, alpha=0.4, interpolation='nearest', extent=extent)# Show the plotplt.show()
这样看是不是就理解多了。由于它是一个正常的x射线结果，所以并没有什么需要特殊说明的。
再看这个例子，这个结果中被标注的是肺炎。Grad-CAM能准确显示出医生为确定是否患有肺炎而必须检查的胸部x光片区域。也就是说我们的模型的确学到了一些东西（红色区域再肺部附近）
删除钩子
要从模型中删除钩子，只需要在返回句柄中调用remove()方法。
backward\_hook.remove\(\) forward\_hook.remove\(\)总结
这篇文章可以帮助你理清Grad-CAM 是如何工作的，以及如何用Pytorch实现它。因为Pytorch包含了强大的钩子函数，所以我们可以在任何模型中使用本文的代码。
Grad-CAM论文：https://arxiv.org/pdf/1610.02391.pdf
本文的数据集：https://www.kaggle.com/datasets/paultimothymooney/chest-xray-pneumonia
...
六、Pytorch搭建Transformer~2
这里是参考了如下代码后对Transformer进行的尽量简洁实现，适合初学者把握模型核心要点
- https://wmathor.com/index.php/archives/1455/
- http://nlp.seas.harvard.edu/annotated-transformer/ (哈佛NLP团队实现版本)
完整代码已发布在github上（包含jupyter实现、.py实现、本文绘制的图还有训练好的模型等），欢迎star：https://github.com/BoXiaolei/MyTransformer_pytorch
1. 数据预处理
import math
import torch
import numpy as np
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as Data
# 姑且把导包也放在这个地方吧# S: 起始标记
# E: 结束标记
# P：意为padding，将当前序列补齐至最长序列长度的占位符
sentence = [ # enc_input dec_input dec_output ['ich mochte ein bier P','S i want a beer .', 'i want a beer . E'], ['ich mochte ein cola P','S i want a coke .', 'i want a coke . E'],
]# 词典，padding用0来表示
# 源词典
src_vocab = {'P':0, 'ich':1,'mochte':2,'ein':3,'bier':4,'cola':5}
src_vocab_size = len(src_vocab) # 6
# 目标词典（包含特殊符）
tgt_vocab = {'P':0,'i':1,'want':2,'a':3,'beer':4,'coke': 5,'S':6,'E':7,'.':8}
# 反向映射词典，idx ——> word
idx2word = {v:k for k,v in tgt_vocab.items()}
tgt_vocab_size = len(tgt_vocab) # 9 src_len = 5 # 输入序列enc_input的最长序列长度，其实就是最长的那句话的token数
tgt_len = 6 # 输出序列dec_input/dec_output的最长序列长度# 构建模型输入的Tensor
def make_data(sentence): enc_inputs, dec_inputs, dec_outputs = [],[],[] for i in range(len(sentence)): enc_input = [src_vocab[word] for word in sentence[i][0].split()] dec_input = [tgt_vocab[word] for word in sentence[i][1].split()] dec_output = [tgt_vocab[word] for word in sentence[i][2].split()] enc_inputs.append(enc_input) dec_inputs.append(dec_input) dec_outputs.append(dec_output)# LongTensor是专用于存储整型的，Tensor则可以存浮点、整数、bool等多种类型 return torch.LongTensor(enc_inputs),torch.LongTensor(dec_inputs),torch.LongTensor(dec_outputs)enc_inputs, dec_inputs, dec_outputs = make_data(sentence)
print(' enc_inputs: \n', enc_inputs) # enc_inputs: [2,5]
print(' dec_inputs: \n', dec_inputs) # dec_inputs: [2,6]
print(' dec_outputs: \n', dec_outputs) # dec_outputs: [2,6]
enc_inputs:
tensor([[1, 2, 3, 4, 0],
[1, 2, 3, 5, 0]])
dec_inputs:
tensor([[6, 1, 2, 3, 4, 8],
[6, 1, 2, 3, 5, 8]])
dec_outputs:
tensor([[1, 2, 3, 4, 8, 7],
[1, 2, 3, 5, 8, 7]])
# 使用Dataset加载数据
class MyDataSet(Data.Dataset): def __init__(self,enc_inputs, dec_inputs, dec_outputs): super(MyDataSet,self).__init__() self.enc_inputs = enc_inputs self.dec_inputs = dec_inputs self.dec_outputs = dec_outputsdef __len__(self): # 我们前面的enc_inputs.shape = [2,5],所以这个返回的是2 return self.enc_inputs.shape[0] # 根据idx返回的是一组 enc_input, dec_input, dec_output def __getitem__(self, idx): return self.enc_inputs[idx], self.dec_inputs[idx], self.dec_outputs[idx] # 构建DataLoader
loader = Data.DataLoader(dataset=MyDataSet(enc_inputs,dec_inputs, dec_outputs),batch_size=2,shuffle=True)
2. 模型参数
# 用来表示一个词的向量长度
d_model = 512# FFN的隐藏层神经元个数
d_ff = 2048# 分头后的q、k、v词向量长度，依照原文我们都设为64
# 原文：queries and kes of dimention d_k,and values of dimension d_v .所以q和k的长度都用d_k来表示
d_k = d_v = 64 # Encoder Layer 和 Decoder Layer的个数
n_layers = 6# 多头注意力中head的个数，原文：we employ h = 8 parallel attention layers, or heads
n_heads = 8- Transformer包含Encoder和Decoder
- Encoder和Decoder各自包含6个Layer
- Encoder Layer中包含 Self Attention 和 FFN 两个Sub Layer
- Decoder Layer中包含 Masked Self Attention、 Cross Attention、 FFN 三个Sub Layer
布局如图：
3. Positional Encoding
用于为输入的词向量进行位置编码
原文：The positional encodings have the same dimension d_model as the embeddings, so that the two can be summed
lass PositionalEncoding(nn.Module): def __init__(self, d_model, dropout=0.1, max_len=5000): # dropout是原文的0.1，max_len原文没找到 '''max_len是假设的一个句子最多包含5000个token''' super(PositionalEncoding, self).__init__() self.dropout = nn.Dropout(p=dropout) # 开始位置编码部分,先生成一个max_len * d_model 的矩阵，即5000 * 512 # 5000是一个句子中最多的token数，512是一个token用多长的向量来表示，5000*512这个矩阵用于表示一个句子的信息 pe = torch.zeros(max_len, d_model) pos = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) # pos：[max_len,1],即[5000,1] # 先把括号内的分式求出来,pos是[5000,1],分母是[256],通过广播机制相乘后是[5000,256] div_term = pos / pow(10000.0,torch.arange(0, d_model, 2).float() / d_model) # 再取正余弦 pe[:, 0::2] = torch.sin(div_term) pe[:, 1::2] = torch.cos(div_term) # 一个句子要做一次pe，一个batch中会有多个句子，所以增加一维用来和输入的一个batch的数据相加时做广播 pe = pe.unsqueeze(0) # [5000,512] -> [1,5000,512] # 将pe作为固定参数保存到缓冲区，不会被更新 self.register_buffer('pe', pe)def forward(self, x): '''x: [batch_size, seq_len, d_model]''' # 5000是我们预定义的最大的seq_len，就是说我们把最多的情况pe都算好了，用的时候用多少就取多少 x = x + self.pe[:, :x.size(1), :] return self.dropout(x) # return: [batch_size, seq_len, d_model], 和输入的形状相同
4. Pad Mask
- 首先我们要清楚，这是一个计算mask的函数，它的返回是一个布尔矩阵，为True的位置是需要被mask掉的，False的位置是不需要动的
- 其次这个函数是理解Transformer代码中非常重要的一环，因为我们输入模型的句子有长有短，我们用占位符P统一补足成了最长的那个句子的长度，而这些占位符是没有意义的，不能让他们吸收到query的注意力，因此我们要把这些位置设为True
- 这个计算出的mask在何时被使用呢？
- 在query和key的转置相乘得出（len_q,len_k）这个注意力分数矩阵以后，将使用本函数得到的mask来掩盖相乘结果矩阵
- 原来的相乘结果矩阵（len_q,len_k）中第 i 行第 j 列的意义是“作为q的序列中第i个词对作为k的序列中第j个词的注意力分数”，而第 i 整行就是q中这个词对k中所有词的注意力，第 j 整列就是q中所有词对k中第j个词的注意力分数，作为padding，q中的所有词都不应该注意它，因此对应列均需设为True
- 为什么只有k的padding位被mask了，q的padding位为什么没被mask？（即此函数的返回矩阵为什么只有最后几列是True，最后几行不应该也是True么）
- 按理来说是这样的，作为padding不该被别人注意，同时它也不该注意别人，计算出的padding对其他词的注意力也是无意义的，我们这里其实是偷了个懒，但这是因为：q中的padding对k中的词的注意力我们是不会用到的，因为我们不会用一个padding字符去预测下一个词，并且它的向量表示不管怎么更新都不会影响到别的q中别的词的计算，所以我们就放任自流了。但k中的padding不一样，如果不管它将无意义的吸收掉大量q中词汇的注意力，使得模型的学习出现偏差。
- p.s. 哈佛团队也是这么实现的
# 为enc_input和dec_input做一个mask，把占位符P的token（就是0） mask掉
# 返回一个[batch_size, len_q, len_k]大小的布尔张量，True是需要mask掉的位置
def get_attn_pad_mask(seq_q, seq_k): batch_size, len_q = seq_q.size() batch_size, len_k = seq_k.size() # seq_k.data.eq(0)返回一个等大的布尔张量，seq_k元素等于0的位置为True,否则为False # 然后扩维以保证后续操作的兼容(广播) pad_attn_mask = seq_k.data.eq(0).unsqueeze(1) # pad_attn_mask: [batch_size,1,len_k] # 要为每一个q提供一份k，所以把第二维度扩展了q次 # 另注意expand并非真正加倍了内存，只是重复了引用，对任意引用的修改都会修改原始值 # 这里是因为我们不会修改这个mask所以用它来节省内存 return pad_attn_mask.expand(batch_size, len_q, len_k) # return: [batch_size, len_q, len_k] # 返回的是batch_size个 len_q * len_k的矩阵，内容是True和False， # 第i行第j列表示的是query的第i个词对key的第j个词的注意力是否无意义，若无意义则为True，有意义的为False（即被padding的位置是True）5. Subsequence Mask
此函数对应Transformer模型架构中Decoder的第一个注意力“Masked Multi-Head self Attention”中的Masked一词，作用是防止模型看到未来时刻的输入
- 怎么理解呢？
- 其实这需要结合实际使用模型进行预测时的数据流动来理解（可以参考本文最后一部分的测试流程图），因为实际的预测中解码器Decoder其实是一个词一个词累积着的输出的(每次比上回多一个)，每输出一个词就会把这个词拼到新的dec_input中，然后再预测下一个词，直到输出终止标记，这是测试过程。而我们为了让模型学到最精确的表示，每一回我们喂给decoder的都是“正确答案”，但是我们的模型必须依照真实的使用场景来训练，所以我们需要让decoder只看到当前时刻以前的输出（即正确答案不能多看，最多看到当前这个题的就行了），而实现这一目的的方法就是：屏蔽掉当前时刻以后的注意力分数，而这个时间上“以后”的概念，现现在数据上就是序列中token的前后顺序，所以当前token之后的都需要被mask
# 用于获取对后续位置的掩码，防止在预测过程中看到未来时刻的输入
# 原文：to prevent positions from attending to subsequent positions
def get_attn_subsequence_mask(seq): """seq: [batch_size, tgt_len]""" # batch_size个 tgt_len * tgt_len的mask矩阵 attn_shape = [seq.size(0), seq.size(1), seq.size(1)]# np.triu 是生成一个 upper triangular matrix 上三角矩阵，k是相对于主对角线的偏移量 # k=1意为不包含主对角线（从主对角线向上偏移1开始） subsequence_mask = np.triu(np.ones(attn_shape), k=1) subsequence_mask = torch.from_numpy(subsequence_mask).byte() # 因为只有0、1所以用byte节省内存 return subsequence_mask # return: [batch_size, tgt_len, tgt_len]因为这个mask只用于解码器中的第一个self attention，q和k都是自己（dec_input），所以是一个方阵
6. ScaledDotProductAttention
此函数用于计算缩放点积注意力，在MultiHeadAttention中被调用
*class ScaledDotProductionAttention(nn.Module):
def init(self):
super(ScaledDotProductionAttention, self).init()
def forward(self, Q, K, V, attn_mask): ''' Q: [batch_size, n_heads, len_q, d_k] K: [batch_size, n_heads, len_k, d_k] V: [batch_size, n_heads, len_v(=len_k), d_v] 全文两处用到注意力，一处是self attention，另一处是co attention，前者不必说，后者的k和v都是encoder的输出，所以k和v的形状总是相同的 attn_mask: [batch_size, n_heads, seq_len, seq_len] ''' # 1) 计算注意力分数QK^T/sqrt(d_k) scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k) # scores: [batch_size, n_heads, len_q, len_k] # 2) 进行 mask 和 softmax # mask为True的位置会被设为-1e9 scores.masked_fill_(attn_mask, -1e9) attn = nn.Softmax(dim=-1)(scores) # attn: [batch_size, n_heads, len_q, len_k] # 3) 乘V得到最终的加权和 context = torch.matmul(attn, V) # context: [batch_size, n_heads, len_q, d_v] ''' 得出的context是每个维度(d_1-d_v)都考虑了在当前维度(这一列)当前token对所有token的注意力后更新的新的值， 换言之每个维度d是相互独立的，每个维度考虑自己的所有token的注意力，所以可以理解成1列扩展到多列 返回的context: [batch_size, n_heads, len_q, d_v]本质上还是batch_size个句子， 只不过每个句子中词向量维度512被分成了8个部分，分别由8个头各自看一部分，每个头算的是整个句子(一列)的512/8=64个维度，最后按列拼接起来 ''' return context # context: [batch_size, n_heads, len_q, d_v]*7. MultiHeadAttention
多头注意力的实现，Transformer的核心
class MultiHeadAttention(nn.Module): def __init__(self): super(MultiHeadAttention, self).__init__() self.W_Q = nn.Linear(d_model, d_model) self.W_K = nn.Linear(d_model, d_model) self.W_V = nn.Linear(d_model, d_model) self.concat = nn.Linear(d_model, d_model)def forward(self, input_Q, input_K, input_V, attn_mask): ''' input_Q: [batch_size, len_q, d_model] len_q是作为query的句子的长度，比如enc_inputs（2,5,512）作为输入，那句子长度5就是len_q input_K: [batch_size, len_k, d_model] input_K: [batch_size, len_v(len_k), d_model] attn_mask: [batch_size, seq_len, seq_len]'''residual, batch_size = input_Q, input_Q.size(0)# 1）linear projection [batch_size, seq_len, d_model] -> [batch_size, n_heads, seq_len, d_k/d_v] Q = self.W_Q(input_Q).view(batch_size, -1, n_heads, d_k).transpose(1, 2) # Q: [batch_size, n_heads, len_q, d_k] K = self.W_K(input_K).view(batch_size, -1, n_heads, d_k).transpose(1, 2) # K: [batch_size, n_heads, len_k, d_k] V = self.W_V(input_V).view(batch_size, -1, n_heads, d_v).transpose(1, 2) # V: [batch_size, n_heads, len_v(=len_k), d_v]# 2）计算注意力 # 自我复制n_heads次，为每个头准备一份mask attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1) # attn_mask: [batch_size, n_heads, seq_len, seq_len] context = ScaledDotProductionAttention()(Q, K, V, attn_mask) # context: [batch_size, n_heads, len_q, d_v]# 3）concat部分 context = torch.cat([context[:,i,:,:] for i in range(context.size(1))], dim=-1) output = self.concat(context) # [batch_size, len_q, d_model] return nn.LayerNorm(d_model).cuda()(output + residual) # output: [batch_size, len_q, d_model]''' 最后的concat部分，网上的大部分实现都采用的是下面这种方式（也是哈佛NLP团队的写法） context = context.transpose(1, 2).reshape(batch_size, -1, d_model) output = self.linear(context) 但是我认为这种方式拼回去会使原来的位置乱序，于是并未采用这种写法，两种写法最终的实验结果是相近的 '''
8. FeedForward Networks
这部分代码很简单，对应模型图中的 Feed Forward和 Add & Norm
class PositionwiseFeedForward(nn.Module): def __init__(self): super(PositionwiseFeedForward, self).__init__() # 就是一个MLP self.fc = nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model) )def forward(self, inputs): '''inputs: [batch_size, seq_len, d_model]''' residual = inputs output = self.fc(inputs) return nn.LayerNorm(d_model).cuda()(output + residual) # return： [batch_size, seq_len, d_model] 形状不变9. Encoder Layer
包含一个MultiHeadAttention和一个FFN
class EncoderLayer(nn.Module): def __init__(self): super(EncoderLayer, self).__init__() self.enc_self_attn = MultiHeadAttention() self.pos_ffn = PositionwiseFeedForward()def forward(self, enc_inputs, enc_self_attn_mask): ''' enc_inputs: [batch_size, src_len, d_model] enc_self_attn_mask: [batch_size, src_len, src_len] ''' # Q、K、V均为 enc_inputs enc_ouputs = self.enc_self_attn(enc_inputs, enc_inputs, enc_inputs, enc_self_attn_mask) # enc_ouputs: [batch_size, src_len, d_model] enc_ouputs = self.pos_ffn(enc_ouputs) # enc_outputs: [batch_size, src_len, d_model] return enc_ouputs # enc_outputs: [batch_size, src_len, d_model]10. Encoder
包含一个源序列词向量嵌入nn.Embedding、一个位置编码PositionalEncoding和6个Encoder Layer
class Encoder(nn.Module): def __init__(self): super(Encoder, self).__init__() # 直接调的现成接口完成词向量的编码，输入是类别数和每一个类别要映射成的向量长度 self.src_emb = nn.Embedding(src_vocab_size, d_model) self.pos_emb = PositionalEncoding(d_model) self.layers = nn.ModuleList([EncoderLayer() for _ in range(n_layers)])def forward(self, enc_inputs): '''enc_inputs: [batch_size, src_len]''' enc_outputs = self.src_emb(enc_inputs) # [batch_size, src_len] -> [batch_size, src_len, d_model] enc_outputs = self.pos_emb(enc_outputs) # enc_outputs: [batch_size, src_len, d_model] # Encoder中是self attention，所以传入的Q、K都是enc_inputs enc_self_attn_mask = get_attn_pad_mask(enc_inputs, enc_inputs) # enc_self_attn_mask: [batch_size, src_len, src_len] for layer in self.layers: enc_outputs = layer(enc_outputs, enc_self_attn_mask) return enc_outputs # enc_outputs: [batch_size, src_len, d_model]11. Decoder Layer
包含两个MultiHeadAttention和一个FFN
class DecoderLayer(nn.Module): def __init__(self): super(DecoderLayer, self).__init__() self.dec_self_attn = MultiHeadAttention() self.dec_enc_attn = MultiHeadAttention() self.pos_ffn = PositionwiseFeedForward() def forward(self, dec_inputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask): ''' dec_inputs: [batch_size, tgt_len, d_model] enc_outputs: [batch_size, src_len, d_model] dec_self_attn_mask: [batch_size, tgt_len, tgt_len] dec_enc_attn_mask: [batch_size, tgt_len, src_len] 前者是Q后者是K ''' dec_outputs = self.dec_self_attn(dec_inputs, dec_inputs, dec_inputs, dec_self_attn_mask) dec_outputs = self.dec_enc_attn(dec_outputs, enc_outputs, enc_outputs, dec_enc_attn_mask) dec_outputs = self.pos_ffn(dec_outputs)return dec_outputs # dec_outputs: [batch_size, tgt_len, d_model]12. Decoder
包含一个目标序列词向量序列嵌入nn.Embedding、一个位置编码PositionalEncoding还有6个Decoder Layer
class Decoder(nn.Module): def __init__(self): super(Decoder, self).__init__() self.tgt_emb = nn.Embedding(tgt_vocab_size, d_model) self.pos_emb = PositionalEncoding(d_model) self.layers = nn.ModuleList([DecoderLayer() for _ in range(n_layers)])def forward(self, dec_inputs, enc_inputs, enc_outputs): ''' 这三个参数对应的不是Q、K、V，dec_inputs是Q，enc_outputs是K和V，enc_inputs是用来计算padding mask的 dec_inputs: [batch_size, tgt_len] enc_inpus: [batch_size, src_len] enc_outputs: [batch_size, src_len, d_model] ''' dec_outputs = self.tgt_emb(dec_inputs) dec_outputs = self.pos_emb(dec_outputs).cuda() dec_self_attn_pad_mask = get_attn_pad_mask(dec_inputs, dec_inputs).cuda() dec_self_attn_subsequence_mask = get_attn_subsequence_mask(dec_inputs).cuda() # 将两个mask叠加，布尔值可以视为0和1，和大于0的位置是需要被mask掉的，赋为True，和为0的位置是有意义的为False dec_self_attn_mask = torch.gt((dec_self_attn_pad_mask +dec_self_attn_subsequence_mask), 0).cuda() # 这是co-attention部分，为啥传入的是enc_inputs而不是enc_outputs呢 dec_enc_attn_mask = get_attn_pad_mask(dec_inputs, enc_inputs)for layer in self.layers: dec_outputs = layer(dec_outputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask) return dec_outputs # dec_outputs: [batch_size, tgt_len, d_model]13. Transformer
包含一个Encoder、一个Decoder、一个nn.Linear
class Transformer(nn.Module): def __init__(self): super(Transformer, self).__init__() self.encoder = Encoder().cuda() self.decoder = Decoder().cuda() self.projection = nn.Linear(d_model, tgt_vocab_size).cuda()def forward(self, enc_inputs, dec_inputs): ''' enc_inputs: [batch_size, src_len] dec_inputs: [batch_size, tgt_len] ''' enc_outputs = self.encoder(enc_inputs) dec_outputs = self.decoder(dec_inputs, enc_inputs, enc_outputs) dec_logits = self.projection(dec_outputs) # dec_logits: [batch_size, tgt_len, tgt_vocab_size] # 解散batch，一个batch中有batch_size个句子，每个句子有tgt_len个词（即tgt_len行）， # 现在让他们按行依次排布，如前tgt_len行是第一个句子的每个词的预测概率， # 再往下tgt_len行是第二个句子的，一直到batch_size * tgt_len行 return dec_logits.view(-1, dec_logits.size(-1)) # [batch_size * tgt_len, tgt_vocab_size] '''最后变形的原因是：nn.CrossEntropyLoss接收的输入的第二个维度必须是类别'''14. 训练
model = Transformer().cuda()
model.train()
# 损失函数,忽略为0的类别不对其计算loss（因为是padding无意义）
criterion = nn.CrossEntropyLoss(ignore_index=0)
optimizer = optim.SGD(model.parameters(), lr=1e-3, momentum=0.99)# 训练开始
for epoch in range(1000): for enc_inputs, dec_inputs, dec_outputs in loader: ''' enc_inputs: [batch_size, src_len] [2,5] dec_inputs: [batch_size, tgt_len] [2,6] dec_outputs: [batch_size, tgt_len] [2,6] ''' enc_inputs, dec_inputs, dec_outputs = enc_inputs.cuda(), dec_inputs.cuda(), dec_outputs.cuda() outputs = model(enc_inputs, dec_inputs) # outputs: [batch_size * tgt_len, tgt_vocab_size] # outputs: [batch_size * tgt_len, tgt_vocab_size], dec_outputs: [batch_size, tgt_len] loss = criterion(outputs, dec_outputs.view(-1)) # 将dec_outputs展平成一维张量# 更新权重 optimizer.zero_grad() loss.backward() optimizer.step()print(f'Epoch [{epoch + 1}/1000], Loss: {loss.item()}') torch.save(model, 'MyTransformer.pth')Epoch [1/1000], Loss: 2.1738522052764893
Epoch [2/1000], Loss: 2.0738959312438965
Epoch [3/1000], Loss: 1.9049569368362427
Epoch [4/1000], Loss: 1.5952812433242798
Epoch [5/1000], Loss: 1.3922237157821655
Epoch [6/1000], Loss: 1.2176579236984253
Epoch [7/1000], Loss: 0.9465160965919495
Epoch [8/1000], Loss: 0.7412980198860168
Epoch [9/1000], Loss: 0.5801363587379456
Epoch [10/1000], Loss: 0.3842979967594147
...... Epoch [990/1000], Loss: 3.3477856504759984e-06
Epoch [991/1000], Loss: 3.576270273697446e-06
Epoch [992/1000], Loss: 3.675609832498594e-06
Epoch [993/1000], Loss: 5.632609827443957e-06
Epoch [994/1000], Loss: 4.202114723739214e-06
Epoch [995/1000], Loss: 3.91402772947913e-06
Epoch [996/1000], Loss: 3.566336090443656e-06
Epoch [997/1000], Loss: 3.2285781799146207e-06
Epoch [998/1000], Loss: 4.897496637568111e-06
Epoch [999/1000], Loss: 3.894158908224199e-06
Epoch [1000/1000], Loss: 3.665677240860532e-06
15. 测试
# 原文使用的是大小为4的beam search，这里为简单起见使用更简单的greedy贪心策略生成预测，不考虑候选，每一步选择概率最大的作为输出
# 如果不使用greedy_decoder，那么我们之前实现的model只会进行一次预测得到['i']并不会自回归，所以我们利用编写好的Encoder-Decoder来手动实现自回归（把上一次Decoder的输出作为下一次的输入，直到预测出终止符）
def greedy_decoder(model, enc_input, start_symbol): """enc_input: [1, seq_len] 对应一句话""" enc_outputs = model.encoder(enc_input) # enc_outputs: [1, seq_len, 512] # 生成一个1行0列的，和enc_inputs.data类型相同的空张量，待后续填充 dec_input = torch.zeros(1, 0).type_as(enc_input.data) # .data避免影响梯度信息 next_symbol = start_symbol flag = True while flag: # dec_input.detach() 创建 dec_input 的一个分离副本 # 生成了一个 只含有next_symbol的（1,1）的张量 # -1 表示在最后一个维度上进行拼接cat # 这行代码的作用是将next_symbol拼接到dec_input中，作为新一轮decoder的输入 dec_input = torch.cat([dec_input.detach(), torch.tensor([[next_symbol]], dtype=enc_input.dtype).cuda()], -1) # dec_input: [1,当前词数] dec_outputs = model.decoder(dec_input, enc_input, enc_outputs) # dec_outputs: [1, tgt_len, d_model] projected = model.projection(dec_outputs) # projected: [1, 当前生成的tgt_len, tgt_vocab_size] # max返回的是一个元组（最大值，最大值对应的索引），所以用[1]取到最大值对应的索引, 索引就是类别，即预测出的下一个词 # keepdim为False会导致减少一维 prob = projected.squeeze(0).max(dim=-1, keepdim=False)[1] # prob: [1], # prob是一个一维的列表，包含目前为止依次生成的词的索引，最后一个是新生成的（即下一个词的类别） # 因为注意力是依照前面的词算出来的，所以后生成的不会改变之前生成的 next_symbol = prob.data[-1] if next_symbol == tgt_vocab['.']: flag = False print(next_symbol) return dec_input # dec_input: [1,tgt_len]# 测试
model = torch.load('MyTransformer.pth')
model.eval（)
with torch.no_grad(): # 手动从loader中取一个batch的数据 enc_inputs, _, _ = next(iter(loader)) enc_inputs = enc_inputs.cuda() for i in range(len(enc_inputs)): greedy_dec_input = greedy_decoder(model, enc_inputs[i].view(1, -1), start_symbol=tgt_vocab['S']) predict = model(enc_inputs[i].view(1, -1), greedy_dec_input) # predict: [batch_size * tgt_len, tgt_vocab_size] predict = predict.data.max(dim=-1, keepdim=False)[1] '''greedy_dec_input是基于贪婪策略生成的，而贪婪解码的输出是基于当前时间步生成的假设的输出。这意味着它可能不是最优的输出，因为它仅考虑了每个时间步的最有可能的单词，而没有考虑全局上下文。 因此，为了获得更好的性能评估，通常会将整个输入序列和之前的假设输出序列传递给模型，以考虑全局上下文并允许模型更准确地生成输出 ''' print(enc_inputs[i], '->', [idx2word[n.item()] for n in predict])tensor(1, device='cuda:0')
tensor(2, device='cuda:0')
tensor(3, device='cuda:0')
tensor(5, device='cuda:0')
tensor(8, device='cuda:0')
tensor([1, 2, 3, 5, 0], device='cuda:0') -> ['i', 'want', 'a', 'coke', '.']
tensor(1, device='cuda:0')
tensor(2, device='cuda:0')
tensor(3, device='cuda:0')
tensor(4, device='cuda:0')
tensor(8, device='cuda:0')
tensor([1, 2, 3, 4, 0], device='cuda:0') -> ['i', 'want', 'a', 'beer', '.']
附1：对注意力分头的探究
# 探究一下多头注意力从(batch_size, seq_len, d_model) 到 (batch_size,n_heads, seq_len, d_k/v)的意义# 1、这是初始的q
q = torch.arange(120).reshape(2,5,12)
print(q)
print('------------------')
batch_size = 2
seq_len = 5
d_model = 12
n_heads = 3
d_k = 4# 2、分成n_heads个头
new_q = q.view(batch_size, -1, n_heads, d_k).transpose(1,2)
# 上面一行代码的形状变化：(2,5,12) -> (2,5,3,4) -> (2,3,5,4)
# 意义变化：最初是batch_size为2，一个batch中有2个句子，一个句子包含5个词，每个词由长度为12的向量表示
# 最后仍然是batch_size为2，但一个batch中有3个头，每个头包含一个句子，每个句子包含5个词，但每个词由长度为4的向量表示print(new_q)
print(new_q.shape) # torch.Size([2, 3, 5, 4])
print('------------------')# 3、将n_heads个头合并
final_q = q.transpose(1,2).contiguous().view(batch_size, -1, d_model)
print(final_q)
print(final_q.shape)
print('------------------')# 按原来的concat实现拼回去元素顺序和最初不同了，因此改成下面这种实现
final_q2 = torch.cat([new_q[:,i,:,:] for i in range(new_q.size(1))], dim=-1)
print(final_q2)
print(final_q2.shape)tensor([[[ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
[ 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23],
[ 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35],
[ 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47],
[ 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59]],[[ 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71],
[ 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83],
[ 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95],
[ 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107],
[108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119]]])
tensor([[[[ 0, 1, 2, 3],
[ 12, 13, 14, 15],
[ 24, 25, 26, 27],
[ 36, 37, 38, 39],
[ 48, 49, 50, 51]],[[ 4, 5, 6, 7],
[ 16, 17, 18, 19],
[ 28, 29, 30, 31],
[ 40, 41, 42, 43],
[ 52, 53, 54, 55]],[[ 8, 9, 10, 11],
[ 20, 21, 22, 23],
[ 32, 33, 34, 35],
[ 44, 45, 46, 47],
[ 56, 57, 58, 59]]],[[[ 60, 61, 62, 63],
[ 72, 73, 74, 75],
[ 84, 85, 86, 87],
[ 96, 97, 98, 99],
[108, 109, 110, 111]],[[ 64, 65, 66, 67],
[ 76, 77, 78, 79],
[ 88, 89, 90, 91],
[100, 101, 102, 103],
[112, 113, 114, 115]],[[ 68, 69, 70, 71],
[ 80, 81, 82, 83],
[ 92, 93, 94, 95],
[104, 105, 106, 107],
[116, 117, 118, 119]]]])
torch.Size([2, 3, 5, 4])
tensor([[[ 0, 12, 24, 36, 48, 1, 13, 25, 37, 49, 2, 14],
[ 26, 38, 50, 3, 15, 27, 39, 51, 4, 16, 28, 40],
[ 52, 5, 17, 29, 41, 53, 6, 18, 30, 42, 54, 7],
[ 19, 31, 43, 55, 8, 20, 32, 44, 56, 9, 21, 33],
[ 45, 57, 10, 22, 34, 46, 58, 11, 23, 35, 47, 59]],[[ 60, 72, 84, 96, 108, 61, 73, 85, 97, 109, 62, 74],
[ 86, 98, 110, 63, 75, 87, 99, 111, 64, 76, 88, 100],
[112, 65, 77, 89, 101, 113, 66, 78, 90, 102, 114, 67],
[ 79, 91, 103, 115, 68, 80, 92, 104, 116, 69, 81, 93],
[105, 117, 70, 82, 94, 106, 118, 71, 83, 95, 107, 119]]])
torch.Size([2, 5, 12])
tensor([[[ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
[ 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23],
[ 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35],
[ 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47],
[ 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59]],[[ 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71],
[ 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83],
[ 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95],
[ 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107],
[108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119]]])
torch.Size([2, 5, 12])
附2：模型中的数据流动
...
七、
...
八、
...