diff --git "a/docs/academic/\347\256\227\346\263\225\347\247\221\346\231\256/Transformer/\345\244\247\346\250\241\345\236\213KV Cache\350\212\202\347\234\201\347\245\236\345\231\250MLA\345\255\246\344\271\240\347\254\224\350\256\260\357\274\210\345\214\205\345\220\253\346\216\250\347\220\206\346\227\266\347\232\204\347\237\251\351\230\265\345\220\270\346\224\266\345\210\206\346\236\220\357\274\211.md" "b/docs/academic/\347\256\227\346\263\225\347\247\221\346\231\256/Transformer/\345\244\247\346\250\241\345\236\213KV Cache\350\212\202\347\234\201\347\245\236\345\231\250MLA\345\255\246\344\271\240\347\254\224\350\256\260\357\274\210\345\214\205\345\220\253\346\216\250\347\220\206\346\227\266\347\232\204\347\237\251\351\230\265\345\220\270\346\224\266\345\210\206\346\236\220\357\274\211.md" new file mode 100644 index 0000000..cced0e1 --- /dev/null +++ "b/docs/academic/\347\256\227\346\263\225\347\247\221\346\231\256/Transformer/\345\244\247\346\250\241\345\236\213KV Cache\350\212\202\347\234\201\347\245\236\345\231\250MLA\345\255\246\344\271\240\347\254\224\350\256\260\357\274\210\345\214\205\345\220\253\346\216\250\347\220\206\346\227\266\347\232\204\347\237\251\351\230\265\345\220\270\346\224\266\345\210\206\346\236\220\357\274\211.md" @@ -0,0 +1,499 @@ + + +> 首先，本文回顾了MHA的计算方式以及KV Cache的原理，然后深入到了DeepSeek V2的MLA的原理介绍，同时对MLA节省的KV Cache比例做了详细的计算解读。接着，带着对原理的理解理清了HuggingFace MLA的全部实现，每行代码都去对应了完整公式中的具体行并且对每个操作前后的Tensor Shape变化也进行了解析。我们可以看到目前的官方实现在存储KV Cache的时候并不是存储隐向量，而是把隐向量都解压缩变成了标准的MHA的KV Cache，实际上是完全不能节省显存的。接着，就继续学习了一下清华大学的ZHANG Mingxing组实现的MLA矩阵吸收的工程实现，在这一节也详细分析了原理包括$W^{UK}$和$W_{UV}$分别如何吸收到$W_{UQ}$和$W_o$中，分析了实现了矩阵吸收的每行代码的原理以及操作发生前后相关Tensor的维度变化。接着，对矩阵吸收代码实现里的矩阵乘法的性质进行分析，可以看到MLA在大多数阶段都是计算密集型而非访存密集型的。最后引用了作者团队的Benchmark结果，以及说明为何不是直接保存吸收后的大投影矩阵，而是在forward里面重新计算两个矩阵的吸收。 + +这里提一下，我维护的几个记录个人学习笔记以及社区中其它大佬们的优秀博客链接的仓库都获得了不少star，感谢读者们的认可，我也会继续在开源社区多做贡献。github主页：https://github.com/BBuf ，欢迎来踩 + +![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/c1dd6cb3f0c04c378780d3269a97e1f2.png) + +# 0x0. 前言 +这篇文章主要是对Deepseek2提出的优化KV Cache的MLA方法做个人理解，特别是关于MLA的矩阵吸收部分，这部分Paper以及官方开源实现没有给出。然后，开源社区大佬确实出手很快，在知乎的《如何看待 DeepSeek 发布的 MoE 大模型 DeepSeek-V2？》问题下清华大学的ZHANG Mingxing组就提出了MLA的矩阵吸收实现并且给出了一个兼容现有Transformers实现的PR（https://huggingface.co/deepseek-ai/DeepSeek-V2-Chat/discussions/12）。 +​ + + +![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/b741f77a480149c4bd12b71a6a426933.png) + + +翻译： + +这个修补过的DeepseekV2Model包含了对DeepseekV2Attention的以下修改，以减少VRAM消耗并提高效率： + +1. 不再缓存解压缩的Key/Value状态，而仅缓存低秩Key-Value联合压缩以及Key的解耦RoPE部分。 为了重用transformers库的缓存实用程序，我们将k_pe视为key_states，将compressed_kv视为value_states。 +2. 采用DeepseekV2论文中描述的吸收技术，通过改变计算Query和Output向量时的乘法顺序。这不仅节省了中间张量的内存消耗，还减少了浮点运算的次数。 +3. 分别计算RoPE部分和非RoPE部分的注意力分数，然后将它们相加。原始实现将Query/Key向量的两部分连接起来，但在缓存压缩Key/Value状态时由于不必要的数据广播和内存往返而被证明效率较低。 + +通过应用上述更改，MLA模块在解码阶段可以在NVIDIA A100-PCIE-40GB GPU上对单个请求实现高达20.4倍的加速，对32个批处理请求实现3.63倍的加速；在NVIDIA GeForce RTX 4080上，对单个和批处理请求分别实现26.2倍和3.52倍的加速。 + +这篇博客就是以我的视角来理解清楚这个PR中矩阵吸收的工程实现以及为何它可以加速现有的Deepseek2 MLA实现。本文先回顾一下MHA的Decode阶段KV Cache由来，然后根据paper的公式去理解Transformers中的DeepseekV2Attention类的实现。再接下来就学习一下ZHANG Mingxing大佬组所提出的MLA矩阵吸收工程实现。 + +Paper Link：https://arxiv.org/pdf/2405.04434 + +# 0x1. MHA 解码阶段KV Cache由来 + +首先回顾一下 MHA 机制在Decode阶段的原理和实现。 + +> 下面的公式来源也是DeepSeek2 paper，我做了更精细的解释。 + +假设batch_size为1，另外由于是解码阶段，输入只有一个token，所以序列的长度也是1，所以输入可以表示为$h_t \in \mathbb{R}^d$。接着假设embedding词表维度为$d$，并且有$n_h$表示注意力头的数量，$d_h$表示每个注意力头的维度。 + +> t表示解码阶段当前是第几个token。 + +然后通过$W^Q, W^K, W^V \in \mathbb{R}^{d_h n_h \times d}$三个参数矩阵得到$q_t, k_t, v_t \in \mathbb{R}^{d_h n_h}$，具体方法就是三个矩阵乘： + +$$ +q_t = W^Q h_t, +\newline +k_t = W^K h_t, +\newline +v_t = W^V h_t, +$$ + +在 MHA 的计算中，这里的 $q_t, k_t, v_t$ 又会分割成 $n_h$ 个注意力头，即： + +$$ +\begin{bmatrix} +q_{t,1}; q_{t,2}; \cdots ; q_{t,n_h} +\end{bmatrix} = q_t +\newline +\begin{bmatrix} +k_{t,1}; k_{t,2}; \cdots ; k_{t,n_h} +\end{bmatrix} = k_t +\newline +\begin{bmatrix} +v_{t,1}; v_{t,2}; \cdots ; v_{t,n_h} +\end{bmatrix} = v_t +$$ + +这里 $q_{t,i}, k_{t,i}, v_{t,i} \in \mathbb{R}^{d_h}$ 分别表示query、key和value的第$i$个头的计算结果。 + +接下来就是计算注意力分数和输出了，公式如下： + +$$ +o_{t,i} = \sum_{j=1}^{t} \text{Softmax}_j \left( \frac{q_{t,i} k_{j,i}}{\sqrt{d_h}} \right) v_{j,i}, +\newline +u_t = W^O [o_{t,1}; o_{t,2}; \cdots ; o_{t,n_h}] +$$ + +这里 $W^O \in \mathbb{R}^{d \times d_h n_h}$ 表示输出映射矩阵。从上面的公式可以看出来，对于当前的第 $t$ 个 token的query，会和$t$之前所有token的key, value做注意力计算，并且由于token by token的生成$t$之前所的有token对应的$k$，$v$我们都可以Cache下来，避免重复计算，这就是KV Cache的由来。 + +对于一个$l$层的标准MHA的网络来说，每个token需要的KV Cache大小为$2n_hd_hl$，其中2表示bf16的字节。 + +为了改进KV Cache，演化了一系列AI Infra的工作，比如Paged Attention, GQA, MLA包括最新的[GQA，MLA之外的另一种KV Cache压缩方式：动态内存压缩（DMC）](https://mp.weixin.qq.com/s/5pd4fF14ZUgYeM4UXA7ujQ)，[vAttention：用于在没有Paged Attention的情况下Serving LLM](https://mp.weixin.qq.com/s/F87-Qoo3xYGbwTTYr68guw) 。 + +# 0x2. DeepSeek2 MLA 原理介绍 + + +![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/b7b65080b2364006b8c441cedcb9ca03.png) + +上面的图是Deepseek2 paper中对集中常见KV Cache压缩方法的对比，可以看到MLA的核心是对keys和values进行低秩联合压缩来减少KV Cache。对应paper的公式9-11。 + + +![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/7294ec4a6602427193766c41c034a555.png) + +其中， + +- $\mathbf{c}_{t}^{KV} \in \mathbb{R}^{d_c}$：表示对 key 和 value 压缩后的隐向量 latent vector ，这里 $d_c (\ll d_h n_h)$ 表示 KV Cache压缩的维度。 +- $\mathbf{W}^{DKV} \in \mathbb{R}^{d_c \times d}$：表示向下映射 down-projection 矩阵 +- $\mathbf{W}^{UK}, \mathbf{W}^{UV} \in \mathbb{R}^{d_h n_h \times d_c}$：表示向上映射 up-projection 矩阵 + +这样在推理时，只需要缓存隐向量 $\mathbf{c}_{t}^{KV}$ 即可，因此 MLA 对应的每一个 token 的 KV Cache 参数只有 $2d_c l$ 个，其中$l$是网络层数，$2$是bfloat16的字节。 + +此外，为了降低训练过程中的激活内存，DeepSeek2还对query进行低秩压缩，即便这并不能降低KV Cache： + + +![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/258030d8230a43878bb80b7dc407f02d.png) + +其中， + +- $\mathbf{c}_{t}^{Q} \in \mathbb{R}^{d'_c}$：表示将 queries 压缩后的隐向量，$d'_c (\ll d_h n_h)$ 表示 query 压缩后的维度 +- $\mathbf{W}^{DQ} \in \mathbb{R}^{d'_c \times d}, \mathbf{W}^{UQ} \in \mathbb{R}^{d_h n_h \times d'_c}$ 分别表示 down-projection 和 up-projection 矩阵 + + +接下来MLA讨论的一个问题是，在上面的压缩过程中我们并没有考虑到RoPE。原始的RoPE需要在query和key中融入相对位置信息。在MLA中，在query中融入相对位置信息是比较容易的，但是由于KV Cache缓存的是压缩后的低秩key-value信息，这里面是没办法融入相对位置信息的。 + +> 关于RoPE为何不兼容MLA，苏神的博客里有更深刻的解释，建议阅读 https://kexue.fm/archives/10091 + +> 下面是对 Paper 的 Decoupled Rotary Position Embedding 章节进行解释。 + +识别图中的所有文本和公式如下： + +由于对 query 和 key 来说，RoPE 都是位置敏感的。如果对 $\mathbf{k}_{t}^{C}$ 采用 RoPE，那么当前生成 token 相关的 RoPE 矩阵会在 $\mathbf{W}^{Q}$ 和 $\mathbf{W}^{UK}$ 之间，并且矩阵乘法不遵循交换律，因此在推理时 $\mathbf{W}^{UK}$ 就无法整合到 $\mathbf{W}^{Q}$ 中。这就意味着，推理时我们必须重新计算所有之前 tokens 的 keys，这将大大降低推理效率。 + +> 这里的$\mathbf{W}^{UK}$ 就整合到 $\mathbf{W}^{Q}$ 请看下面截图的解释，来自苏神的博客。我会在下一大节再仔细讨论这个原理。 +> +>![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/8afd59f0cb32462f8ccf8fa2e3dcd6cb.png) + + +因此，DeepSeek2提出了解耦 RoPE 策略，具体来说： + +使用额外的多头 queries $\mathbf{q}_{t, i}^{R} \in \mathbb{R}^{d_h^R}$ 以及共享的 key $\mathbf{k}_{t}^{R} \in \mathbb{R}^{d_h^R}$ 来携带 RoPE 信息，其中 $d_h^R$ 表示解耦的 queries 和 key 的一个 head 的维度。 + +基于这种解耦的 RoPE 策略，MLA 遵循的计算逻辑为： + + +![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/99a75067db544bc1a1a81d1714eb36b5.png) + +其中： + +- $\mathbf{W}^{QR} \in \mathbb{R}^{d_h^R n_h \times d_c'}$ 和 $\mathbf{W}^{KR} \in \mathbb{R}^{d_h^R \times d}$ 分别表示计算解耦后的 queries 和 key 的矩阵 +- RoPE($\cdot$) 表示应用 RoPE 的操作； $[ \cdot ; \cdot ]$ 表示拼接操作 + +推理时，只需要缓存解耦后的 key 即可，当然还有上面的隐向量 $\mathbf{c}_{t}^{KV}$，因此对于单个 token 的 KV Cache 只包含 $(d_c + d_h^R)l$ 个元素，这里没考虑层数和bf16的字节数。具体可以看一下Table 1的数据对比： + + + +![](https://img-blog.csdnimg.cn/direct/bd1018cce7834e6bbc50ad2f008f2644.png) + +翻译一下： + +> 表1 | 各种注意力机制中每个token的KV Cache对比。$n_h$ 表示注意力头的数量，$d_h$ 表示每个注意力头的维度，$l$ 表示层数，$n_g$ 表示GQA中的组数，$d_c$ 和 $d_h^R$ 分别表示KV压缩维度和MLA中解耦后queries和key的每头维度。KV Cache的数量以元素的数量来衡量，而不考虑存储精度。对于DeepSeek-V2，$d_c$ 被设置为$4d_h$ 而 $d_h^R$ 被设置为$\frac{d_h}{2}$。因此，其KV Cache等于只有2.25组的GQA，但其性能强于MHA。 + + +原理的话应该就是这些了，接下来就带着原理阅读DeepseekV2Attention的实现。 + + +这里再特别说明一下Paper中相比于Dense的Deepseek 67B（或者LLaMa3 70B）节省93.3% KV Cache的计算方法： + + + +![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/56a60c837c6f46e9894ceba2414f29da.png) + + +首先是层数，DeepSeek2是60层，而Deepseek 67B为95层，层数的节省比例为 60 / 95 + +然后是单层的KV Cache，比例是(4.5 x 128) / (2 x 8 x 128)，其中2表示K和V，8表示num_attention_heads，128表示head_size，4.5则是上面MLA中的9/2的压缩隐向量。 + +此外，DeepSeek2针对KV Cache使用了6Bit量化，节省比例为 6 / 16 + +把这三个比例乘起来，再用1减掉就是93.3%的由来。 + +> 这里的6bit感觉是考虑了量化参数 scale 和 zero-point, 如果采用4bit量化, 而scale / zero-point 为fp32, 则当group_size=32时, 根据group-wise量化规则，每32个元素对应一组float32的scale和zero_point，那么每个元素平摊的位宽就多了2Bit，量化位宽等同于6Bit。 + +# 0x3. MLA HuggingFace官方实现代码解读 + +为了便于描述代码，这里直接把完整的公式贴出来，Paper的附录C： + + +![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/d6f72941452243feb85546253393f393.png) + + +对照上面的原理介绍下面对 DeepseekV2Attention 模块进行解读，代码链接：https://huggingface.co/deepseek-ai/DeepSeek-V2-Chat/blob/main/modeling_deepseek.py#L680 + +首先来看初始化部分，为了快速理解代码直接忽略掉RoPE计算相关的代码部分。 + +```python +# Copied from transformers.models.llama.modeling_llama.LlamaAttention with Llama->DeepseekV2 +class DeepseekV2Attention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: DeepseekV2Config, layer_idx: Optional[int] = None): + super().__init__() + + self.attention_dropout = config.attention_dropout + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + + self.max_position_embeddings = config.max_position_embeddings + self.rope_theta = config.rope_theta + # 对应 query 压缩后的隐向量的维度 d'_c + self.q_lora_rank = config.q_lora_rank + # 对应$d_h^R$， 表示应用了rope的 queries 和 key 的一个 head 的维度。 + self.qk_rope_head_dim = config.qk_rope_head_dim + # 对应 key-value 压缩后的隐向量维度 d_c + self.kv_lora_rank = config.kv_lora_rank + # value 的一个注意力头的隐藏层为度 + self.v_head_dim = config.v_head_dim + # 表示query和key的隐藏向量中应用rope部分的维度 + self.qk_nope_head_dim = config.qk_nope_head_dim + # 每一个注意力头的维度应该是两部分只和 + self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim + + self.is_causal = True + + # MLA 中对 Q 投影矩阵也做了一个低秩分解，对应生成 q_a_proj 和 q_b_proj 两个矩阵。 + # 其中 q_a_proj 大小为 [hidden_size, q_lora_rank] = [5120, 1536]， + # 对应上面公式中的W^DQ + self.q_a_proj = nn.Linear( + self.hidden_size, config.q_lora_rank, bias=config.attention_bias + ) + self.q_a_layernorm = DeepseekV2RMSNorm(config.q_lora_rank) + # q_b_proj 大小为 [q_lora_rank, num_heads * q_head_dim] = + # [q_lora_rank, num_attention_heads * (qk_nope_head_dim + qk_rope_head_dim)] = [1536, 128*(128+64)] = [1536, 24576] + # 对应上述公式中的W^UQ和W^QR合并后的大矩阵 + self.q_b_proj = nn.Linear( + config.q_lora_rank, self.num_heads * self.q_head_dim, bias=False + ) + # 与Q向量类似，KV向量的生成也是先投影到一个低维的 compressed_kv 向量（对应c_t^{KV}） + # 再升维展开。具体的代码涉及 kv_a_proj_with_mqa 和 kv_b_proj 两个参数矩阵。 + # 其中 kv_a_proj_with_mqa 大小为 [hidden_size， kv_lora_rank + qk_rope_head_dim] + # = [5120, 512 + 64] = [5120, 576]，对应上述公式中的W^{DKV}和W^{KR}。 + self.kv_a_proj_with_mqa = nn.Linear( + self.hidden_size, + config.kv_lora_rank + config.qk_rope_head_dim, + bias=config.attention_bias, + ) + self.kv_a_layernorm = DeepseekV2RMSNorm(config.kv_lora_rank) + # kv_b_proj 大小为 [kv_lora_rank， num_heads * (q_head_dim - qk_rope_head_dim + v_head_dim)] + # = [512, 128*((128+64)-64+128)] = [512, 32768]，对应上述公式中的W^{UK}和W^{UV}。 + # 由于 W^{UK} 只涉及 non rope 的部分所以维度中把 qk_rope_head_dim 去掉了。 + self.kv_b_proj = nn.Linear( + config.kv_lora_rank, + self.num_heads + * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim), + bias=False, + ) + # 对应完整公式的第 47 行 + self.o_proj = nn.Linear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=config.attention_bias, + ) +``` + +根据0x2节的原理介绍，现在已经可以把 DeepseekV2Attention 模块里面所有的权重矩阵都和初始化的代码对应起来了，如果你想继续看下去一定要理解到初始化的每行代码。 + +为了方便理解forward代码时回看公式，这里再重复贴一下完整公式： + + +![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/d6f72941452243feb85546253393f393.png) + + +接下来再看一下forward的代码，这对应了完整公式里面的详细计算过程： + +```python +def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if "padding_mask" in kwargs: + warnings.warn( + "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" + ) + # hidden_states对应公式中的h_t，的shape是(batch_size, seq_length, + # hidden_size)，其中 hidden_size 具体为 5120，假设batch_size和seq_length都为1 + bsz, q_len, _ = hidden_states.size() + + # 计算Q：对应完整公式中的 37-39 行，先降维再升维，好处是相比直接使用大小为 [5120, 24576] 的矩阵 + # [5120, 1536] * [1536, 24576] 这样的低秩分解在存储空间和计算量上都大幅度降低 + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2) + # 切分 rope 和非 rope 部分，完整公式中 40 行反过来 + q_nope, q_pe = torch.split( + q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) + + # 对应公式中的 41 和 43 行只是还没有加 rope + # 一个优化的 MLA KVCache 实现只需要缓存这个 compressed_kv 就行 + # kv_a_proj_with_mqa shape 为[hidden_size， kv_lora_rank + qk_rope_head_dim] + # = [5120, 512 + 64] = [5120, 576] + # 所以compressed_kv的shape就是[1, 1, 576] + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + # 对应完整公式的 44 行反过来 + compressed_kv, k_pe = torch.split( + compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + # 这里的 k_pe 和 上面的 q_pe 要扔给 RoPE模块，所以需要重整下shape + k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2) + # 对应公式中的 42 和 45 行，将 MLA 展开成标准 MHA 的形式 + kv = ( + self.kv_b_proj(self.kv_a_layernorm(compressed_kv)) + .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) + .transpose(1, 2) + ) + # 因为 kv_b_proj 打包了 W^{UK} 和 W^{UV} 把他们分离出来 + k_nope, value_states = torch.split( + kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1 + ) + # 获取key/value的序列长度 + kv_seq_len = value_states.shape[-2] + if past_key_value is not None: + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + # 给需要 rope 的部分加 rope + cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) + q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids) + + # 更新和拼接历史 KVCache，可以看到这里存储的是展开后的 MHA KVCache + # 其中 q_head_dim 等于 qk_nope_head_dim + qk_rope_head_dim + query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) + query_states[:, :, :, : self.qk_nope_head_dim] = q_nope + query_states[:, :, :, self.qk_nope_head_dim :] = q_pe + + key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) + key_states[:, :, :, : self.qk_nope_head_dim] = k_nope + key_states[:, :, :, self.qk_nope_head_dim :] = k_pe + # Transformers库中标准的 KV Cache 更新代码 + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) + + # 后续就是标准的多头自注意力计算了，为了篇幅，忽略这部分代码 + ... +``` + +读完这一节基本就能将MLA完整公式的每一行都搞清楚了，并且我们可以看到目前的官方实现在存储KV Cache的时候并不是存储隐向量，而是把隐向量都解压缩变成了标准的MHA的KV Cache，实际上是完全不能节省显存的。 + +# 0x4. 矩阵吸收 + +这一节就是学习和理解一下清华大学的ZHANG Mingxing组实现的MLA矩阵吸收（ https://zhuanlan.zhihu.com/p/700214123 ）。它的代码是直接应用在HF实现上的，所以可以很方便进行应用。为了理解再重复贴一下完整公式： + + +![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/d6f72941452243feb85546253393f393.png) + +以及Paper中提到的矩阵吸收， + + +![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/af87d35856ad490c90ce7819581aa7e6.png) + +## 0x4.1 W^{UK}的吸收 + +对于 $\mathbf{W}^{UK}$ 矩阵我们有： + +$\text{atten\_weights} = \mathbf{q}_{t}^\top \mathbf{k}_{t} = (\mathbf{W}^{UQ} \mathbf{c}_{t}^{Q})^\top \mathbf{W}^{UK} \mathbf{c}_{t}^{KV} = \mathbf{c}_{t}^{Q^\top} \mathbf{W}^{UQ^\top} \mathbf{W}^{UK} \mathbf{c}_{t}^{KV}$ + + + +也就是说我们实际上不需要将低维的 $\mathbf{c}_{t}^{KV}$ 展开再计算，而是直接将 $\mathbf{W}^{UK}$ 通过结合律先和左边做乘法。 + +```python +# 以下和上一节的MLA forward部分实现相同 +# hidden_states对应公式中的h_t，的shape是(batch_size, seq_length, +# hidden_size)，其中 hidden_size 具体为 5120，假设batch_size为1，seq_length为q_len +bsz, q_len, _ = hidden_states.size() + +# 计算Q：对应完整公式中的 37-39 行，先降维再升维，好处是相比直接使用大小为 [5120, 24576] 的矩阵 +# [5120, 1536] * [1536, 24576] 这样的低秩分解在存储空间和计算量上都大幅度降低 +q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) +q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2) +# 切分 rope 和非 rope 部分，完整公式中 40 行反过来 +q_nope, q_pe = torch.split( + q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 +) + +# 对应公式中的 41 和 43 行只是还没有加 rope +# 一个优化的 MLA KVCache 实现只需要缓存这个 compressed_kv 就行，不过后面实际上展开 +# hidden_states 的 shape 为 (1, past_len, hidden_size) +# kv_a_proj_with_mqa shape 为[hidden_size， kv_lora_rank + qk_rope_head_dim] +# = [5120, 512 + 64] = [5120, 576] +# 所以compressed_kv的shape就是[1, past_len, 576] +compressed_kv = self.kv_a_proj_with_mqa(hidden_states) +# 对应完整公式的 44 行反过来 +compressed_kv, k_pe = torch.split( + compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 +) +# 这里的 k_pe 和 上面的 q_pe 要扔给 RoPE模块，所以需要重整下shape +k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2) +``` + +以下部分是$W^{UK}$的吸收需要做的改动，省略掉了把compressed_kv和k_pe加入到Transformers KV Cache的改动： + +```python +# 从 kv_b_proj 中分离的 W^{UK} 和 W^{UV} 两部分，他们要分别在不同的地方吸收 +kv_b_proj = self.kv_b_proj.weight.view(self.num_heads, -1, self.kv_lora_rank) +q_absorb = kv_b_proj[:, :self.qk_nope_head_dim,:] +out_absorb = kv_b_proj[:, self.qk_nope_head_dim:, :] + +cos, sin = self.rotary_emb(q_pe) +q_pe = apply_rotary_pos_emb(q_pe, cos, sin, q_position_ids) +# W^{UK} 即 q_absorb 被 q_nope 吸收 +q_nope = torch.einsum('hdc,bhqd->bhqc', q_absorb, q_nope) +# 吸收后 attn_weights 直接基于 compressed_kv 计算不用展开。 +attn_weights = torch.matmul(q_pe, k_pe.transpose(2, 3)) + torch.einsum('bhqc,blc->bhql', q_nope, compressed_kv) +attn_weights *= self.softmax_scale +``` + +主要是吸收这里的代码需要仔细理解： +- 从0x4节的讲解已经知道kv_b_proj就是$W^{UK}$ 和 $W^{UV}$两部分，这里是把$W^{UK}$吸收到$W^{UQ}$，所以需要先把两者分离出来。注意到 `self.kv_b_proj` weight shape为 `[kv_lora_rank， num_heads * (q_head_dim - qk_rope_head_dim + v_head_dim)] = [512, 128*((128+64)-64+128)] = [512, 32768]`，所以`kv_b_proj`的shape为`[num_heads,q_head_dim - qk_rope_head_dim + v_head_dim , kv_lora_rank]`, `q_absorb`的shape为`[num_heads, qk_nope_head_dim , kv_lora_rank]=[128, 128, 512]`，同样`out_absorb`的shape为`[num_heads, v_head_dim , kv_lora_rank]=[128, 128, 512]`。 +- `q_nope = torch.einsum('hdc,bhqd->bhqc', q_absorb, q_nope) `这行代码中，`q_nope`的shape是`[batch_size, num_heads, q_len, q_head_dim]`。所以这行代码就是一个矩阵乘法，把$W^{UK}$吸收到$W^{UQ}$。 +- 吸收后 attn_weights 直接基于 compressed_kv 计算不用展开。对应`torch.einsum('bhqc,blc->bhql', q_nope, compressed_kv)`这行代码。其中q_nope 的维度是 `[batch_size, num_heads, q_len, kv_lora_rank]`，compressed_kv 是 `[batch_size, past_len, kv_lora_rank]`，输出的维度是 `[batch_size, num_heads, q_len, past_len]`。 +- 此外，我们还可以观察到`torch.matmul(q_pe, k_pe.transpose(2, 3))`这行代码是分开计算了RoPE部分的q和k的注意力计算再求和，没有和原本的实现一样将加上了 rope 的 q_pe/k_pe 和没加 rope 的 q_nope/k_nope 拼接起来一起，也就是下面的代码。作者团队把这个拆分叫做 Move Elision 的优化，后续有性能对比。 + +```python +# 更新和拼接历史 KVCache，可以看到这里存储的是展开后的 MHA KVCache +# 其中 q_head_dim 等于 qk_nope_head_dim + qk_rope_head_dim +query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) +query_states[:, :, :, : self.qk_nope_head_dim] = q_nope +query_states[:, :, :, self.qk_nope_head_dim :] = q_pe + +key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim) +key_states[:, :, :, : self.qk_nope_head_dim] = k_nope +key_states[:, :, :, self.qk_nope_head_dim :] = k_pe +``` + +> 除了压缩KV Cache之外，我们还可以观察到上面涉及到的2个矩阵乘法实际上都来到了计算密集的领域，例如对于 ` torch.einsum('bhqc,blc->bhql', q_nope, compressed_kv) ` 。由于不同 head 的 q_nope 部分 share 了共同的 compressed_kv 部分，实际计算的是 batch_size 个 [num_heads * q_len, kv_lora_rank] 和 [past_len, kv_lora_rank] 的矩阵乘法。计算等价于一个 MQA 操作，计算强度正比于 num_heads 的也就是 128。 + +## 0x4.2 W^{UV}的吸收 +对于$W^{UV}$我们有： + +$v= W^{UV}c_t^{KV}$ (对应公式的45行) +$u=$ attn_weights 和 $v$ 的矩阵乘法 +$o = uW_o$ + +这里把 attn_weights 记作 $M$，那么有： + +$o = MW^{UV}c_t^{KV}W_o$ + +和$W^{UK}$的吸收过程类似，利用结合律改变计算顺序，那么有： + +$o = Mc_t^{KV}W^{UV}W_o$ + +我们可以把$W^{UV}$吸收到$W_o$中，对应的代码实现： + +```python +# attn_weight的shape是[batch_size, num_heads, q_len, past_len] +# compressed_kv的shape是[batch_size, past_len, kv_lora_rank] +# attn_output的shape是[batch_size, num_heads, q_len, kv_lora_rank] +attn_output = torch.einsum('bhql,blc->bhqc', attn_weights, compressed_kv) +# out_absorb的shape是[num_heads, v_head_dim , kv_lora_rank] +# out_absorb.mT的shape是[num_heads, kv_lora_rank, v_head_dim] +# 最终attn_output的shape是[batch_size, num_heads, q_len, v_head_dim] +attn_output = torch.matmul(attn_output, out_absorb.mT) +``` + +> 注意：.mT 方法用于获取张量的转置（transpose）。对于二维张量（矩阵），转置操作会交换其行和列。而对于高维张量，.mT 会交换最后两个维度。 + +> 同样，这里除了压缩KV Cache之外，我们还可以观察到上面涉及到的2个矩阵乘法实际上也来到了计算密集的领域，例如对于 `attn_output = torch.einsum('bhql,blc->bhqc', attn_weights, compressed_kv) ` 。由于不同 head 的 attn_weights 部分 share 了共同的 compressed_kv 部分，实际计算的是 batch_size 个 [num_heads * q_len, kv_lora_rank] 和 [past_len, kv_lora_rank] 的矩阵乘法。计算等价于一个 MQA 操作，计算强度正比于 num_heads 的也就是 128。因此相比 MHA，吸收后的 MLA 计算强度要大得多，因此也可以更加充分的利用 GPU 算力。 + +## 0x4.3 MLA MatMul的性质 +上面几乎分析了每个矩阵乘法的计算shape，可以发现除了在对q做计算时涉及到gemv之外，也就是`q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))`，其它地方的矩阵乘运算q_len维度都是和num_heads在一起做计算，而num_heads在Deepseek2的配置里面已经是128了，导致其它的Matmul几乎都落在了计算密集的范畴。 + +综上，对于MLA模块来说，有很大比例的MatMul都达到了计算密集的范畴，这改变了之前MHA的访存密集的性质。然而，在整个网络中间，由于有MoE模块的存在，如果BatchSize不够大（无法激活所有的expert，导致计算和访存比=计算密度低）还是无法整体达到计算密集的范围，但MLA节省的KV Cache本就可以让DeepSeek2的Batch大幅度提升，所以均摊到每个token的带宽需求相比于Dense的LLaMa3 70B也会大幅度下降。 +## 0x4.4 Benchmark + +最后，作者团队在他们的Blog中给出了一些Benchmark结果，可以看到这个矩阵吸收的有效性。 + + +![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/8bf647d03a004a079de1ef2acdaecf2d.png) + + +![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/e743e80b9583473ea67c4642a4bca7fc.png) + +图中的标签分别表示 原始的解压缩版本CacheDecompressed (CD)，KV缓存压缩后的CacheCompressed（CC），吸收后直接使用 compressed_kv 计算的 Absorbed_CacheCompressed (A_CC) 版本，和增加了 move elision 优化的最终版本 Absorbed_CacheCompressed_MoveElision (A_CC_ME)。 + +## 0x4.5 是否需要重计算 +以$W^{UK}$的吸收为例子，这里实际上是说在矩阵吸收的时候是否要提前把$W^{UK}$和$W^{UQ}$的矩阵乘结果保存下来，而不是在forward的时候重计算。作者在评论区回复过这个问题，意思就是直接在forward的时候重计算速度会更优。博客里面的解释如下： + + +![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/0296a658ddb54c99a7d5f2ee26edfe8d.png) + +# 0x5. 总结 +这就是本篇博客的全部内容了，这里总结一下。首先，本文回顾了MHA的计算方式以及KV Cache的原理，然后深入到了DeepSeek V2的MLA的原理介绍，同时对MLA节省的KV Cache比例做了详细的计算解读。接着，带着对原理的理解理清了HuggingFace MLA的全部实现，每行代码都去对应了完整公式中的具体行并且对每个操作前后的Tensor Shape变化也进行了解析。我们可以看到目前的官方实现在存储KV Cache的时候并不是存储隐向量，而是把隐向量都解压缩变成了标准的MHA的KV Cache，实际上是完全不能节省显存的。接着，就继续学习了一下清华大学的ZHANG Mingxing组实现的MLA矩阵吸收的工程实现，在这一节也详细分析了原理包括$W^{UK}$和$W_{UV}$分别如何吸收到$W_{UQ}$和$W_o$中，分析了实现了矩阵吸收的每行代码的原理以及操作发生前后相关Tensor的维度变化。接着，对矩阵吸收代码实现里的矩阵乘法的性质进行分析，可以看到MLA在大多数阶段都是计算密集型而非访存密集型的。最后引用了作者团队的Benchmark结果，以及说明为何不是直接保存吸收后的大投影矩阵，而是在forward里面重新计算两个矩阵的吸收。 + + +# 0x6. 参考资料 +- https://www.zhihu.com/question/655172528 +- https://arxiv.org/pdf/2405.04434 + diff --git "a/docs/project/\351\203\250\347\275\262\344\274\230\345\214\226/AI \347\247\273\345\212\250\347\253\257\347\256\227\346\263\225\344\274\230\345\214\226/60\350\241\214\344\273\243\347\240\201\345\212\240\351\200\23720\345\200\215_NEON\345\256\236\347\216\260\346\267\261\345\272\246\345\255\246\344\271\240OD\344\273\273\345\212\241\345\220\216\345\244\204\347\220\206\347\273\230\346\241\206.md" "b/docs/project/\351\203\250\347\275\262\344\274\230\345\214\226/AI \347\247\273\345\212\250\347\253\257\347\256\227\346\263\225\344\274\230\345\214\226/60\350\241\214\344\273\243\347\240\201\345\212\240\351\200\23720\345\200\215_NEON\345\256\236\347\216\260\346\267\261\345\272\246\345\255\246\344\271\240OD\344\273\273\345\212\241\345\220\216\345\244\204\347\220\206\347\273\230\346\241\206.md" new file mode 100644 index 0000000..690aaa8 --- /dev/null +++ "b/docs/project/\351\203\250\347\275\262\344\274\230\345\214\226/AI \347\247\273\345\212\250\347\253\257\347\256\227\346\263\225\344\274\230\345\214\226/60\350\241\214\344\273\243\347\240\201\345\212\240\351\200\23720\345\200\215_NEON\345\256\236\347\216\260\346\267\261\345\272\246\345\255\246\344\271\240OD\344\273\273\345\212\241\345\220\216\345\244\204\347\220\206\347\273\230\346\241\206.md" @@ -0,0 +1,250 @@ +**【前言】** 本文版权属于GiantPandaCV，未经允许，请勿转载！ +最近在学neon汇编加速，由于此前OD任务发现在检测后处理部分使用OpenCV较为占用资源且耗时，遂尝试使用NEON做后处理绘框，以达到加速并降低CPU资源消耗的目的。 +### 一、实现思路 +假设对一张Mat图像进行操作（其实也不仅仅是Mat对象，理论上只要知道图像通道的首指针即可），在ARM端使用NEON instrinc指令集里实现一个后处理绘框的功能，可以简单罗列成以下几步： +**1. 定义参数：** 首先确定图像的宽度和高度，图像的首地址指针，以及边界（边框）的厚度。 +**2. 向量寄存器加载：** 使用NEON的加载指令从内存中加载像素数据到向量寄存器中。 +**3. 处理上下边框：** + + - 对于顶部边界，遍历整个第一行的像素，并使用NEON的存储指令将特定颜色值写回到这些位置（比如想绘制的是绿框，那么需要将B通道的绘框元素数据更改为0，G通道为255，R通道为0）。 + - 同样地，对于底部边界，遍历最后一行的像素并执行相同的操作。 + + **4.处理左右边框：** +这个稍微复杂一些，因为需要处理每一行的开始和结束位置。一种方法是使用循环，每次处理一行，然后更新寄存器中的值以反映特定颜色。我们可以使用NEON的广播指令来创建一个包含特定颜色所有分量的向量，然后使用存储指令将其写入到图像的左侧和右侧边界。 +**5.边框优化：** +由于很多检测框的宽度很难保证一定是SIMD WIDTH的倍数，这就造成了在绘图时一些不必要的麻烦，举个例子，假设检测框的width是97，SIMD WIDTH的长度是16（一次性处理16个元素的向量寄存器），那么97/16=6······1，刚好多出了1个pixel，此时需要某些处理措施规避这种情况。 +### 二、实现过程 +#### 2.1 定义参数 +首先确定图像的宽度和高度，本次测试所获得的检测框均由这篇博文中的end2end模型中获得【1】，也就是在绘框前，我们会得到一个vector数组，均为通过nms获得的检测框，这个数组数据排列格式如下： +![](https://img-blog.csdnimg.cn/img_convert/06b2b83ca74f9055884d8a57caa2c755.png) +一个box对应四个元素，其实box是按照obj的score排列，但为了方便讲解，我们假设他是按从左到右顺序排列，由于测试的图片均为COCO2017 Val中的数据，图片尺寸中值远大于320，为了美观，此篇博文默认绘框边界（边框）的厚度为2，也就是占满2个pixel。 +函数定义如下： + +```cpp +void neon_rectangle_blod(uint8_t *img, uint16_t img_step, uint16_t x, uint16_t y, uint16_t w, uint16_t h, uint8_t blue, uint8_t green, uint8_t red) +``` +函数形参解释：$*img$为图像首指针，$img step$指图像的width，$x,y$指检测框左上角，$w,h$指检测框的宽高， $blue, green, red$指三色通道需要填充的数值。 +#### 2.2 向量寄存器加载 +这一步需要将图像BGR通道元素加载到寄存器，由于图像一般为uint8格式，这里可以使用最大的寄存器，把位宽拉满，也就是一次性操作16个元素，调用NEON instrinc中的vld3q_u8加载图像BGR数据到uint8x16x3_t寄存器中，再将单个通道的数据分发到到单个uint8x16_t寄存器中，伪代码如下: + +```cpp +// 假设img是指向图像BGR数据的指针 +uint8x16x3_t bgr_data = vld3q_u8((uint8_t *) img); + +// 分别将BGR通道的数据分发到单独的uint8x16_t寄存器 +uint8x16_t reg_b = bgr_data.val[0]; // 蓝色通道 +uint8x16_t reg_g = bgr_data.val[1]; // 绿色通道 +uint8x16_t reg_r = bgr_data.val[2]; // 红色通道 + +// 后续对每个通道进行单独的操作 +...... +``` +#### 2.3 处理上下边框 +我们需要定位到上下边框的起始位置，获取起始位置的地址，再将地址往后以16个pixel为一个SIMD_WIDTH塞入寄存器，将寄存器中的B,G,R通道进行向量赋值，表示一次性处理16个数据流位宽，代码如下: + +```cpp +// 绘制矩形的上下边界 + for (uint16_t i = 0; i < w; i += 16) + { + // 计算当前行的起始地址 + uint8_t *top_row1 = img + (y * img_step + x + i) * 3; + uint8_t *bottom_row1 = img + ((y + h) * img_step + x + i) * 3; + + // 使用NEON指令集并行加载和存储颜色 + uint8x16x3_t pixels_top1 = vld3q_u8(top_row1); + uint8x16x3_t pixels_bottom1 = vld3q_u8(bottom_row1); + + // 绘制顶部和底部线条 + pixels_top1.val[0] = neon_color_b; // 蓝色通道 + pixels_top1.val[1] = neon_color_g; // 绿色通道 + pixels_top1.val[2] = neon_color_r; // 红色通道 + + pixels_bottom1.val[0] = neon_color_b; + pixels_bottom1.val[1] = neon_color_g; + pixels_bottom1.val[2] = neon_color_r; + + vst3q_u8(top_row1, pixels_top1); + vst3q_u8(bottom_row1, pixels_bottom1); + + // 计算当前行的起始地址 + uint8_t *top_row2 = img + ((y + 1) * img_step + x + i) * 3; + uint8_t *bottom_row2 = img + ((y + h - 1) * img_step + x + i) * 3; + + // 使用NEON指令集并行加载和存储颜色 + uint8x16x3_t pixels_top2 = vld3q_u8(top_row2); + uint8x16x3_t pixels_bottom2 = vld3q_u8(bottom_row2); + + // 绘制顶部和底部线条 + pixels_top2.val[0] = neon_color_b; // 蓝色通道 + pixels_top2.val[1] = neon_color_g; // 绿色通道 + pixels_top2.val[2] = neon_color_r; // 红色通道 + + pixels_bottom2.val[0] = neon_color_b; + pixels_bottom2.val[1] = neon_color_g; + pixels_bottom2.val[2] = neon_color_r; + + vst3q_u8(top_row2, pixels_top2); + vst3q_u8(bottom_row2, pixels_bottom2); + } +``` +#### 2.4 处理左右边框 +这里就有点难受了，因为是ARM架构通用的汇编，不像一些厂家有专门处理竖直方向的寄存器或者额外的硬件加速模块，所以这一步只能老老实实一个pixel一个pixel的去涂，因此和OpenCV的处理方式没有太大差异，代码如下: + +```cpp +// 绘制矩形的左右边界 + for (uint16_t j = 0; j < h; j++) + { + // 计算当前列的起始地址 + uint8_t *left_col1 = img + ((y + j) * img_step + x) * 3; + uint8_t *right_col1 = img + ((y + j) * img_step + (x + w)) * 3; + // 设置左边和右边列的颜色 + left_col1[0] = right_col1[0] = blue; + left_col1[1] = right_col1[1] = green; + left_col1[2] = right_col1[2] = red; + + // 计算当前列的起始地址 + uint8_t *left_col2 = img + ((y + j) * img_step + x + 1) * 3; + uint8_t *right_col2 = img + ((y + j) * img_step + (x + w) - 1) * 3; + // 设置左边和右边列的颜色 + left_col2[0] = right_col2[0] = blue; + left_col2[1] = right_col2[1] = green; + left_col2[2] = right_col2[2] = red; + } +``` +#### 2.5 优化边框 +这里提供一种思路，既然没办法确保检测框的宽度刚好是SIMD_WIDTH的倍数，那我们就将宽度扩充或者减小到SIMD_WIDTH的倍数，但为了美观处理，不管是扩充还是减小宽度，我们都离不开一个操作，那就是中心对齐，以扩宽为例，如下图所示： +![](https://img-blog.csdnimg.cn/img_convert/42882337344bbb6b2b1c71eebc578030.png) +那么，就有很好的方式去应对这种情况，我们假设检测框的width对SIMD_WIDTH进行mod操作，如果余数小于 +$SIMD_`WIDTH/2$，对检测框width进行缩小操作，反之，则进行扩充操作，代码如下： + +```cpp +void check_point(int *x1, int *x2, int nstride) +{ + int mod, w, xc, nw; + w = *x2 - *x1; + xc = *x1 + (int)(w / 2); + mod = w % nstride; + if (mod > (nstride / 2)) + { + *x1 = xc - (int)((w + nstride - mod) / 2); + *x2 = xc + (int)((w + nstride - mod) / 2); + } + else + { + nw = w - mod; + *x1 = xc - int(nw / 2); + *x2 = xc + int(nw / 2); + } +} +``` +### 三、测试结果 +测试机器为4+32内存的树莓派4B，共带有4颗A72核，我们分别使用NEON和OpenCV作为【1】中end2end模型出框后的后处理绘框函数，测试数据为COCO2017 Val数据集，将两个程序用taskset -c先绑定在编号为0的核上，得出两者在处理5000张图的处理速度差异，如下所示: +![](https://img-blog.csdnimg.cn/img_convert/df9da9a26f89ef3959403a7171c53f26.png) +其中，cost time为推理完5000张图的所有耗时，单位为ms，average cost time为处理单张图片的耗时，单位为us，我们可以看到，在单个A72上，NEON实现的绘框函数要比OpenCV快了20倍左右。 +此外，OpenCV的强大源于多核并行，为了能更加客观且全面的测试出两者的性能差异，我们在OpenCV版本的基础上，不断增加核进行测试，得出以下测试图例： +![](https://img-blog.csdnimg.cn/img_convert/3ee7c14aac7139cb84b26f8750cbf392.png) +图中P/ms表示1ms能处理多少图，越高表示每毫秒处理图越多，单图绘框速度越快，从图可以看出，单核运行的NEON绘框的速度依旧稳稳碾压多核并行的OpenCV。 +OpenCV绘框效果如下: +![](https://img-blog.csdnimg.cn/img_convert/22dbe287f5403d718ecc0891d2edebeb.jpeg) +NEON汇编绘框效果如下: +![](https://img-blog.csdnimg.cn/img_convert/b2a8ada1a54194f0d852896a5615c425.jpeg) +### 四、完整代码 + +```cpp +void check_point(int *x1, int *x2, int nstride) +{ + int mod, w, xc, nw; + w = *x2 - *x1; + xc = *x1 + (int)(w / 2); + mod = w % nstride; + if (mod > (nstride / 2)) + { + *x1 = xc - (int)((w + nstride - mod) / 2); + *x2 = xc + (int)((w + nstride - mod) / 2); + } + else + { + nw = w - mod; + *x1 = xc - int(nw / 2); + *x2 = xc + int(nw / 2); + } +} + +void neon_rectangle_blod(uint8_t *img, uint16_t img_step, uint16_t x, uint16_t y, uint16_t w, uint16_t h, uint8_t blue, uint8_t green, uint8_t red) +{ + // 创建一个全1的8位向量，用于绘制矩形的颜色 + uint8x16_t neon_color_b = vdupq_n_u8(blue); + uint8x16_t neon_color_g = vdupq_n_u8(green); + uint8x16_t neon_color_r = vdupq_n_u8(red); + + // 绘制矩形的上下边界 + for (uint16_t i = 0; i < w; i += 16) + { + // 计算当前行的起始地址 + uint8_t *top_row1 = img + (y * img_step + x + i) * 3; + uint8_t *bottom_row1 = img + ((y + h) * img_step + x + i) * 3; + + // 使用NEON指令集并行加载和存储颜色 + uint8x16x3_t pixels_top1 = vld3q_u8(top_row1); + uint8x16x3_t pixels_bottom1 = vld3q_u8(bottom_row1); + + // 绘制顶部和底部线条 + pixels_top1.val[0] = neon_color_b; // 蓝色通道 + pixels_top1.val[1] = neon_color_g; // 绿色通道 + pixels_top1.val[2] = neon_color_r; // 红色通道 + + pixels_bottom1.val[0] = neon_color_b; + pixels_bottom1.val[1] = neon_color_g; + pixels_bottom1.val[2] = neon_color_r; + + vst3q_u8(top_row1, pixels_top1); + vst3q_u8(bottom_row1, pixels_bottom1); + + // 计算当前行的起始地址 + uint8_t *top_row2 = img + ((y + 1) * img_step + x + i) * 3; + uint8_t *bottom_row2 = img + ((y + h - 1) * img_step + x + i) * 3; + + // 使用NEON指令集并行加载和存储颜色 + uint8x16x3_t pixels_top2 = vld3q_u8(top_row2); + uint8x16x3_t pixels_bottom2 = vld3q_u8(bottom_row2); + + // 绘制顶部和底部线条 + pixels_top2.val[0] = neon_color_b; // 蓝色通道 + pixels_top2.val[1] = neon_color_g; // 绿色通道 + pixels_top2.val[2] = neon_color_r; // 红色通道 + + pixels_bottom2.val[0] = neon_color_b; + pixels_bottom2.val[1] = neon_color_g; + pixels_bottom2.val[2] = neon_color_r; + + vst3q_u8(top_row2, pixels_top2); + vst3q_u8(bottom_row2, pixels_bottom2); + } + + // 绘制矩形的左右边界 + for (uint16_t j = 0; j < h; j++) + { + // 计算当前列的起始地址 + uint8_t *left_col1 = img + ((y + j) * img_step + x) * 3; + uint8_t *right_col1 = img + ((y + j) * img_step + (x + w)) * 3; + // 设置左边和右边列的颜色 + left_col1[0] = right_col1[0] = blue; + left_col1[1] = right_col1[1] = green; + left_col1[2] = right_col1[2] = red; + + // 计算当前列的起始地址 + uint8_t *left_col2 = img + ((y + j) * img_step + x + 1) * 3; + uint8_t *right_col2 = img + ((y + j) * img_step + (x + w) - 1) * 3; + // 设置左边和右边列的颜色 + left_col2[0] = right_col2[0] = blue; + left_col2[1] = right_col2[1] = green; + left_col2[2] = right_col2[2] = red; + } +} +``` +### 五、总结 +本篇博文主要讲述后处理绘框的汇编实现方式，在树莓派上的单核以及多核A72上都实现了加速，但时间关系未于其他开发板做比较，从去年开始，似乎4大+4小变成了业界主流，既4颗A76+4颗A57或者4颗A76+4颗A53，ARM端CPU算力要远远强过四颗A72，至于这种汇编实现方式，在这些开发板上能加速多少，确实不好说，有兴趣的朋友可以用这几十行代码去测试下~ +### 六、参考 +[1] https://zhuanlan.zhihu.com/p/672633849 +[2] https://zhuanlan.zhihu.com/p/698551682 +[3] https://developer.arm.com/documentation/ diff --git "a/docs/project/\351\203\250\347\275\262\344\274\230\345\214\226/AI \351\203\250\347\275\262\345\217\212\345\205\266\345\256\203\344\274\230\345\214\226\347\256\227\346\263\225/\345\246\202\344\275\225\344\275\277\347\224\250\342\200\234LoRa\342\200\235\347\232\204\346\226\271\345\274\217\345\212\240\350\275\275Onnx\346\250\241\345\236\213\357\274\232StableDiffusion\347\233\270\345\205\263\346\250\241\345\236\213 \347\232\204C++\346\216\250\347\220\206.md" "b/docs/project/\351\203\250\347\275\262\344\274\230\345\214\226/AI \351\203\250\347\275\262\345\217\212\345\205\266\345\256\203\344\274\230\345\214\226\347\256\227\346\263\225/\345\246\202\344\275\225\344\275\277\347\224\250\342\200\234LoRa\342\200\235\347\232\204\346\226\271\345\274\217\345\212\240\350\275\275Onnx\346\250\241\345\236\213\357\274\232StableDiffusion\347\233\270\345\205\263\346\250\241\345\236\213 \347\232\204C++\346\216\250\347\220\206.md" new file mode 100644 index 0000000..14b87d8 --- /dev/null +++ "b/docs/project/\351\203\250\347\275\262\344\274\230\345\214\226/AI \351\203\250\347\275\262\345\217\212\345\205\266\345\256\203\344\274\230\345\214\226\347\256\227\346\263\225/\345\246\202\344\275\225\344\275\277\347\224\250\342\200\234LoRa\342\200\235\347\232\204\346\226\271\345\274\217\345\212\240\350\275\275Onnx\346\250\241\345\236\213\357\274\232StableDiffusion\347\233\270\345\205\263\346\250\241\345\236\213 \347\232\204C++\346\216\250\347\220\206.md" @@ -0,0 +1,1015 @@ +## 如何使用“LoRa”的方式加载Onnx模型：StableDiffusion相关模型的C++推理 + +本文主要干了以下几个事： + +1.基于 onnxruntime，将 StableDiffusionInpaintPipeline、StableDiffusionControlNetImg2ImgPipeline(stablediffusion + controlnet + LoRa) C++工程化； + +2.输出一个 C++版本的 ddim-schduler 库； + +3.提供一种“LoRa”的 onnx 模型加载方式； + +4.所有相关代码、模型开源 + +> 项目地址: https://github.com/TalkUHulk/ai.deploy.box + +> 模型地址: https://huggingface.co/TalkUHulk/AiDB + +## StableDiffusionInpaint + +### 模型导出 + +StableDiffusionInpaint 的 onnx 导出非常简单，optimum 已经做好了集成，支持命令行直接导出，具体参考可参考optimum-cli： + +```shell +optimum-cli export onnx —task stable-diffusion —model stable-diffusion-inpainting stable-diffusion-inpainting-onnx +``` + +这样得到了四个 onnx 模型（unet、 vae encoder、decoder 和 text encoder）。 + +### tokenizer&scheduler + +与检测、分类等传统 cv 方法不同，我们如果想在 c++中串起整个 pipeline，还缺少 c++版本的 tokenizer 和 scheduler。有很多优秀的开源 C++版本的 tokenizer，这里我选用了tokenizers_cpp，地址：https://github.com/mlc-ai/tokenizers-cpp。tokenizers-cpp 接口简单，并且可直接使用 🤗hugging face 中开源的的 tokenizer.json 配置文件。 + +```cpp +auto tokenizer = Tokenizer::FromBlobJSON( + LoadBytesFromFile("./tokenizers/tokenizer.json")); +std::string startoftext = "<|startoftext|>"; +std::string endoftext = "<|endoftext|>"; +std::string prompt = startoftext + "a lovely girl" + endoftext; +std::vector text_input_ids = tokenizer->Encode(prompt); +``` + +而对于 scheduler，目前没找到很好用的 c++版本，所以作者实现了一个 C++版本的 ddim_scheduler，并做了开源ddim_scheduler_cpp，rep地址：https://github.com/TalkUHulk/ddim_scheduler_cpp。ddim_scheduler_cpp 底层基于 Eigen 实现，与 diffusers 接口保持一致，可直接替换。 + +```cpp +// init from json +auto scheduler = DDIMScheduler("scheduler_config.json"); + +// set num_inference_steps +scheduler.set_timesteps(10); + +// get timesteps +std::vector timesteps; +scheduler.get_timesteps(timesteps); + +// random init for example +std::vector sample(1 * 4 * 64 * 64); +std::vector model_output(1 * 4 * 64 * 64); + +for(int i = 0; i < 4 * 64 * 64; i++){ + sample[i] = distribution(generator); + model_output[i] = distribution(generator); +} + +// step +std::vector pred_sample; +for(auto t: timesteps){ + scheduler.step(model_output, {1, 4, 3, 3}, sample, {1, 4, 3, 3}, pred_sample, t); +} + +``` + +### C++推理 + +目前，我们将所有必须的 C++物料都集齐了。借助作者之前开源的一个开源工具AiDB（[rep地址](https://mp.weixin.qq.com/s/D3mj9Dj2nmqeUIJMy8BIag)），只需要简单配置，直接可以使用 C++加载并推理 onnx 模型。 + +```cpp +auto scheduler = Scheduler::DDIMScheduler("scheduler_config.json"); +auto tokenizer = Tokenizer::FromBlobJSON( + LoadBytesFromFile("tokenizer.json")); +std::string startoftext = "<|startoftext|>"; +std::string endoftext = "<|endoftext|>"; +std::string prompt = startoftext + "A cute cat" + endoftext; + +std::vector text_input_ids = tokenizer->Encode(prompt); + +std::string uncond_tokens = startoftext + "" + endoftext; + +std::vector uncond_input = tokenizer->Encode(uncond_tokens); + +auto text_enc = AIDB::Interpreter::createInstance("text_encoder", "onnx"); + +std::vector> prompt_embeds; +std::vector> prompt_embeds_shape; + +text_enc->forward(text_input_ids.data(), 77, 0, 0, prompt_embeds, prompt_embeds_shape); + +std::vector> negative_prompt_embeds; +std::vector> negative_prompt_embeds_shape; +text_enc->forward(uncond_input.data(), 77, 0, 0, negative_prompt_embeds, negative_prompt_embeds_shape); + +std::vector prompt_embeds_cat(2 * 77 * 768, 0); +memcpy(prompt_embeds_cat.data(), negative_prompt_embeds[0].data(), 77 * 768 * sizeof(float)); +memcpy(prompt_embeds_cat.data() + 77 * 768, prompt_embeds[0].data(), 77 * 768 * sizeof(float)); + +auto num_inference_steps = 10; +scheduler.set_timesteps(num_inference_steps); +std::vector timesteps; +scheduler.get_timesteps(timesteps); + +auto vae_enc = AIDB::Interpreter::createInstance("sd_inpaint_vae_encoder", "onnx"); +auto vae_dec = AIDB::Interpreter::createInstance("sd_inpaint_vae_decoder", "onnx"); +auto unet = AIDB::Interpreter::createInstance("sd_inpaint_unet", "onnx"); + +std::vector latents(1 * 4 * 64 * 64); + +AIDB::Utility::randn(latents.data(), latents.size()); + +auto image = cv::imread("dog.png"); +auto mask = cv::imread("dog_mask.png", 0); + +// 图像预处理 +int target = 512; +float src_ratio = float(image.cols) / float(image.rows); +float target_ratio = 1.0f; + +int n_w, n_h, pad_w = 0, pad_h = 0; +float _scale_h, _scale_w; + +if(src_ratio > target_ratio){ + n_w = target; + + n_h = floor(float(n_w) / float(image.cols) * float(image.rows) + 0.5f); + pad_h = target - n_h; + _scale_h = _scale_w = float(n_w) / float(image.cols); +} else if(src_ratio < target_ratio){ + n_h = target; + n_w = floor(float(n_h) / float(image.rows) * float(image.cols) + 0.5f); + pad_w = target - n_w; + _scale_h = _scale_w = float(n_h) / float(image.rows); +} else{ + n_w = target; + n_h = target; + _scale_h = _scale_w = float(n_w) / float(image.cols); +} + +cv::resize(image, image, cv::Size(n_w, n_h)); +cv::copyMakeBorder(image, image, 0, pad_h, 0, pad_w, cv::BORDER_CONSTANT, cv::Scalar(255, 255, 255)); + +cv::resize(mask, mask, cv::Size(n_w, n_h)); +cv::copyMakeBorder(mask, mask, 0, pad_h, 0, pad_w, cv::BORDER_CONSTANT, cv::Scalar(255, 255, 255)); + +cv::threshold(mask, mask, 127.5, 1, cv::THRESH_BINARY); + +cv::cvtColor(image, image, cv::COLOR_BGR2RGB); +image.convertTo(image, CV_32F); +image = image / 127.5 - 1.0; +cv::Mat mask_image = cv::Mat::zeros(image.rows, image.cols, CV_32FC3); +image.copyTo(mask_image, 1 - mask); + +cv::Mat blob; +cv::Mat blob_mask; +cv::dnn::blobFromImage(mask_image, blob); + +cv::dnn::blobFromImage(mask, blob_mask, 1.0f, cv::Size(64, 64)); + +std::vector> masked_image_latents; +std::vector> masked_image_latents_shape; + +vae_enc->forward(blob.data, 512, 512, 3, masked_image_latents, masked_image_latents_shape); + +auto scaling_factor = 0.18215f; +std::for_each(masked_image_latents.begin(), masked_image_latents.end(), + [=](std::vector& masked_image_latent) { + std::for_each(masked_image_latent.begin(), masked_image_latent.end(), [=](float &item){ item *= scaling_factor;}); + } +); + +auto init_noise_sigma = scheduler.get_init_noise_sigma(); +std::for_each(latents.begin(), latents.end(), [=](float &item){item*=init_noise_sigma;}); +auto guidance_scale = 7.5f; +int step = 0; +// 循环处理 +for(auto t: timesteps){ + auto tic = std::chrono::system_clock::now(); + + std::vector latent_model_input(2 * 9 * 64 * 64, 0); + memcpy(latent_model_input.data(), latents.data(), 4 * 64 * 64 * sizeof(float)); + memcpy(latent_model_input.data() + 4 * 64 * 64, blob_mask.data, 1 * 64 * 64 * sizeof(float)); + memcpy(latent_model_input.data() + 5 * 64 * 64, masked_image_latents[0].data(), 4 * 64 * 64 * sizeof(float)); + memcpy(latent_model_input.data() + 9 * 64 * 64, latent_model_input.data(), 9 * 64 * 64 * sizeof(float)); + + std::vector> noise_preds; + std::vector> noise_preds_shape; + std::vector input; + std::vector> input_shape; + + input.push_back(latent_model_input.data()); + input_shape.push_back({2, 9, 64, 64}); + + std::vector timestep = {(long long)t}; + input.push_back(timestep.data()); + input_shape.push_back({1}); + + input.push_back(prompt_embeds_cat.data()); + input_shape.push_back({2, 77, 768}); + + unet->forward(input, input_shape, noise_preds, noise_preds_shape); + + // noise_preds [2,4,64,64] noise_pred_uncond | noise_pred_text + std::vector noise_pred(1 * 4 * 64 * 64, 0); + for(int i = 0; i < noise_pred.size(); i++){ + noise_pred[i] = noise_preds[0][i] + guidance_scale * (noise_preds[0][i + 4 * 64 * 64] - noise_preds[0][i]); + } + std::vector pred_sample; + scheduler.step(noise_pred, {1, 4, 64, 64}, latents, {1, 4, 64, 64}, pred_sample, t); + latents.clear(); + latents.assign(pred_sample.begin(), pred_sample.end()); + + auto toc = std::chrono::system_clock::now(); + std::chrono::duration elapsed = toc - tic; + // 计算并输出进度百分比 + std::cout << "\rStep " << step++ << " " <' << ']'; + std::flush(std::cout); + } + +std::for_each(latents.begin(), latents.end(), [=](float &item){item /= scaling_factor;}); + +std::vector> sample; +std::vector> sample_shape; +vae_dec->forward(latents.data(), 64, 64, 4, sample, sample_shape); +cv::Mat sd_image; +AIDB::Utility::stable_diffusion_process(sample[0].data(), sample_shape[0][2], sample_shape[0][3], sample_shape[0][1], sd_image); + +cv::imwrite("stable_diffusion_inpainting.jpg", sd_image); + +``` + +![](https://files.mdnice.com/user/48619/93f36c70-2854-4f6a-ba16-057668aef724.jpg) + +至此，我们已经成功搭建起 stablediffusioninpaint 的 C++ pipeline，但更常用、更有趣的是 controlnet 和 lora 与 stablediffusion 的结合。下面我们尝试搭建 StableDiffusionControlNetImg2ImgPipeline 的 C++推理代码，并支持 LoRa 加载。 + +## StableDiffusionControlNetImg2ImgPipeline + +### 模型导出 + +目前 optimum 还未提供 stablediffusion + controlnet +LoRa 的 onnx 模型导出选项，所以这里我们先将模型导出。 +这里我们有两种导出方案，分别导出 controlNet 和 Unet，以及将二者合并为一个模型。 + +![](https://files.mdnice.com/user/48619/b92dc0ac-48fe-45a5-a052-258b5b09c324.png) + +先看一下 controlNet 的整体架构，controlNet 和 Unet 的耦合比较深，如果我们分开导出，两个模型的输出和输入数量都会非常多，比入 Unet 部分有 down_block_res_sample_0 ～ down_block_res_sample_11、mid_block_res_sample 等 16 个输入，这样在写 inference 代码的时候就会比较繁琐。所以我们选择将两个模型合并为一个。但这样也有有另一个问题，比如我首先使用 controlNet-canny 导出了 onnx 模型，同时又想使用 controlNet-hed，那 unet 部分是不是要重复导出？这里有几个方法解决，我们后面再说明。 + +此处使用Yoji Shinkawa Style LoRA（🤗 https://civitai.com/models/12324/yoji-shinkawa-style-lora) + +导出代码： + +```python +from diffusers import ( + ControlNetModel, + StableDiffusionControlNetImg2ImgPipeline, +) +from diffusers.models.attention_processor import AttnProcessor + +is_torch_less_than_1_11 = version.parse(version.parse(torch.__version__).base_version) < version.parse("1.11") +is_torch_2_0_1 = version.parse(version.parse(torch.__version__).base_version) == version.parse("2.0.1") + +class UNetControlNetModel(torch.nn.Module): + def __init__( + self, + unet, + controlnet: ControlNetModel, + ): + super().__init__() + self.unet = unet + self.controlnet = controlnet + + def forward( + self, + sample, + timestep, + encoder_hidden_states, + controlnet_cond, + conditioning_scale, + ): + for i, (_controlnet_cond, _conditioning_scale) in enumerate( + zip(controlnet_cond, conditioning_scale) + ): + down_samples, mid_sample = self.controlnet( + sample, + timestep, + encoder_hidden_states=encoder_hidden_states, + controlnet_cond=_controlnet_cond, + conditioning_scale=_conditioning_scale, + return_dict=False, + ) + + # merge samples + if i == 0: + down_block_res_samples, mid_block_res_sample = down_samples, mid_sample + else: + down_block_res_samples = [ + samples_prev + samples_curr + for samples_prev, samples_curr in zip(down_block_res_samples, down_samples) + ] + mid_block_res_sample += mid_sample + + noise_pred = self.unet( + sample, + timestep, + encoder_hidden_states=encoder_hidden_states, + down_block_additional_residuals=down_block_res_samples, + mid_block_additional_residual=mid_block_res_sample, + return_dict=False, + )[0] + return noise_pred + +def onnx_export( + model, + model_args: tuple, + output_path: Path, + ordered_input_names, + output_names, + dynamic_axes, + opset, + use_external_data_format=False, +): + output_path.parent.mkdir(parents=True, exist_ok=True) + with torch.inference_mode(), torch.autocast("cuda"): + if is_torch_less_than_1_11: + export( + model, + model_args, + f=output_path.as_posix(), + input_names=ordered_input_names, + output_names=output_names, + dynamic_axes=dynamic_axes, + do_constant_folding=True, + use_external_data_format=use_external_data_format, + enable_onnx_checker=True, + opset_version=opset, + ) + else: + export( + model, + model_args, + f=output_path.as_posix(), + input_names=ordered_input_names, + output_names=output_names, + dynamic_axes=dynamic_axes, + do_constant_folding=True, + opset_version=opset, + ) + +with torch.no_grad(): + dtype = torch.float32 + device = "cpu" + # init controlnet + controlnet = ControlNetModel.from_pretrained("sd-controlnet-canny", torch_dtype=dtype).to(device) + if is_torch_2_0_1: + controlnet.set_attn_processor(AttnProcessor()) + + pipeline = StableDiffusionControlNetImg2ImgPipeline.from_pretrained( + "stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=dtype, safety_checker=None + ).to(device) + + pipeline.load_lora_weights("stable-diffusion-v1-5/LoRa/", "shinkawa_youji_offset") + output_path = Path("exp_lora") + if is_torch_2_0_1: + pipeline.unet.set_attn_processor(AttnProcessor()) + pipeline.vae.set_attn_processor(AttnProcessor()) + + # # TEXT ENCODER + num_tokens = pipeline.text_encoder.config.max_position_embeddings + text_hidden_size = pipeline.text_encoder.config.hidden_size + text_input = pipeline.tokenizer( + "A sample prompt", + padding="max_length", + max_length=pipeline.tokenizer.model_max_length, + truncation=True, + return_tensors="pt", + ) + onnx_export( + pipeline.text_encoder, + # casting to torch.int32 until the CLIP fix is released: https://github.com/huggingface/transformers/pull/18515/files + model_args=(text_input.input_ids.to(device=device, dtype=torch.int32)), + output_path=output_path / "text_encoder" / "model.onnx", + ordered_input_names=["input_ids"], + output_names=["last_hidden_state", "pooler_output"], + dynamic_axes={ + "input_ids": {0: "batch", 1: "sequence"}, + }, + opset=14, + ) + del pipeline.text_encoder + + ## VAE ENCODER + vae_encoder = pipeline.vae + vae_in_channels = vae_encoder.config.in_channels + vae_sample_size = vae_encoder.config.sample_size + # need to get the raw tensor output (sample) from the encoder + vae_encoder.forward = lambda sample: vae_encoder.encode(sample).latent_dist.sample() + + onnx_export( + vae_encoder, + model_args=(torch.randn(1, vae_in_channels, vae_sample_size, vae_sample_size).to(device=device, dtype=dtype),), + output_path=output_path / "vae_encoder" / "model.onnx", + ordered_input_names=["sample"], + output_names=["latent_sample"], + dynamic_axes={ + "sample": {0: "batch", 1: "channels", 2: "height", 3: "width"}, + }, + opset=14, + ) + + # # UNET + unet_controlnet = UNet2DConditionControlNetModel(pipeline.unet, controlnet) + unet_in_channels = pipeline.unet.config.in_channels + unet_sample_size = pipeline.unet.config.sample_size + num_tokens = pipeline.text_encoder.config.max_position_embeddings + text_hidden_size = pipeline.text_encoder.config.hidden_size + + img_size = 8 * unet_sample_size + unet_path = output_path / "unet" / "model.onnx" + + onnx_export( + unet_controlnet, + model_args=( + torch.randn(2, unet_in_channels, unet_sample_size, unet_sample_size).to(device=device, dtype=dtype), + torch.tensor([1.0]).to(device=device, dtype=dtype), + torch.randn(2, num_tokens, text_hidden_size).to(device=device, dtype=dtype), + torch.randn(2, 3, img_size, img_size).to(device=device, dtype=dtype), + torch.tensor([1.0]).to(device=device, dtype=dtype), + ), + output_path=unet_path, + ordered_input_names=[ + "sample", + "timestep", + "encoder_hidden_states", + "controlnet_cond", + "conditioning_scale", + ], + output_names=["noise_pred"], # has to be different from "sample" for correct tracing + dynamic_axes={ + "sample": {0: "2B", 2: "H", 3: "W"}, + "encoder_hidden_states": {0: "2B"}, + "controlnet_cond": {0: "2B", 2: "8H", 3: "8W"}, + }, + opset=14, + use_external_data_format=True, # UNet is > 2GB, so the weights need to be split + ) + + unet_model_path = str(unet_path.absolute().as_posix()) + unet_opt_graph = onnx.load(unet_model_path) + onnx.save_model( + unet_opt_graph, + unet_model_path, + save_as_external_data=True, + all_tensors_to_one_file=True, + location="model.onnx_data", + convert_attribute=False, + ) + del pipeline.unet + + # VAE DECODER + vae_decoder = pipeline.vae + vae_latent_channels = vae_decoder.config.latent_channels + # forward only through the decoder part + vae_decoder.forward = vae_encoder.decode + onnx_export( + vae_decoder, + model_args=( + torch.randn(1, vae_latent_channels, unet_sample_size, unet_sample_size).to(device=device, dtype=dtype), + ), + output_path=output_path / "vae_decoder" / "model.onnx", + ordered_input_names=["latent_sample"], + output_names=["sample"], + dynamic_axes={ + "latent_sample": {0: "batch", 1: "channels", 2: "height", 3: "width"}, + }, + opset=14, + ) + +``` + +这里有几个点需要注意。 + +#### OP 问题 + +pytorch2.0 以上，需要做以下设置才可以成功导出 + +```python +pipeline.unet.set_attn_processor(AttnProcessor()) +pipeline.vae.set_attn_processor(AttnProcessor()) +controlnet.set_attn_processor(AttnProcessor()) +``` + +具体可以参考 diffusers->models->attention_processor.py 中的相关代码。Pytorch2.0 以上 scaled dot-product attention 计算会默认使用torch.nn.functional.scaled_dot_product_attention，而 onnx 导出时不支持该 OP。 + +``` +torch.onnx.errors.UnsupportedOperatorError: Exporting the operator 'aten::scaled_dot_product_attention' to ONNX opset version 14 is not supported. +``` + +因此需要做替换，diffusers 很贴心的把相关代码实现好，我们直接使用即可。 + +#### 模型大小>2GB + +ONNX 模型本质就是一个 Protobuf 序列化后的二进制文件，而 Protobuf 的文件大小限制为 2GB。因此对于 Unet 相关模型来说，存储大小已经超过了限制。onnx 给出的方案是单独存储 weights、bias 这些权重。 这里做下详细说明。 +先来看下onnx.proto(文件地址：https://github.com/onnx/onnx/blob/main/onnx/onnx.proto)中的定义： + +``` +message TensorProto { +.... +repeated int64 dims = 1; +optional int32 data_type = 2; +.... +// Data can be stored inside the protobuf file using type-specific fields or raw_data. +// Alternatively, raw bytes data can be stored in an external file, using the external_data field. +// external_data stores key-value pairs describing data location. Recognized keys are: +// - "location" (required) - POSIX filesystem path relative to the directory where the ONNX +// protobuf model was stored +// - "offset" (optional) - position of byte at which stored data begins. Integer stored as string. +// Offset values SHOULD be multiples 4096 (page size) to enable mmap support. +// - "length" (optional) - number of bytes containing data. Integer stored as string. +// - "checksum" (optional) - SHA1 digest of file specified in under 'location' key. +repeated StringStringEntryProto external_data = 13; + +// Location of the data for this tensor. MUST be one of: +// - DEFAULT - data stored inside the protobuf message. Data is stored in raw_data (if set) otherwise in type-specified field. +// - EXTERNAL - data stored in an external location as described by external_data field. +enum DataLocation { + DEFAULT = 0; + EXTERNAL = 1; +} +// If value not set, data is stored in raw_data (if set) otherwise in type-specified field. +optional DataLocation data_location = 14; +} +``` + +我们可以通过 data_location 来判断某个参数的位置，然后读取 external_data 参数加载权重，接下来我们在代码中手动加载： + +```python +import onnx +from onnx.external_data_helper import _get_all_tensors + +onnx_model = onnx.load("unet/model.onnx", load_external_data=False) +for tensor in _get_all_tensors(onnx_model): + if tensor.HasField("data_location") and tensor.data_location == onnx.TensorProto.EXTERNAL: + info = {} + for item in tensor.external_data: + info[item.key] = item.value + with open(info["location"], "rb") as data_file: + data_file.seek(info["offset"]) + tensor.raw_data = data_file.read(info["length"]) + +``` + +摘出其中一个 tensor 的 external_data 详细说明： + +``` +[ key: "location" + value: "model.onnx_data" + , key: "offset" + value: "0" + , key: "length" + value: "46080" +] +``` + +location 记录了权重存储的文件名，offset 是该权重在文件中的偏移量，length 是权重的长度。有了以上信息，onnx 内部就可以直接 load 权重，解决 2GB 限制问题。 + +仔细的同学会观察到，导出的 uent 目录下有，除了.onnx 模型，还有非常非常多的 weight/bias 等文件。这其实就是每一个权重数据。如此碎片化，我们使用或者版本管理起来非常不方便。我们使用以下代码，将所有的权重合并到一个文件中： + +```python +unet_opt_graph = onnx.load(unet_model_path) +onnx.save_model( + unet_opt_graph, + unet_model_path, + save_as_external_data=True, + all_tensors_to_one_file=True, + location="model.onnx_data", + convert_attribute=False, +) +``` + +这样所有的权重就会保存到一个 model.onnx_data 文件里。 + +### C++推理 + +与上文类似，借助 AiDB，使用 C++串起整个 pipeline + +```cpp +std::vector noise(1 * 4 * 64 * 64); +AIDB::Utility::randn(noise.data(), noise.size()); + +auto strength = 1.0f; // 0~1 之间重绘比例。越低越接近输入图片。 +auto scheduler = Scheduler::DDIMScheduler("scheduler_config.json"); +auto scaling_factor = 0.18215f; +auto tokenizer = Tokenizer::FromBlobJSON( + LoadBytesFromFile("tokenizer.json")); +std::string startoftext = "<|startoftext|>"; +std::string endoftext = "<|endoftext|>"; +std::string trigger = argv[1]; +std::string prompt = startoftext + trigger + endoftext; +std::vector text_input_ids = tokenizer->Encode(prompt); + +std::string uncond_tokens = startoftext + "longbody, lowres, cropped, worst quality, low quality, multiple people" + endoftext; + +std::vector uncond_input = tokenizer->Encode(uncond_tokens); +auto text_enc = AIDB::Interpreter::createInstance("text_encoder", "onnx"); + +std::vector> prompt_embeds; +std::vector> prompt_embeds_shape; + +text_enc->forward(text_input_ids.data(), 77, 0, 0, prompt_embeds, prompt_embeds_shape); + +std::vector> negative_prompt_embeds; +std::vector> negative_prompt_embeds_shape; +text_enc->forward(uncond_input.data(), 77, 0, 0, negative_prompt_embeds, negative_prompt_embeds_shape); + +std::vector prompt_embeds_cat(2 * 77 * 768, 0); +memcpy(prompt_embeds_cat.data(), negative_prompt_embeds[0].data(), 77 * 768 * sizeof(float)); +memcpy(prompt_embeds_cat.data() + 77 * 768, prompt_embeds[0].data(), 77 * 768 * sizeof(float)); + +auto num_inference_steps = 10; +scheduler.set_timesteps(num_inference_steps); +std::vector timesteps; +scheduler.get_timesteps(timesteps); + +// Figuring initial time step based on strength +auto init_timestep = min(int(num_inference_steps * strength), num_inference_steps); +auto t_start = max(num_inference_steps - init_timestep, 0); + +timesteps.assign(timesteps.begin() + t_start, timesteps.end()); + +num_inference_steps = timesteps.size(); + +auto image = cv::imread("portrait.png"); +int target = 512; +float src_ratio = float(image.cols) / float(image.rows); +float target_ratio = 1.0f; + +int n_w, n_h, pad_w = 0, pad_h = 0; +float _scale_h, _scale_w; + +if(src_ratio > target_ratio){ + n_w = target; + + n_h = floor(float(n_w) / float(image.cols) * float(image.rows) + 0.5f); + pad_h = target - n_h; + _scale_h = _scale_w = float(n_w) / float(image.cols); +} else if(src_ratio < target_ratio){ + n_h = target; + n_w = floor(float(n_h) / float(image.rows) * float(image.cols) + 0.5f); + pad_w = target - n_w; + _scale_h = _scale_w = float(n_h) / float(image.rows); +} else{ + n_w = target; + n_h = target; + _scale_h = _scale_w = float(n_w) / float(image.cols); +} + +cv::resize(image, image, cv::Size(n_w, n_h)); +cv::copyMakeBorder(image, image, 0, pad_h, 0, pad_w, cv::BORDER_CONSTANT, cv::Scalar(255, 255, 255)); + +auto low_threshold = 150; +auto high_threshold = 200; +cv::Mat canny; +cv::Canny(image, canny, low_threshold, high_threshold); + +std::vector bgr_channels{canny, canny, canny}; +cv::merge(bgr_channels, canny); + +image.convertTo(image, CV_32F); +image = image / 127.5 - 1.0; +cv::Mat blob; +cv::dnn::blobFromImage(image, blob); + +canny.convertTo(canny, CV_32F); +cv::Mat blob_canny; +cv::dnn::blobFromImage(canny, blob_canny, 1.0f / 255.0f); + +auto vae_enc = AIDB::Interpreter::createInstance("sd_vae_encoder_with_controlnet", "onnx"); + +auto vae_dec = AIDB::Interpreter::createInstance("sd_vae_decoder_with_controlnet", "onnx"); + +auto unet = AIDB::Interpreter::createInstance2("sd_unet_with_controlnet_with_lora", "shinkawa", "onnx"); + +// Prepare latent variables +std::vector> image_latents; +std::vector> image_latents_shape; + +vae_enc->forward(blob.data, 512, 512, 3, image_latents, image_latents_shape); + +auto latents = image_latents[0]; + + +std::for_each(latents.begin(), latents.end(), [=](float &item){item *= scaling_factor;}); + +auto latent_timestep = timesteps[0]; +std::vector init_latents; +scheduler.add_noise(latents, {1, 4, 64, 64}, noise, {1, 4, 64, 64}, latent_timestep, init_latents); + +auto guidance_scale = 7.5f; + +std::vector controlnet_keep(timesteps.size(), 1.0); +float controlnet_conditioning_scale = 0.5f; // 0~1 之间的 ControlNet 约束比例。越高越贴近约束。 + +int step = 0; +for(auto t: timesteps){ + auto tic = std::chrono::system_clock::now(); + double cond_scale = controlnet_conditioning_scale * controlnet_keep[step]; + std::vector latent_model_input(2 * 4 * 64 * 64, 0); + memcpy(latent_model_input.data(), init_latents.data(), 4 * 64 * 64 * sizeof(float)); + memcpy(latent_model_input.data() + 4 * 64 * 64, init_latents.data(), 4 * 64 * 64 * sizeof(float)); + + + std::vector> down_and_mid_blok_samples; + std::vector> down_and_mid_blok_samples_shape; + std::vector input; + std::vector> input_shape; + + // sample + input.push_back(latent_model_input.data()); + input_shape.push_back({2, 4, 64, 64}); + + // t ✅ + std::vector timestep = {(float)t}; + input.push_back(timestep.data()); + input_shape.push_back({1}); + + // encoder_hidden_states ✅ + input.push_back(prompt_embeds_cat.data()); + input_shape.push_back({2, 77, 768}); + + std::vector controlnet_cond(2 * 3 * 512 * 512, 0); + memcpy(controlnet_cond.data(), blob_canny.data, 3 * 512 * 512 * sizeof(float)); + memcpy(controlnet_cond.data() + 3 * 512 * 512, blob_canny.data, 3 * 512 * 512 * sizeof(float)); + + // controlnet_cond ✅ + input.push_back(controlnet_cond.data()); + input_shape.push_back({2, 3, 512, 512}); + + // conditioning_scale ✅ + std::vector cond_scales = {(float)(cond_scale)}; + input.push_back(cond_scales.data()); + input_shape.push_back({1}); + + std::vector> noise_preds; + std::vector> noise_preds_shape; + unet->forward(input, input_shape, noise_preds, noise_preds_shape); + + // noise_preds [2,4,64,64] noise_pred_uncond | noise_pred_text + std::vector noise_pred(1 * 4 * 64 * 64, 0); + for(int i = 0; i < noise_pred.size(); i++){ + noise_pred[i] = noise_preds[0][i] + guidance_scale * (noise_preds[0][i + 4 * 64 * 64] - noise_preds[0][i]); + } + std::vector pred_sample; + scheduler.step(noise_pred, {1, 4, 64, 64}, init_latents, {1, 4, 64, 64}, pred_sample, t); + init_latents.clear(); + init_latents.assign(pred_sample.begin(), pred_sample.end()); + auto toc = std::chrono::system_clock::now(); + std::chrono::duration elapsed = toc - tic; + std::cout << "\rStep " << step++ << " " <' << ']'; + std::flush(std::cout); +} +std::for_each(init_latents.begin(), init_latents.end(), [=](float &item){item /= scaling_factor;}); + +std::vector> sample; +std::vector> sample_shape; +vae_dec->forward(init_latents.data(), 64, 64, 4, sample, sample_shape); + +cv::Mat sd_image(sample_shape[0][2], sample_shape[0][3], CV_8UC3); +AIDB::Utility::stable_diffusion_process(sample[0].data(), sample_shape[0][2], sample_shape[0][3], sample_shape[0][1], sd_image); + +cv::imwrite("stable_diffusion_controlnet_img2img_" + trigger + ".jpg", sd_image); +``` + +![](https://files.mdnice.com/user/48619/2eaef51b-02ec-44d5-a255-9116fcc97555.jpeg) + +### LoRA 方式加载 + +回到上文提到的问题，以上例子使用 controlNet-canny 导出 onnx 模型，如果我们又想使用 controlNet-hed，或者使用更多的 LoRa 呢？是否一定必须重新导出整个模型， +是否可以用“LoRa”的方式加载模型呢。答案是肯定的，查看 onnruntime 的接口，官方提供了如下接口: + +```cpp +SessionOptions& SessionOptions::AddExternalInitializers(const std::vector& names,const std::vector& ort_values) +``` + +利用此接口，我们可以实现“LoRa 方式”的模型加载。这里以“LoRa”举例，controlNet 同理。 + +先做一点简单的知识储备，ONNX 模型本质就是一个 Protobuf 序列化后的二进制文件，所以理论上我们可以做任意合理的修改。根据 onnx.proto 的定义，首先来看一下 onnx 模型的结构。 onnx 主要包含以下几个类：ModelProto，NodeProto，GraphProto，TensorProto 等。ModelProto 作为 top-level 类，用于绑定 ML 模型并将其计算图与元数据相关联。NodeProto 用来描述了 graph 中的 node。TensorProto 则用来组织 tensor 的具体信息。所以 onnx 的结构大概可以用下图表示： + +![](https://files.mdnice.com/user/48619/6dfee3a7-5d77-4350-b06f-9e4f5d5e94d0.png) + +这样我们就有了一个大概的思路，读取 LoRa 模型，解析 LoRa 模型中 tensor，因为网络结构都是相同的，我们直接通过 onnxruntime 的 AddExternalInitializers 接口，来替换原始网络中的 LoRa 部分。 + +#### onnx 模型读取 + +使用 protobuf 读取 onnx 模型，而不是使用 ort： + +```cpp +std::ifstream fin(lora_path, std::ios::in | std::ios::binary); +onnx::ModelProto onnx_model; +onnx_model.ParseFromIstream(&fin); + +auto graph = onnx_model.graph(); +const auto& initializer = graph.initializer(); +std::vector init_names; +std::vector initializer_data; +for(auto& tensor: initializer){ + init_names.push_back(tensor.name()); + tensors.emplace_back(tensor); +} + +fin.close(); +``` + +#### OP名称 + +原始模型与 onnx 导出的模型的名字是不一致的，我们需要找到映射关系，才能正确加载。 +首先加载 🤗safetensors 格式的模型 + +```python +from safetensors.torch import load_file +state_dict = load_file("LoRa.safetensors") +``` + +此时 state_dict 中的 key 并不是模型 onnx 导出前的 key，这里需要做一个转换。直接参考 diffusers 的代码： + +```python +from diffusers.loaders.lora_conversion_utils import _convert_kohya_lora_to_diffusers, _maybe_map_sgm_blocks_to_diffusers +from diffusers.utils.state_dict_utils import convert_unet_state_dict_to_peft, convert_state_dict_to_diffusers + +indent = " " +unet_config = None +# Map SDXL blocks correctly. +if unet_config is not None: + # use unet config to remap block numbers + state_dict = _maybe_map_sgm_blocks_to_diffusers(state_dict, unet_config) +state_dict, network_alphas = _convert_kohya_lora_to_diffusers(state_dict) + +keys = list(state_dict.keys()) + +unet_name = "unet" +text_encoder_name = "text_encoder" + +unet_keys = [k for k in keys if k.startswith(unet_name)] +unet_state_dict = {k: v for k, v in state_dict.items() if k in unet_keys} + +unet_state_dict = convert_unet_state_dict_to_peft(unet_state_dict) + +text_encoder_lora_state_dict = {} +if any(text_encoder_name in key for key in keys): + text_encoder_keys = [k for k in keys if k.startswith(text_encoder_name) and k.split(".")[0] == text_encoder_name] + text_encoder_lora_state_dict = { + k.replace(f"{text_encoder_name}.", ""): v for k, v in state_dict.items() if k in text_encoder_keys + } + + text_encoder_lora_state_dict = convert_state_dict_to_diffusers(text_encoder_lora_state_dict) + + +unet_lora_state_dict = {} +for key, value in unet_state_dict.items(): + name = key + if "lora_A" in key: + name = key.replace("lora_A", "lora_layer.down") + elif "lora_B" in key: + name = key.replace("lora_B", "lora_layer.up") + + unet_lora_state_dict[name] = value +``` + +执行以上代码，可以得到 torch.onnx.export 前模型的 key:value。接下来就是和 onnx 模型中的 name 找到对应关系。 + +其实 onnx 模型中已经储存了对应的对应关系，我们使用以下代码先观察下 onnx 模型中村了什么信息(这里只输出了 lora 相关的): + +```python +onnx_model = onnx.load("unet.onnx", load_external_data=False) +for node in onnx_model.graph.node: + print(onnx.helper.printable_node(node, indent, subgraphs=True)) + +``` + +部分输出： + +![](https://files.mdnice.com/user/48619/494336eb-0a38-44a2-b6da-7ed875b64632.png) + +可以看到每个node的对应关系，格式如下torch-op-name = OP(param, onnx-tensor-name)。按照以上规则，可以找到两种模型opname的映射，将这种关系保存下来： + +```python +unet_mapping = {} +for node in onnx_model.graph.node: + contents = onnx.helper.printable_node(node, indent, subgraphs=True)[0] + content = re.findall(r'%.*?(?= |,)|%.*?(?=\))', contents) # torch-op-name = OP(形参， onnx-weight-name) + if ".".join(content[0].split("/")[1:-1] + ["weight"]) in unet_lora_state_dict: + unet_mapping[".".join(content[0].split("/")[1:-1] + ["weight"])] = content[-1][1:] +``` + +#### LoRa保存 + + 最后就是如何组织新的LoRa模型了。这里为了方便，我们构造一个“假的”onnx模型，仅仅存储LoRa的权重，name以上一节映射后为准。 + +```python +initializer = [] +for key, value in unet_lora_state_dict.items(): + initializer.append( + helper.make_tensor( + name=unet_mapping[key], + data_type=helper.TensorProto.DataType.FLOAT, + dims=value.T.shape if value.ndim == 2 else value.shape, + vals=value.T.float().numpy().astype(np.float32).tobytes() + if value.ndim == 2 else value.float().numpy().astype(np.float32).tobytes(), + raw=True + ) + ) + +graph = helper.make_graph( + name=LoRa, + inputs=[], + outputs=[], + nodes=[], + initializer=initializer +) + +opset = [ + helper.make_operatorsetid(LoRa, 14) +] + +model = helper.make_model(graph, opset_imports=opset) +onnx.save_model(model, "LoRa.lora") +``` + +#### LoRa校验 + + 以上3步已经得到了新的模型，但为了确认我们的方式是否正确，我们拿一个已经导出的Unet模型和对应的LoRa权重做一下校验 + +```python +import onnx +from onnx.external_data_helper import load_external_data_for_model, _get_all_tensors +import numpy as np + + +onnx_model1 = onnx.load("unet-model.onnx", load_external_data=True) +onnx_model2 = onnx.load("lora.lora") + +lora_state_dict = {} +for t in _get_all_tensors(onnx_model2): + lora_state_dict[t.name] = np.frombuffer(t.raw_data, dtype=np.float32) + + +for t in _get_all_tensors(onnx_model1): + if t.name in lora_state_dict: + np.testing.assert_almost_equal(np.frombuffer(t.raw_data, dtype=np.float32), + lora_state_dict[t.name], + decimal=6) +``` + +确认没问题，我们的准备工作也算完成。下面完成C++代码部分。 + +#### LoRa加载 + + 读取新的LoRa模型，将权重的name和raw_data读取出来，然后创建对应的tensor，最后调用session_options.AddExternalInitializers一起初始化即可。需要注意的是，onnxruntime的CreateTensor操作是浅拷贝，所以在写法上注意局部变量的生存周期。 + +```cpp +if(!param._lora_path.empty()){ + std::vector init_names; + std::vector initializer_data; + std::vector tensors; + auto allocator = Ort::AllocatorWithDefaultOptions(); + + for(const auto& lora_path: param._lora_path){ + std::ifstream fin(lora_path, std::ios::in | std::ios::binary); + onnx::ModelProto onnx_model; + onnx_model.ParseFromIstream(&fin); + + auto graph = onnx_model.graph(); + const auto& initializer = graph.initializer(); + + for(auto& tensor: initializer){ + init_names.push_back(tensor.name()); + tensors.emplace_back(tensor); + } + fin.close(); + } + for(const auto& tensor: tensors){ + std::vector shape(tensor.dims_size(), 0); + for(int i = 0; i < tensor.dims_size(); i++){ + shape[i] = tensor.dims(i); + } + Ort::Value input_tensor = Ort::Value::CreateTensor(allocator.GetInfo(), + (void *)(tensor.raw_data().c_str()), + tensor.raw_data().length(), + shape.data(), + shape.size(), + ONNXTensorElementDataType(tensor.data_type())); + + initializer_data.push_back(std::move(input_tensor)); + } + _session_options.AddExternalInitializers(init_names, initializer_data); + _session = std::make_shared(_env, param._model_path.c_str(), _session_options); +} +``` + +作者在C站找了几个相同结构的LoRa，分别为blindbox、mix4和moxin，测试一下效果 + + +![](https://files.mdnice.com/user/48619/daa213c9-b9a9-434f-b296-0945e267fedb.jpg) + + + +> 以上代码和模型都已开源，更多详情，敬请登陆github，欢迎Star。 diff --git a/mkdocs.yml b/mkdocs.yml index f58455a..87dbf08 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -95,6 +95,7 @@ nav: - GQA，MLA之外的另一种KV Cache压缩方式：动态内存压缩（DMC）: academic/算法科普/Transformer/GQA，MLA之外的另一种KV Cache压缩方式：动态内存压缩（DMC）.md - OverlapMamba 具备超强泛化能力的定位方法: academic/算法科普/Transformer/OverlapMamba 具备超强泛化能力的定位方法.md - vAttention：用于在没有Paged Attention的情况下Serving LLM: academic/算法科普/Transformer/vAttention：用于在没有Paged Attention的情况下Serving LLM.md + - 大模型KV Cache节省神器MLA学习笔记（包含推理时的矩阵吸收分析）: academic/算法科普/Transformer/大模型KV Cache节省神器MLA学习笔记（包含推理时的矩阵吸收分析）.md - 半监督&无监督（质量较低）: - WSCL: academic/算法科普/半监督&无监督/WSCL.md - 弱监督方法在森林病虫害检测中的应用: academic/算法科普/半监督&无监督/弱监督方法在森林病虫害检测中的应用.md @@ -690,6 +691,7 @@ nav: - Windows 上基于 TensorRT 的 YOLOV6 部署保姆级教程: project/部署优化/AI 移动端算法优化/Windows 上基于 TensorRT 的 YOLOV6 部署保姆级教程.md - 用于ARM Cortex-M系列的芯片的神经网络推理库CMSIS-NN详解: project/部署优化/AI 移动端算法优化/用于ARM Cortex-M系列的芯片的神经网络推理库CMSIS-NN详解.md - NEON做色域变化_ 用单核性能无限逼近八核并行OpenCV: project/部署优化/AI 移动端算法优化/NEON做色域变化_ 用单核性能无限逼近八核并行OpenCV.md + - 60行代码加速20倍_NEON实现深度学习OD任务后处理绘框: project/部署优化/AI 移动端算法优化/60行代码加速20倍_NEON实现深度学习OD任务后处理绘框 - 深度学习编译器（质量较高）: - 多面体模型在深度学习编译器的应用: project/部署优化/深度学习编译器/多面体模型在深度学习编译器的应用.md - 【从零开始学深度学习编译器】一，深度学习编译器及TVM介绍: project/部署优化/深度学习编译器/【从零开始学深度学习编译器】一，深度学习编译器及TVM介绍.md @@ -793,6 +795,7 @@ nav: - 大模型部署框架 FastLLM 实现细节解析: project/部署优化/AI 部署及其它优化算法/大模型部署框架 FastLLM 实现细节解析.md - AiDB一个集合了6大推理框架的AI工具箱加速你的模型部署: project/部署优化/AI 部署及其它优化算法/AiDB一个集合了6大推理框架的AI工具箱加速你的模型部署.md - nndeploy：一款最新上线的支持多平台、简单易用、高性能的机器学习部署框架: project/部署优化/AI 部署及其它优化算法/nndeploy：一款最新上线的支持多平台、简单易用、高性能的机器学习部署框架.md + - 如何使用“LoRa”的方式加载Onnx模型：StableDiffusion相关模型 的C++推理: project/部署优化/AI 部署及其它优化算法/如何使用“LoRa”的方式加载Onnx模型：StableDiffusion相关模型 的C++推理.md - csapp-第九章-虚拟内存-读书笔记: project/部署优化/AI 部署及其它优化算法/csapp-第九章-虚拟内存-读书笔记.md - csapp 第八章 异常控制流 读书笔记: project/部署优化/AI 部署及其它优化算法/csapp 第八章 异常控制流 读书笔记.md - RESOURCES: