1. 项目概述解码Transformer的隐秘角落第一次接触Transformer架构时我像大多数人一样被self-attention的优雅设计所震撼。但在实际构建多模态模型的过程中那些论文里一笔带过的实现细节却成了最折磨人的暗坑。比如Decoder的掩码机制在训练和推理时的微妙差异或是位置编码与视觉特征融合时的维度陷阱——这些在教科书式的教程里往往被简化处理却直接影响着模型的实际表现。经过三个跨模态项目的实战洗礼包括一个百万级参数的图文生成系统我逐渐意识到Transformer的强大不仅在于宏观架构的创新更在于那些容易被忽视的工程细节。本文将聚焦Decoder工作机制、掩码的时空控制以及训练/推理的模式差异这三个关键维度结合PyTorch代码实例揭示那些只可意会的实现要点。2. 核心需求解析为什么这些细节至关重要2.1 多模态场景的特殊挑战当Transformer从纯文本领域扩展到视觉-语言多模态任务时Decoder需要同时处理两种截然不同的数据流来自视觉Backbone的网格特征如CNN或ViT输出的patch embeddings和文本的词嵌入序列。这两种特征在时空维度上的不对齐性使得标准的Decoder实现需要做出以下关键调整跨模态注意力权重分配视觉特征通常具有更高的空间冗余度需要在计算query-key乘积时引入适当的温度系数调节序列长度动态变化图像patch数量固定而文本token长度可变导致位置编码需要特殊处理掩码的复合作用既要防止文本解码时的未来信息泄露又要控制视觉特征参与的时空范围2.2 工业级部署的实际需求在真实生产环境中我们往往面临与学术研究不同的约束条件# 典型的多模态Decoder输入结构示例 inputs { visual_feats: torch.randn(batch, 196, 768), # ViT的14x14 patch特征 text_embeddings: torch.randn(batch, seq_len, 768), visual_padding_mask: mask_pad, # 处理不同尺寸图像 text_attention_mask: mask_seq # 处理变长文本 }这种异构输入要求Decoder必须具备动态路由能力而标准实现往往假设输入是同构的序列。3. Decoder的掩码机制深度剖析3.1 因果掩码的时空二象性在纯文本任务中因果掩码Causal Mask是一个简单的下三角矩阵用于阻止当前位置关注后续token。但在处理视觉输入时这种假设面临两个根本性挑战空间局部性约束相邻图像patch之间天然存在强相关性完全阻断未来信息可能损害特征整合模态交互需求文本token可能需要关注图像的任意区域严格时序掩码会限制这种跨模态关联解决方案是采用分阶段掩码策略def create_hybrid_mask(text_len, img_len, device): # 文本部分的严格因果掩码 text_mask torch.tril(torch.ones(text_len, text_len)).to(device) # 图像部分的全局可见掩码 img_mask torch.ones(img_len, img_len).to(device) # 交叉区域处理 cross_mask torch.ones(text_len, img_len).to(device) # 文本可看全图 img_to_text torch.zeros(img_len, text_len).to(device) # 图像不看文本 return torch.block_diag(text_mask, img_mask) torch.cat([cross_mask, img_to_text], dim0)3.2 动态掩码的工程实现技巧实际部署时会遇到几个典型问题内存爆炸预生成所有可能的掩码组合在长序列场景不可行批处理冲突同一batch内样本的序列长度差异导致掩码形状不统一高效实现方案class DynamicMaskGenerator: def __init__(self, max_len512): self.cache {} # 缓存常见组合 def __call__(self, text_lens, img_lens): key (tuple(text_lens), tuple(img_lens)) if key not in self.cache: masks [] for t_len, i_len in zip(text_lens, img_lens): masks.append(create_hybrid_mask(t_len, i_len)) self.cache[key] torch.stack(masks) return self.cache[key]关键提示在TPU等加速器上动态生成掩码可能比预分配内存更高效这与GPU最佳实践相反4. 训练与推理的模式陷阱4.1 自回归解码的隐藏成本训练时使用的Teacher Forcing方式与推理时的自回归生成存在本质差异这导致三个常见误区缓存失效未正确实现KV缓存导致推理速度下降10倍长度发散变长生成时batch内计算资源浪费数值累积长序列生成时的softmax数值稳定性问题改进后的推理实现class EfficientDecoder(TransformerDecoder): def forward(self, x, cacheNone, offset0): if cache is None: # 训练模式 return super().forward(x) # 推理时增量更新 q self.self_attn.q_proj(x) k self.self_attn.k_proj(x) v self.self_attn.v_proj(x) if cache is not None: k torch.cat([cache[k], k], dim1) v torch.cat([cache[v], v], dim1) cache.update({k: k, v: v}) attn (q k.transpose(-2, -1)) * self.scale attn attn.softmax(dim-1) return attn v4.2 多模态对齐的验证策略为确保训练/推理一致性必须建立交叉验证机制确定性测试固定种子比较两种模式的输出差异梯度检查对比teacher forcing与autoregressive的梯度分布边界案例特别测试空输入、单token输入等极端情况验证代码示例def validate_consistency(model, test_input): # 训练模式输出 model.train() with torch.no_grad(): train_out model(test_input) # 模拟推理模式 model.eval() cache None outputs [] for i in range(max_len): out, cache model(test_input[:, i:i1], cache) outputs.append(out) infer_out torch.cat(outputs, dim1) # 比较关键指标 print(f输出差异: {torch.norm(train_out - infer_out)}) print(f注意力分布差异: {compare_attention(train_out, infer_out)})5. 实战中的血泪经验5.1 位置编码的维度诅咒在将ViT与文本Decoder结合时位置编码的维度处理有三大陷阱尺度不匹配图像patch位置编码通常用2D正弦编码而文本用1D归一化冲突图像坐标通常归一化到[0,1]而文本用绝对位置模态混淆错误共享位置编码导致特征污染解决方案示例class MultimodalPositionEmbedding(nn.Module): def __init__(self, d_model, max_text_len512, patch_size16): super().__init__() # 文本位置编码 self.text_pos nn.Parameter(torch.randn(max_text_len, d_model)) # 图像位置编码 self.img_pos nn.Sequential( nn.Conv2d(3, d_model, kernel_sizepatch_size, stridepatch_size), PositionEncodingSine(d_model // 2) # 2D编码 ) def forward(self, text_emb, img_emb, img): text_emb text_emb self.text_pos[:text_emb.size(1)] img_pos self.img_pos(img).flatten(2).transpose(1, 2) img_emb img_emb img_pos return text_emb, img_emb5.2 注意力计算的精度博弈混合精度训练时注意力矩阵计算容易成为数值不稳定性的源头logit溢出FP16下softmax前的值超过65504会导致inf梯度消失极小的注意力权重在反向传播时失效掩码污染错误地将掩码的-INF参与计算稳定实现方案def stable_attention(q, k, v, maskNone): scale q.size(-1)**0.5 scores (q k.transpose(-2, -1)) / scale if mask is not None: scores scores.masked_fill(mask 0, -1e4) # 非极值 # 行归一化前减去最大值 max_values scores.amax(dim-1, keepdimTrue) scores scores - max_values exp_scores torch.exp(scores) if mask is not None: exp_scores exp_scores * mask.float() return (exp_scores / (exp_scores.sum(dim-1, keepdimTrue) 1e-6)) v6. 性能优化实战技巧6.1 内存高效的注意力实现当处理高分辨率图像时注意力矩阵可能消耗GB级内存。通过以下技巧可降低内存占用分块计算将大矩阵拆分为子块逐块处理内存共享重复利用中间结果缓冲区稀疏激活基于门控机制动态修剪低权重连接示例实现class MemoryEfficientAttention(nn.Module): def __init__(self, d_model, num_heads, chunk_size64): super().__init__() self.chunk_size chunk_size self.to_qkv nn.Linear(d_model, d_model * 3) self.scale (d_model // num_heads) ** 0.5 def forward(self, x, maskNone): q, k, v self.to_qkv(x).chunk(3, dim-1) out [] for i in range(0, q.size(1), self.chunk_size): q_chunk q[:, i:iself.chunk_size] attn (q_chunk k.transpose(-2, -1)) / self.scale if mask is not None: attn attn.masked_fill(mask[:, i:iself.chunk_size] 0, -1e9) out.append(attn.softmax(dim-1) v) return torch.cat(out, dim1)6.2 跨设备的分布式策略在多GPU训练时标准的数据并行可能导致注意力计算效率低下通信开销all-gather操作引入的同步延迟内存重复各卡保存完整的KV缓存副本负载不均衡变长序列导致的计算资源浪费优化后的分布式方案class DistributedMultiModalDecoder(nn.Module): def __init__(self, config): super().__init__() self.device_mesh DeviceMesh(cuda, torch.arange(world_size)) self.shard_spec [ None, # batch维度不切分 dp, # 序列维度数据并行 None # 特征维度不切分 ] def forward(self, x): x distribute_tensor(x, self.device_mesh, self.shard_spec) # 各设备只计算本地分片 local_q self.q_proj(x) global_k all_gather(self.k_proj(x)) global_v all_gather(self.v_proj(x)) attn (local_q global_k.transpose(-2, -1)) / self.scale return attn global_v7. 调试与验证方法论7.1 注意力模式可视化诊断建立系统的注意力分析工具能快速定位问题跨头一致性检查不同注意力头间的相关性分析异常值检测识别出格权重分布梯度流向追踪绘制关键参数的梯度传播路径诊断代码框架class AttentionDiagnoser: def __init__(self, model): self.hooks [] self.activations {} for name, layer in model.named_modules(): if isinstance(layer, nn.MultiheadAttention): self.hooks.append(layer.register_forward_hook( lambda m, inp, out, namename: self.activations.update({name: out[1]}) )) def analyze(self, input_sample): with torch.no_grad(): model(input_sample) results {} for name, attn in self.activations.items(): # 计算各头间的相似度 cross_head_sim F.cosine_similarity( attn.unsqueeze(1), attn.unsqueeze(2), dim-1 ) results[f{name}_head_sim] cross_head_sim # 检查异常值 flat_attn attn.flatten() outliers (flat_attn 2 * flat_attn.mean()).sum() results[f{name}_outliers] outliers return results7.2 数值稳定性监控系统在训练过程中实时监控关键指标梯度范数各层梯度幅度的变化曲线激活统计量注意力权重的均值/方差动态损失曲面参数更新方向的曲率变化监控实现示例class TrainingMonitor: def __init__(self, model): self.model model self.metrics { grad_norms: defaultdict(list), activation_stats: defaultdict(dict) } for name, param in model.named_parameters(): param.register_hook( lambda grad, namename: self.metrics[grad_norms][name].append(grad.norm().item()) ) def log_activations(self, outputs): for name, tensor in outputs.items(): if isinstance(tensor, torch.Tensor): self.metrics[activation_stats][name].update({ mean: tensor.mean().item(), std: tensor.std().item(), max: tensor.max().item() }) def plot_metrics(self): # 绘制梯度变化曲线 plt.figure(figsize(12, 6)) for name, norms in self.metrics[grad_norms].items(): plt.plot(norms, labelname) plt.legend() plt.show()8. 前沿扩展方向8.1 稀疏化与动态计算最新研究显示Transformer的计算效率仍有提升空间内容感知稀疏基于输入复杂度动态调整计算量拓扑自适应根据任务需求重组注意力连接模式混合精度路由不同路径采用不同数值精度实验性实现框架class DynamicTransformerBlock(nn.Module): def __init__(self, d_model, num_heads): super().__init__() self.router nn.Linear(d_model, num_heads) self.experts nn.ModuleList([ nn.MultiheadAttention(d_model, 1) for _ in range(num_heads) ]) def forward(self, x): # 动态路由 gate_scores self.router(x).softmax(dim-1) # 专家并行计算 all_outputs [] for i, expert in enumerate(self.experts): out, _ expert(x, x, x) all_outputs.append(out * gate_scores[..., i].unsqueeze(-1)) return sum(all_outputs)8.2 跨架构知识迁移将Transformer的洞见应用于其他架构CNN-Transformer混合在浅层使用卷积提取局部特征RNN记忆增强用LSTM维护长期依赖关系图结构适配将注意力机制扩展到非序列数据混合架构示例class HybridEncoder(nn.Module): def __init__(self): super().__init__() self.cnn nn.Sequential( nn.Conv2d(3, 64, kernel_size7, stride2), nn.ReLU(), nn.MaxPool2d(kernel_size3, stride2) ) self.transformer TransformerEncoder( d_model512, nhead8, num_layers6 ) def forward(self, x): spatial_feats self.cnn(x).flatten(2).permute(0, 2, 1) return self.transformer(spatial_feats) 
长河君 · 2026/10/28 20:01:47 · 阅读 1.2k Transformer多模态解码器工程实践与优化技巧
1. 项目概述解码Transformer的隐秘角落第一次接触Transformer架构时我像大多数人一样被self-attention的优雅设计所震撼。但在实际构建多模态模型的过程中那些论文里一笔带过的实现细节却成了最折磨人的暗坑。比如Decoder的掩码机制在训练和推理时的微妙差异或是位置编码与视觉特征融合时的维度陷阱——这些在教科书式的教程里往往被简化处理却直接影响着模型的实际表现。经过三个跨模态项目的实战洗礼包括一个百万级参数的图文生成系统我逐渐意识到Transformer的强大不仅在于宏观架构的创新更在于那些容易被忽视的工程细节。本文将聚焦Decoder工作机制、掩码的时空控制以及训练/推理的模式差异这三个关键维度结合PyTorch代码实例揭示那些只可意会的实现要点。2. 核心需求解析为什么这些细节至关重要2.1 多模态场景的特殊挑战当Transformer从纯文本领域扩展到视觉-语言多模态任务时Decoder需要同时处理两种截然不同的数据流来自视觉Backbone的网格特征如CNN或ViT输出的patch embeddings和文本的词嵌入序列。这两种特征在时空维度上的不对齐性使得标准的Decoder实现需要做出以下关键调整跨模态注意力权重分配视觉特征通常具有更高的空间冗余度需要在计算query-key乘积时引入适当的温度系数调节序列长度动态变化图像patch数量固定而文本token长度可变导致位置编码需要特殊处理掩码的复合作用既要防止文本解码时的未来信息泄露又要控制视觉特征参与的时空范围2.2 工业级部署的实际需求在真实生产环境中我们往往面临与学术研究不同的约束条件# 典型的多模态Decoder输入结构示例 inputs { visual_feats: torch.randn(batch, 196, 768), # ViT的14x14 patch特征 text_embeddings: torch.randn(batch, seq_len, 768), visual_padding_mask: mask_pad, # 处理不同尺寸图像 text_attention_mask: mask_seq # 处理变长文本 }这种异构输入要求Decoder必须具备动态路由能力而标准实现往往假设输入是同构的序列。3. Decoder的掩码机制深度剖析3.1 因果掩码的时空二象性在纯文本任务中因果掩码Causal Mask是一个简单的下三角矩阵用于阻止当前位置关注后续token。但在处理视觉输入时这种假设面临两个根本性挑战空间局部性约束相邻图像patch之间天然存在强相关性完全阻断未来信息可能损害特征整合模态交互需求文本token可能需要关注图像的任意区域严格时序掩码会限制这种跨模态关联解决方案是采用分阶段掩码策略def create_hybrid_mask(text_len, img_len, device): # 文本部分的严格因果掩码 text_mask torch.tril(torch.ones(text_len, text_len)).to(device) # 图像部分的全局可见掩码 img_mask torch.ones(img_len, img_len).to(device) # 交叉区域处理 cross_mask torch.ones(text_len, img_len).to(device) # 文本可看全图 img_to_text torch.zeros(img_len, text_len).to(device) # 图像不看文本 return torch.block_diag(text_mask, img_mask) torch.cat([cross_mask, img_to_text], dim0)3.2 动态掩码的工程实现技巧实际部署时会遇到几个典型问题内存爆炸预生成所有可能的掩码组合在长序列场景不可行批处理冲突同一batch内样本的序列长度差异导致掩码形状不统一高效实现方案class DynamicMaskGenerator: def __init__(self, max_len512): self.cache {} # 缓存常见组合 def __call__(self, text_lens, img_lens): key (tuple(text_lens), tuple(img_lens)) if key not in self.cache: masks [] for t_len, i_len in zip(text_lens, img_lens): masks.append(create_hybrid_mask(t_len, i_len)) self.cache[key] torch.stack(masks) return self.cache[key]关键提示在TPU等加速器上动态生成掩码可能比预分配内存更高效这与GPU最佳实践相反4. 训练与推理的模式陷阱4.1 自回归解码的隐藏成本训练时使用的Teacher Forcing方式与推理时的自回归生成存在本质差异这导致三个常见误区缓存失效未正确实现KV缓存导致推理速度下降10倍长度发散变长生成时batch内计算资源浪费数值累积长序列生成时的softmax数值稳定性问题改进后的推理实现class EfficientDecoder(TransformerDecoder): def forward(self, x, cacheNone, offset0): if cache is None: # 训练模式 return super().forward(x) # 推理时增量更新 q self.self_attn.q_proj(x) k self.self_attn.k_proj(x) v self.self_attn.v_proj(x) if cache is not None: k torch.cat([cache[k], k], dim1) v torch.cat([cache[v], v], dim1) cache.update({k: k, v: v}) attn (q k.transpose(-2, -1)) * self.scale attn attn.softmax(dim-1) return attn v4.2 多模态对齐的验证策略为确保训练/推理一致性必须建立交叉验证机制确定性测试固定种子比较两种模式的输出差异梯度检查对比teacher forcing与autoregressive的梯度分布边界案例特别测试空输入、单token输入等极端情况验证代码示例def validate_consistency(model, test_input): # 训练模式输出 model.train() with torch.no_grad(): train_out model(test_input) # 模拟推理模式 model.eval() cache None outputs [] for i in range(max_len): out, cache model(test_input[:, i:i1], cache) outputs.append(out) infer_out torch.cat(outputs, dim1) # 比较关键指标 print(f输出差异: {torch.norm(train_out - infer_out)}) print(f注意力分布差异: {compare_attention(train_out, infer_out)})5. 实战中的血泪经验5.1 位置编码的维度诅咒在将ViT与文本Decoder结合时位置编码的维度处理有三大陷阱尺度不匹配图像patch位置编码通常用2D正弦编码而文本用1D归一化冲突图像坐标通常归一化到[0,1]而文本用绝对位置模态混淆错误共享位置编码导致特征污染解决方案示例class MultimodalPositionEmbedding(nn.Module): def __init__(self, d_model, max_text_len512, patch_size16): super().__init__() # 文本位置编码 self.text_pos nn.Parameter(torch.randn(max_text_len, d_model)) # 图像位置编码 self.img_pos nn.Sequential( nn.Conv2d(3, d_model, kernel_sizepatch_size, stridepatch_size), PositionEncodingSine(d_model // 2) # 2D编码 ) def forward(self, text_emb, img_emb, img): text_emb text_emb self.text_pos[:text_emb.size(1)] img_pos self.img_pos(img).flatten(2).transpose(1, 2) img_emb img_emb img_pos return text_emb, img_emb5.2 注意力计算的精度博弈混合精度训练时注意力矩阵计算容易成为数值不稳定性的源头logit溢出FP16下softmax前的值超过65504会导致inf梯度消失极小的注意力权重在反向传播时失效掩码污染错误地将掩码的-INF参与计算稳定实现方案def stable_attention(q, k, v, maskNone): scale q.size(-1)**0.5 scores (q k.transpose(-2, -1)) / scale if mask is not None: scores scores.masked_fill(mask 0, -1e4) # 非极值 # 行归一化前减去最大值 max_values scores.amax(dim-1, keepdimTrue) scores scores - max_values exp_scores torch.exp(scores) if mask is not None: exp_scores exp_scores * mask.float() return (exp_scores / (exp_scores.sum(dim-1, keepdimTrue) 1e-6)) v6. 性能优化实战技巧6.1 内存高效的注意力实现当处理高分辨率图像时注意力矩阵可能消耗GB级内存。通过以下技巧可降低内存占用分块计算将大矩阵拆分为子块逐块处理内存共享重复利用中间结果缓冲区稀疏激活基于门控机制动态修剪低权重连接示例实现class MemoryEfficientAttention(nn.Module): def __init__(self, d_model, num_heads, chunk_size64): super().__init__() self.chunk_size chunk_size self.to_qkv nn.Linear(d_model, d_model * 3) self.scale (d_model // num_heads) ** 0.5 def forward(self, x, maskNone): q, k, v self.to_qkv(x).chunk(3, dim-1) out [] for i in range(0, q.size(1), self.chunk_size): q_chunk q[:, i:iself.chunk_size] attn (q_chunk k.transpose(-2, -1)) / self.scale if mask is not None: attn attn.masked_fill(mask[:, i:iself.chunk_size] 0, -1e9) out.append(attn.softmax(dim-1) v) return torch.cat(out, dim1)6.2 跨设备的分布式策略在多GPU训练时标准的数据并行可能导致注意力计算效率低下通信开销all-gather操作引入的同步延迟内存重复各卡保存完整的KV缓存副本负载不均衡变长序列导致的计算资源浪费优化后的分布式方案class DistributedMultiModalDecoder(nn.Module): def __init__(self, config): super().__init__() self.device_mesh DeviceMesh(cuda, torch.arange(world_size)) self.shard_spec [ None, # batch维度不切分 dp, # 序列维度数据并行 None # 特征维度不切分 ] def forward(self, x): x distribute_tensor(x, self.device_mesh, self.shard_spec) # 各设备只计算本地分片 local_q self.q_proj(x) global_k all_gather(self.k_proj(x)) global_v all_gather(self.v_proj(x)) attn (local_q global_k.transpose(-2, -1)) / self.scale return attn global_v7. 调试与验证方法论7.1 注意力模式可视化诊断建立系统的注意力分析工具能快速定位问题跨头一致性检查不同注意力头间的相关性分析异常值检测识别出格权重分布梯度流向追踪绘制关键参数的梯度传播路径诊断代码框架class AttentionDiagnoser: def __init__(self, model): self.hooks [] self.activations {} for name, layer in model.named_modules(): if isinstance(layer, nn.MultiheadAttention): self.hooks.append(layer.register_forward_hook( lambda m, inp, out, namename: self.activations.update({name: out[1]}) )) def analyze(self, input_sample): with torch.no_grad(): model(input_sample) results {} for name, attn in self.activations.items(): # 计算各头间的相似度 cross_head_sim F.cosine_similarity( attn.unsqueeze(1), attn.unsqueeze(2), dim-1 ) results[f{name}_head_sim] cross_head_sim # 检查异常值 flat_attn attn.flatten() outliers (flat_attn 2 * flat_attn.mean()).sum() results[f{name}_outliers] outliers return results7.2 数值稳定性监控系统在训练过程中实时监控关键指标梯度范数各层梯度幅度的变化曲线激活统计量注意力权重的均值/方差动态损失曲面参数更新方向的曲率变化监控实现示例class TrainingMonitor: def __init__(self, model): self.model model self.metrics { grad_norms: defaultdict(list), activation_stats: defaultdict(dict) } for name, param in model.named_parameters(): param.register_hook( lambda grad, namename: self.metrics[grad_norms][name].append(grad.norm().item()) ) def log_activations(self, outputs): for name, tensor in outputs.items(): if isinstance(tensor, torch.Tensor): self.metrics[activation_stats][name].update({ mean: tensor.mean().item(), std: tensor.std().item(), max: tensor.max().item() }) def plot_metrics(self): # 绘制梯度变化曲线 plt.figure(figsize(12, 6)) for name, norms in self.metrics[grad_norms].items(): plt.plot(norms, labelname) plt.legend() plt.show()8. 前沿扩展方向8.1 稀疏化与动态计算最新研究显示Transformer的计算效率仍有提升空间内容感知稀疏基于输入复杂度动态调整计算量拓扑自适应根据任务需求重组注意力连接模式混合精度路由不同路径采用不同数值精度实验性实现框架class DynamicTransformerBlock(nn.Module): def __init__(self, d_model, num_heads): super().__init__() self.router nn.Linear(d_model, num_heads) self.experts nn.ModuleList([ nn.MultiheadAttention(d_model, 1) for _ in range(num_heads) ]) def forward(self, x): # 动态路由 gate_scores self.router(x).softmax(dim-1) # 专家并行计算 all_outputs [] for i, expert in enumerate(self.experts): out, _ expert(x, x, x) all_outputs.append(out * gate_scores[..., i].unsqueeze(-1)) return sum(all_outputs)8.2 跨架构知识迁移将Transformer的洞见应用于其他架构CNN-Transformer混合在浅层使用卷积提取局部特征RNN记忆增强用LSTM维护长期依赖关系图结构适配将注意力机制扩展到非序列数据混合架构示例class HybridEncoder(nn.Module): def __init__(self): super().__init__() self.cnn nn.Sequential( nn.Conv2d(3, 64, kernel_size7, stride2), nn.ReLU(), nn.MaxPool2d(kernel_size3, stride2) ) self.transformer TransformerEncoder( d_model512, nhead8, num_layers6 ) def forward(self, x): spatial_feats self.cnn(x).flatten(2).permute(0, 2, 1) return self.transformer(spatial_feats)