基于Transformer的高效推理优化技术详解

mindmap
root((Transformer高效推理))
    推理挑战
    计算优化
    内存优化
    模型压缩
    量化技术
    并行计算
    专用硬件
    自适应推理
    多模态融合
    未来趋势

graph TD
    A[Transformer计算复杂度]
    A --> B[自注意力机制]
    A --> C[前馈网络]
    A --> D[层归一化]
    A --> E[位置编码]
    
    B -- 复杂度 --> B1[O(L²·d)]
    B -- 瓶颈 --> B2[QK^T矩阵乘法]
    
    C -- 复杂度 --> C1[O(L·d²)]
    C -- 瓶颈 --> C2[高维特征线性变换]
    
    D -- 复杂度 --> D1[O(L·d)]
    
    E -- 复杂度 --> E1[O(L·d)]

import torch
import torch.nn as nn
import torch.nn.functional as F

class LocalAttention(nn.Module):
    def __init__(self, d_model, num_heads, window_size=64, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        self.window_size = window_size
        
        self.q_linear = nn.Linear(d_model, d_model)
        self.k_linear = nn.Linear(d_model, d_model)
        self.v_linear = nn.Linear(d_model, d_model)
        self.out_linear = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, q, k, v, mask=None):
        batch_size = q.size(0)
        seq_len = q.size(1)
        
        # 线性变换
        q = self.q_linear(q).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        k = self.k_linear(k).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        v = self.v_linear(v).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        
        # 计算局部注意力
        # 为每个位置创建局部窗口
        output = []
        for i in range(0, seq_len, self.window_size):
            end_idx = min(i + self.window_size, seq_len)
            
            # 获取当前窗口的q, k, v
            q_window = q[:, :, i:end_idx, :]
            k_window = k[:, :, i:end_idx, :]
            v_window = v[:, :, i:end_idx, :]
            
            # 计算注意力分数
            scores = torch.matmul(q_window, k_window.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_k, dtype=torch.float32))
            
            # 应用mask（如果有的话）
            if mask is not None and mask.shape[2] == seq_len:
                mask_window = mask[:, i:end_idx, i:end_idx].unsqueeze(1)
                scores = scores.masked_fill(mask_window == 0, -1e9)
            
            # softmax归一化
            attn = F.softmax(scores, dim=-1)
            attn = self.dropout(attn)
            
            # 计算注意力加权和
            context = torch.matmul(attn, v_window)
            output.append(context)
        
        # 拼接所有窗口的结果
        output = torch.cat(output, dim=2)
        
        # 输出线性变换
        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
        output = self.out_linear(output)
        
        return output

# 使用示例
# local_attn = LocalAttention(d_model=768, num_heads=12, window_size=64)
# output = local_attn(q, k, v, mask)

import torch
from collections import deque

class KVCacheManager:
    def __init__(self, max_cache_size=None, cache_quantization=None):
        self.max_cache_size = max_cache_size
        self.cache_quantization = cache_quantization  # 可选值: None, 'int8', 'fp16'
        self.key_caches = {}
        self.value_caches = {}
        self.sequence_lengths = deque()  # 用于LRU缓存管理
        self.sequence_ids = set()
    
    def _quantize_tensor(self, tensor):
        if self.cache_quantization == 'int8':
            # INT8量化
            scale = tensor.abs().max() / 127
            quantized_tensor = (tensor / scale).round().to(torch.int8)
            return quantized_tensor, scale
        elif self.cache_quantization == 'fp16':
            # FP16量化
            return tensor.to(torch.float16), None
        else:
            # 不量化
            return tensor, None
    
    def _dequantize_tensor(self, quantized_tensor, scale=None):
        if self.cache_quantization == 'int8':
            # INT8反量化
            return quantized_tensor.to(torch.float32) * scale
        elif self.cache_quantization == 'fp16':
            # FP16反量化
            return quantized_tensor.to(torch.float32)
        else:
            # 不需要反量化
            return quantized_tensor
    
    def add_sequence(self, sequence_id):
        if sequence_id in self.sequence_ids:
            # 如果序列已存在，先移除
            self.remove_sequence(sequence_id)
        
        # 检查缓存大小是否超过限制
        if self.max_cache_size is not None and len(self.sequence_ids) >= self.max_cache_size:
            # 移除最早的序列（LRU策略）
            oldest_sequence_id = self.sequence_lengths.popleft()
            self.remove_sequence(oldest_sequence_id)
        
        # 为新序列创建空缓存
        self.key_caches[sequence_id] = []
        self.value_caches[sequence_id] = []
        self.sequence_lengths.append(sequence_id)
        self.sequence_ids.add(sequence_id)
    
    def remove_sequence(self, sequence_id):
        if sequence_id in self.sequence_ids:
            del self.key_caches[sequence_id]
            del self.value_caches[sequence_id]
            self.sequence_ids.remove(sequence_id)
            # 从deque中移除（效率不高，但为了简单起见）
            if sequence_id in self.sequence_lengths:
                self.sequence_lengths.remove(sequence_id)
    
    def update_cache(self, sequence_id, key, value):
        if sequence_id not in self.sequence_ids:
            self.add_sequence(sequence_id)
        
        # 量化并存储KV对
        quantized_key, key_scale = self._quantize_tensor(key)
        quantized_value, value_scale = self._quantize_tensor(value)
        
        # 存储量化后的数据和缩放因子
        self.key_caches[sequence_id].append((quantized_key, key_scale))
        self.value_caches[sequence_id].append((quantized_value, value_scale))
    
    def get_cache(self, sequence_id):
        if sequence_id not in self.sequence_ids:
            return None, None
        
        # 获取缓存的KV对并反量化
        keys = []
        values = []
        
        for quantized_key, key_scale in self.key_caches[sequence_id]:
            keys.append(self._dequantize_tensor(quantized_key, key_scale))
        
        for quantized_value, value_scale in self.value_caches[sequence_id]:
            values.append(self._dequantize_tensor(quantized_value, value_scale))
        
        # 拼接所有KV对
        if keys:
            concatenated_keys = torch.cat(keys, dim=1)
            concatenated_values = torch.cat(values, dim=1)
            return concatenated_keys, concatenated_values
        else:
            return None, None
    
    def get_sequence_length(self, sequence_id):
        if sequence_id not in self.sequence_ids:
            return 0
        return len(self.key_caches[sequence_id])
    
    def clear_cache(self):
        self.key_caches.clear()
        self.value_caches.clear()
        self.sequence_lengths.clear()
        self.sequence_ids.clear()

# 使用示例
# kv_cache_manager = KVCacheManager(max_cache_size=100, cache_quantization='fp16')
# 
# # 添加序列
# kv_cache_manager.add_sequence(sequence_id='user_123')
# 
# # 更新缓存
# kv_cache_manager.update_cache(sequence_id='user_123', key=current_key, value=current_value)
# 
# # 获取缓存
# cached_key, cached_value = kv_cache_manager.get_cache(sequence_id='user_123')

import torch
import torch.nn as nn
import torch.nn.utils.prune as prune
import numpy as np

class ChannelPruner:
    def __init__(self, model, pruning_ratio=0.5):
        self.model = model
        self.pruning_ratio = pruning_ratio
        self.pruned_layers = []
    
    def _compute_channel_importance(self, layer, input_data):
        # 对于卷积层或线性层，计算每个通道的重要性
        with torch.no_grad():
            if isinstance(layer, nn.Conv2d) or isinstance(layer, nn.Linear):
                # 使用权重的L1范数作为通道重要性指标
                weight = layer.weight.data
                
                if isinstance(layer, nn.Conv2d):
                    # 对于卷积层，计算每个输出通道的L1范数
                    channel_importance = torch.sum(torch.abs(weight), dim=(1, 2, 3))
                else:
                    # 对于线性层，计算每个输出特征的L1范数
                    channel_importance = torch.sum(torch.abs(weight), dim=1)
                
                return channel_importance
            else:
                return None
    
    def _prune_layer(self, layer, channel_indices_to_keep):
        # 剪枝层的通道
        if isinstance(layer, nn.Conv2d) or isinstance(layer, nn.Linear):
            # 保存原始权重
            original_weight = layer.weight.data.clone()
            
            # 剪枝权重
            if isinstance(layer, nn.Conv2d):
                # 对于卷积层，剪枝输出通道
                layer.weight.data = layer.weight.data[channel_indices_to_keep, :, :, :]
            else:
                # 对于线性层，剪枝输出特征
                layer.weight.data = layer.weight.data[channel_indices_to_keep, :]
            
            # 如果有偏置，也剪枝偏置
            if layer.bias is not None:
                layer.bias.data = layer.bias.data[channel_indices_to_keep]
            
            # 更新层的输出特征数量
            if isinstance(layer, nn.Conv2d):
                layer.out_channels = len(channel_indices_to_keep)
            else:
                layer.out_features = len(channel_indices_to_keep)
            
            return original_weight, channel_indices_to_keep
        else:
            return None, None
    
    def _update_subsequent_layer(self, subsequent_layer, pruned_layer_info):
        # 更新后续层的输入通道以匹配剪枝后的输出通道
        original_weight, channel_indices_to_keep = pruned_layer_info
        
        if original_weight is None or channel_indices_to_keep is None:
            return
        
        if isinstance(subsequent_layer, nn.Conv2d) or isinstance(subsequent_layer, nn.Linear):
            # 如果是卷积层或线性层，需要更新权重的输入通道
            if isinstance(subsequent_layer, nn.Conv2d):
                # 对于卷积层，调整输入通道
                subsequent_layer.weight.data = subsequent_layer.weight.data[:, channel_indices_to_keep, :, :]
                subsequent_layer.in_channels = len(channel_indices_to_keep)
            else:
                # 对于线性层，调整输入特征
                subsequent_layer.weight.data = subsequent_layer.weight.data[:, channel_indices_to_keep]
                subsequent_layer.in_features = len(channel_indices_to_keep)
    
    def prune(self, calibration_data):
        # 遍历模型中的所有层
        layers = list(self.model.named_modules())
        
        for i, (name, layer) in enumerate(layers):
            # 只处理卷积层和线性层
            if isinstance(layer, nn.Conv2d) or isinstance(layer, nn.Linear):
                # 计算通道重要性
                channel_importance = self._compute_channel_importance(layer, calibration_data)
                
                if channel_importance is not None:
                    # 确定要保留的通道数量
                    num_channels = len(channel_importance)
                    num_channels_to_keep = int(num_channels * (1 - self.pruning_ratio))
                    
                    # 选择重要性最高的通道
                    _, indices = torch.topk(channel_importance, num_channels_to_keep)
                    indices = indices.sort()[0]  # 排序以保持顺序
                    
                    # 剪枝当前层
                    pruned_layer_info = self._prune_layer(layer, indices)
                    
                    # 记录剪枝的层
                    self.pruned_layers.append((name, layer, pruned_layer_info))
                    
                    # 更新下一层（如果存在且是卷积层或线性层）
                    if i + 1 < len(layers):
                        next_name, next_layer = layers[i + 1]
                        if isinstance(next_layer, nn.Conv2d) or isinstance(next_layer, nn.Linear):
                            self._update_subsequent_layer(next_layer, pruned_layer_info)
        
        # 返回剪枝后的模型
        return self.model

# 使用示例
# pruner = ChannelPruner(model, pruning_ratio=0.3)
# pruned_model = pruner.prune(calibration_data)
# # 剪枝后通常需要进行微调以恢复性能
# # train_model(pruned_model, fine_tune_data)

import torch
import torch.nn as nn
import numpy as np
from tqdm import tqdm

class PostTrainingQuantizer:
    def __init__(self, model, bits=8, calibration_samples=128):
        self.model = model
        self.bits = bits
        self.calibration_samples = calibration_samples
        self.quant_params = {}
        self.device = next(model.parameters()).device
        
        # 确定量化范围
        self.min_val = -2**(bits-1)
        self.max_val = 2**(bits-1) - 1
    
    def _quantize_tensor(self, tensor, scale, zero_point):
        # 线性量化：tensor_q = round(tensor / scale + zero_point)
        tensor_q = torch.round(tensor / scale + zero_point)
        tensor_q = torch.clamp(tensor_q, self.min_val, self.max_val)
        return tensor_q.to(torch.int8)
    
    def _dequantize_tensor(self, tensor_q, scale, zero_point):
        # 线性反量化：tensor = (tensor_q - zero_point) * scale
        return (tensor_q.float() - zero_point) * scale
    
    def _collect_stats(self, data_loader):
        # 收集激活值统计信息用于校准
        self.model.eval()
        
        # 为每个层创建激活值缓存
        activation_caches = {}
        
        # 注册钩子收集激活值
        hooks = []
        
        def register_hook(name, layer):
            if isinstance(layer, nn.Conv2d) or isinstance(layer, nn.Linear):
                def hook_fn(module, input, output):
                    if name not in activation_caches:
                        activation_caches[name] = []
                    # 只缓存第一个样本批次的数据以节省内存
                    if len(activation_caches[name]) < self.calibration_samples:
                        activation_caches[name].append(output.detach().cpu())
                hooks.append(layer.register_forward_hook(hook_fn))
        
        # 遍历模型并注册钩子
        for name, layer in self.model.named_modules():
            register_hook(name, layer)
        
        # 运行校准数据
        with torch.no_grad():
            for i, (data, _) in enumerate(tqdm(data_loader, desc="Collecting activation stats")):
                if i >= self.calibration_samples:
                    break
                data = data.to(self.device)
                self.model(data)
        
        # 移除钩子
        for hook in hooks:
            hook.remove()
        
        return activation_caches
    
    def _calibrate(self, activation_caches):
        # 校准量化参数
        for name, caches in activation_caches.items():
            # 合并所有校准样本的激活值
            activations = torch.cat(caches, dim=0)
            
            # 计算激活值的最小值和最大值
            min_val = torch.min(activations)
            max_val = torch.max(activations)
            
            # 防止除以零
            if min_val == max_val:
                scale = 1.0
                zero_point = 0
            else:
                # 计算缩放因子和零点
                scale = (max_val - min_val) / (self.max_val - self.min_val)
                # 零点计算确保最小值映射到量化范围的最小值
                zero_point = torch.round(self.min_val - min_val / scale)
                # 确保零点在量化范围内
                zero_point = torch.clamp(zero_point, self.min_val, self.max_val)
            
            self.quant_params[name] = {'scale': scale, 'zero_point': zero_point}
    
    def _quantize_weights(self):
        # 量化模型权重
        for name, param in self.model.named_parameters():
            if 'weight' in name:
                # 获取权重数据
                weight = param.data
                
                # 计算权重的最小值和最大值
                min_val = torch.min(weight)
                max_val = torch.max(weight)
                
                # 防止除以零
                if min_val == max_val:
                    scale = 1.0
                    zero_point = 0
                else:
                    # 计算缩放因子和零点
                    scale = (max_val - min_val) / (self.max_val - self.min_val)
                    # 零点计算确保最小值映射到量化范围的最小值
                    zero_point = torch.round(self.min_val - min_val / scale)
                    # 确保零点在量化范围内
                    zero_point = torch.clamp(zero_point, self.min_val, self.max_val)
                
                # 量化权重
                weight_q = self._quantize_tensor(weight, scale, zero_point)
                
                # 存储量化参数
                self.quant_params[name] = {'scale': scale, 'zero_point': zero_point}
                
                # 使用反量化后的权重替换原始权重（仅用于演示，实际应用中需要专门的量化推理引擎）
                param.data = self._dequantize_tensor(weight_q, scale, zero_point)
    
    def quantize(self, data_loader):
        # 1. 收集激活值统计信息
        activation_caches = self._collect_stats(data_loader)
        
        # 2. 校准量化参数
        self._calibrate(activation_caches)
        
        # 3. 量化模型权重
        self._quantize_weights()
        
        # 返回量化后的模型和量化参数
        return self.model, self.quant_params

# 使用示例
# quantizer = PostTrainingQuantizer(model, bits=8, calibration_samples=128)
# quantized_model, quant_params = quantizer.quantize(calibration_loader)
# # 实际应用中，需要使用支持INT8推理的引擎加载量化模型
# # 例如TensorRT、ONNX Runtime、TFLite等

graph TD
    A[量化技术与实践] --> B[量化基础与分类]
    A --> C[低精度训练与推理]
    A --> D[整型量化技术]
    A --> E[量化工具与框架]
    
    B --> B1[按量化粒度分类]
    B --> B2[按量化位宽分类]
    B --> B3[按量化时机分类]
    B --> B4[按量化方法分类]
    
    C --> C1[半精度训练（FP16）]
    C --> C2[Brain Float 16（BF16）]
    C --> C3[TF32训练]
    C --> C4[低精度推理优化]
    
    D --> D1[校准技术]
    D --> D2[QAT技术]
    D --> D3[INT8推理优化]
    D --> D4[低位宽量化]
    
    E --> E1[主流量化框架]
    E --> E2[专用量化工具]
    E --> E3[量化评估与分析工具]
    E --> E4[量化模型部署工具]

import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
import os

class SimpleTransformerModel(nn.Module):
    def __init__(self, d_model=512, nhead=8, num_layers=6):
        super().__init__()
        self.embedding = nn.Embedding(10000, d_model)
        self.transformer = nn.Transformer(d_model=d_model, nhead=nhead, num_encoder_layers=num_layers, 
                                         num_decoder_layers=num_layers)
        self.fc = nn.Linear(d_model, 10000)
    
    def forward(self, src, tgt):
        src_emb = self.embedding(src) * torch.sqrt(torch.tensor(self.embedding.embedding_dim))
        tgt_emb = self.embedding(tgt) * torch.sqrt(torch.tensor(self.embedding.embedding_dim))
        
        # 添加位置编码（简化版）
        seq_len_src, batch_size, d_model = src_emb.size()
        pos_enc_src = torch.arange(0, seq_len_src, device=src.device).unsqueeze(1).repeat(1, batch_size).unsqueeze(2)
        src_emb = src_emb + pos_enc_src * 0.01
        
        seq_len_tgt, _, _ = tgt_emb.size()
        pos_enc_tgt = torch.arange(0, seq_len_tgt, device=tgt.device).unsqueeze(1).repeat(1, batch_size).unsqueeze(2)
        tgt_emb = tgt_emb + pos_enc_tgt * 0.01
        
        # Transformer的forward方法期望的输入形状是 (seq_len, batch_size, features)
        # 但默认的Transformer模块可能期望不同的维度顺序，这里根据实际情况调整
        output = self.transformer(src_emb, tgt_emb)
        output = self.fc(output)
        return output

# 训练函数
def train(rank, world_size, dataset, args):
    # 设置分布式环境
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    dist.init_process_group(backend='nccl', rank=rank, world_size=world_size)
    
    # 创建模型并移动到指定设备
    torch.cuda.set_device(rank)
    model = SimpleTransformerModel().to(rank)
    
    # 使用DDP包装模型
    ddp_model = DDP(model, device_ids=[rank])
    
    # 创建优化器和损失函数
    optimizer = optim.Adam(ddp_model.parameters(), lr=args.lr)
    criterion = nn.CrossEntropyLoss()
    
    # 创建分布式采样器和数据加载器
    sampler = DistributedSampler(dataset, shuffle=True)
    dataloader = DataLoader(dataset, batch_size=args.batch_size, sampler=sampler)
    
    # 训练循环
    for epoch in range(args.epochs):
        sampler.set_epoch(epoch)  # 确保每个epoch的shuffle不同
        total_loss = 0
        
        for i, (src, tgt) in enumerate(dataloader):
            # 将数据移动到指定设备
            src = src.to(rank)
            tgt = tgt.to(rank)
            
            # 前向传播
            optimizer.zero_grad()
            output = ddp_model(src, tgt[:, :-1])  # tgt[:, :-1]作为输入
            
            # 计算损失
            # output形状: (seq_len, batch_size, vocab_size)
            # tgt[:, 1:]形状: (batch_size, seq_len)
            # 需要调整形状以匹配损失函数
            loss = criterion(output.view(-1, output.size(-1)), tgt[:, 1:].reshape(-1))
            
            # 反向传播
            loss.backward()
            
            # 梯度裁剪（可选）
            torch.nn.utils.clip_grad_norm_(ddp_model.parameters(), max_norm=1.0)
            
            # 更新参数
            optimizer.step()
            
            total_loss += loss.item()
            
            # 只在rank 0上打印进度
            if rank == 0 and i % 10 == 0:
                print(f'Epoch {epoch+1}, Step {i}, Loss: {loss.item():.4f}')
        
        # 只在rank 0上打印 epoch 损失
        if rank == 0:
            avg_loss = total_loss / len(dataloader)
            print(f'Epoch {epoch+1} completed, Average Loss: {avg_loss:.4f}')
    
    # 清理分布式环境
    dist.destroy_process_group()

# 主函数
def main():
    # 定义参数
    class Args:
        def __init__(self):
            self.lr = 1e-4
            self.batch_size = 32
            self.epochs = 10
    
    args = Args()
    
    # 创建虚拟数据集（实际应用中应替换为真实数据集）
    class DummyDataset(torch.utils.data.Dataset):
        def __init__(self, size=1000, seq_len=10):
            self.size = size
            self.seq_len = seq_len
        
        def __len__(self):
            return self.size
        
        def __getitem__(self, idx):
            # 生成随机序列作为源和目标（在实际应用中应使用真实数据）
            src = torch.randint(0, 10000, (self.seq_len,))
            tgt = torch.cat([src[1:], torch.tensor([0])], dim=0)  # 简单的移位目标
            return src, tgt
    
    dataset = DummyDataset()
    
    # 设置分布式训练
    world_size = torch.cuda.device_count()  # 使用所有可用的GPU
    print(f'Using {world_size} GPUs for training')
    
    # 使用mp.spawn启动多个进程
    mp.spawn(train, 
             args=(world_size, dataset, args),
             nprocs=world_size,
             join=True)

# 使用示例
if __name__ == '__main__':
    main()

import torch
import torch.nn as nn
import torch.nn.functional as F

class EarlyExitTransformer(nn.Module):
    def __init__(self, d_model=512, nhead=8, num_layers=6, num_classes=1000):
        super().__init__()
        self.embedding = nn.Embedding(10000, d_model)
        self.transformer_layers = nn.ModuleList([
            nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead)
            for _ in range(num_layers)
        ])
        # 为每个Transformer层添加退出头
        self.exit_heads = nn.ModuleList([
            nn.Linear(d_model, num_classes)
            for _ in range(num_layers)
        ])
        # 最终的输出层
        self.final_head = nn.Linear(d_model, num_classes)
        # 用于确定是否退出的可学习参数
        self.exit_thresholds = nn.ParameterList([
            nn.Parameter(torch.tensor(0.9))  # 初始阈值设为0.9
            for _ in range(num_layers)
        ])
    
    def forward(self, src, return_all=False):
        # src形状: (batch_size, seq_len)
        batch_size = src.size(0)
        
        # 嵌入层
        src_emb = self.embedding(src) * torch.sqrt(torch.tensor(self.embedding.embedding_dim))
        
        # 添加位置编码（简化版）
        seq_len, _, d_model = src_emb.size()
        pos_enc = torch.arange(0, seq_len, device=src.device).unsqueeze(1).repeat(1, batch_size).unsqueeze(2)
        src_emb = src_emb + pos_enc * 0.01
        
        # 交换维度以适应Transformer层的输入要求 (seq_len, batch_size, d_model)
        src_emb = src_emb.permute(1, 0, 2)
        
        # 存储所有退出点的输出
        all_outputs = []
        
        # 遍历所有Transformer层
        x = src_emb
        exit_points = torch.zeros(batch_size, dtype=torch.int64, device=src.device)
        final_outputs = torch.zeros(batch_size, self.exit_heads[0].out_features, device=src.device)
        exited = torch.zeros(batch_size, dtype=torch.bool, device=src.device)
        
        for i, (layer, exit_head, threshold) in enumerate(zip(self.transformer_layers, self.exit_heads, self.exit_thresholds)):
            # 通过当前Transformer层
            x = layer(x)
            
            # 计算当前层的特征表示（取序列的第一个位置作为全局表示）
            features = x[0, :, :]  # (batch_size, d_model)
            
            # 通过退出头计算logits
            logits = exit_head(features)
            
            # 计算置信度
            confidences, _ = torch.max(F.softmax(logits, dim=-1), dim=-1)
            
            # 记录所有退出点的输出
            all_outputs.append(logits)
            
            # 对于尚未退出的样本，检查是否达到退出阈值
            if not exited.all():
                # 找出尚未退出且置信度超过阈值的样本
                new_exits = (~exited) & (confidences > threshold)
                
                # 记录这些样本的退出点和输出
                exit_points[new_exits] = i
                final_outputs[new_exits] = logits[new_exits]
                
                # 更新已退出的样本标记
                exited = exited | new_exits
                
                # 如果所有样本都已退出，可以提前结束
                if exited.all() and not return_all:
                    break
        
        # 对于最后仍未退出的样本，使用最终的输出层
        if not exited.all():
            final_outputs[~exited] = self.final_head(features[~exited])
            exit_points[~exited] = num_layers
        
        if return_all:
            return final_outputs, exit_points, all_outputs
        else:
            return final_outputs, exit_points

# 使用示例
# model = EarlyExitTransformer(d_model=512, nhead=8, num_layers=6, num_classes=1000)
# outputs, exit_points = model(input_ids)
# # exit_points 包含每个样本退出的层索引

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultimodalFusionTransformer(nn.Module):
    def __init__(self, text_dim, image_dim, audio_dim, hidden_dim=768, num_heads=12, num_layers=6, num_classes=1000):
        super().__init__()
        
        # 模态特定编码器
        self.text_encoder = nn.Linear(text_dim, hidden_dim)
        self.image_encoder = nn.Linear(image_dim, hidden_dim)
        self.audio_encoder = nn.Linear(audio_dim, hidden_dim)
        
        # 模态特定位置编码
        self.text_pos_embedding = nn.Parameter(torch.randn(1, 100, hidden_dim))  # 假设文本最大长度为100
        self.image_pos_embedding = nn.Parameter(torch.randn(1, 196, hidden_dim))  # 假设图像特征大小为14x14=196
        self.audio_pos_embedding = nn.Parameter(torch.randn(1, 200, hidden_dim))  # 假设音频特征最大长度为200
        
        # 模态类型嵌入（用于区分不同模态的特征）
        self.text_type_embedding = nn.Parameter(torch.randn(1, 1, hidden_dim))
        self.image_type_embedding = nn.Parameter(torch.randn(1, 1, hidden_dim))
        self.audio_type_embedding = nn.Parameter(torch.randn(1, 1, hidden_dim))
        
        # 融合Transformer编码器
        encoder_layer = nn.TransformerEncoderLayer(d_model=hidden_dim, nhead=num_heads)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        
        # 分类头
        self.classifier = nn.Linear(hidden_dim, num_classes)
        
        # Dropout层
        self.dropout = nn.Dropout(0.1)
    
    def forward(self, text_features, image_features, audio_features):
        # 文本特征编码
        text_emb = self.text_encoder(text_features)  # (batch_size, text_len, text_dim) -> (batch_size, text_len, hidden_dim)
        text_len = text_emb.size(1)
        text_emb = text_emb + self.text_pos_embedding[:, :text_len, :]  # 添加位置编码
        text_emb = text_emb + self.text_type_embedding  # 添加模态类型编码
        text_emb = self.dropout(text_emb)
        
        # 图像特征编码
        image_emb = self.image_encoder(image_features)  # (batch_size, image_len, image_dim) -> (batch_size, image_len, hidden_dim)
        image_len = image_emb.size(1)
        image_emb = image_emb + self.image_pos_embedding[:, :image_len, :]  # 添加位置编码
        image_emb = image_emb + self.image_type_embedding  # 添加模态类型编码
        image_emb = self.dropout(image_emb)
        
        # 音频特征编码
        audio_emb = self.audio_encoder(audio_features)  # (batch_size, audio_len, audio_dim) -> (batch_size, audio_len, hidden_dim)
        audio_len = audio_emb.size(1)
        audio_emb = audio_emb + self.audio_pos_embedding[:, :audio_len, :]  # 添加位置编码
        audio_emb = audio_emb + self.audio_type_embedding  # 添加模态类型编码
        audio_emb = self.dropout(audio_emb)
        
        # 融合所有模态的特征
        # 首先需要将batch_size和seq_len维度交换，以适应Transformer的输入要求
        text_emb = text_emb.permute(1, 0, 2)  # (text_len, batch_size, hidden_dim)
        image_emb = image_emb.permute(1, 0, 2)  # (image_len, batch_size, hidden_dim)
        audio_emb = audio_emb.permute(1, 0, 2)  # (audio_len, batch_size, hidden_dim)
        
        # 拼接所有模态的特征
        multimodal_emb = torch.cat([text_emb, image_emb, audio_emb], dim=0)  # (total_len, batch_size, hidden_dim)
        
        # 通过Transformer编码器
        output = self.transformer_encoder(multimodal_emb)
        
        # 取所有特征的平均值作为最终表示
        # 也可以使用cls标记或其他池化方式
        pooled_output = torch.mean(output, dim=0)  # (batch_size, hidden_dim)
        
        # 通过分类头得到最终预测
        logits = self.classifier(pooled_output)  # (batch_size, num_classes)
        
        return logits

# 使用示例
# model = MultimodalFusionTransformer(text_dim=768, image_dim=2048, audio_dim=128, hidden_dim=768, num_heads=12, num_layers=6, num_classes=1000)
# 
# # 假设我们有以下输入特征
# text_features = torch.randn(32, 50, 768)  # (batch_size, text_len, text_dim)
# image_features = torch.randn(32, 196, 2048)  # (batch_size, image_len, image_dim)
# audio_features = torch.randn(32, 100, 128)  # (batch_size, audio_len, audio_dim)
# 
# # 前向传播
# logits = model(text_features, image_features, audio_features)