import torch import torch.nn as nn from transformers import Wav2Vec2BertModel from llama_nar import LlamaNAREmb from transformers import LlamaConfig import time import torch.nn.functional as F from huggingface_hub import hf_hub_download class Wav2Vec2BERT_Llama(nn.Module): def __init__(self): super().__init__() # 1. 加载预训练模型 self.wav2vec2bert = Wav2Vec2BertModel.from_pretrained("facebook/w2v-bert-2.0", output_hidden_states=True) # 2. 选择性冻结参数 for name, param in self.wav2vec2bert.named_parameters(): # 冻结所有FFN1 (保留FFN2的适应能力) if 'ffn1' in name: param.requires_grad = False # 冻结多头注意力中的K,V投影 if any(proj in name for proj in ['linear_k', 'linear_v']): param.requires_grad = False # 冻结distance_embedding if 'distance_embedding' in name: param.requires_grad = False # 冻结所有卷积相关模块 if any(conv_name in name for conv_name in [ 'conv_module', 'pointwise_conv', 'depthwise_conv', 'feature_extractor', 'pos_conv_embed', 'conv_layers' ]): param.requires_grad = False # 3. 减小Llama模型规模 self.llama_nar = LlamaNAREmb( config=LlamaConfig( hidden_size=512, num_attention_heads=8, num_hidden_layers=8, ), num_heads=8, num_layers=8, hidden_size=512 ) # 4. 降维投影层 self.projection = nn.Sequential( nn.Linear(1024, 512), nn.LayerNorm(512) ) # 5. 简化分类头 self.classifier = nn.Sequential( nn.Linear(512, 128), nn.ReLU(), nn.Dropout(0.1), nn.Linear(128, 2) ) # 6. 减小embedding维度 self.label_embedding = nn.Embedding(num_embeddings=2, embedding_dim=512) # 7. 简化特征处理层 self.feature_processor = nn.Sequential( nn.Linear(512, 512), nn.LayerNorm(512), nn.ReLU(), nn.Dropout(0.1) ) # 8. 减小特殊token的维度 self.special_tokens = nn.Parameter(torch.randn(4, 512)) def _fuse_layers(self, hidden_states): # 修改特征融合方法 def downsample_sequence(sequence, factor=10): """对序列进行下采样""" batch_size, seq_len, hidden_size = sequence.shape # 确保序列长度可以被因子整除 new_len = seq_len // factor padded_len = new_len * factor if seq_len > padded_len: sequence = sequence[:, :padded_len, :] # 重塑张量并进行平均池化 [batch_size, new_len, factor, hidden_size] reshaped = sequence.reshape(batch_size, new_len, factor, hidden_size) downsampled = torch.mean(reshaped, dim=2) # [batch_size, new_len, hidden_size] return downsampled # 1. 获取最后一层特征并进行下采样 last_layer = hidden_states[-1] # [batch_size, seq_len, 1024] downsampled_features = downsample_sequence(last_layer) # [batch_size, seq_len//10, 1024] # 2. 投影到512维度 projected_features = self.projection(downsampled_features) # [batch_size, seq_len//10, 512] return projected_features # 不再需要unsqueeze，因为已经保留了序列维度 def forward(self, batch): main_output = self.wav2vec2bert( **batch['main_features'] ) fused_features = self._fuse_layers(main_output.hidden_states) fused_features = self.feature_processor(fused_features) if ('prompt_labels' in batch and batch['prompt_labels'] is not None and 'prompt_features' in batch and batch['prompt_features'] and len(batch['prompt_features']) > 0): batch_size, num_prompts = batch['prompt_labels'].shape # 重塑特征以批量处理 prompt_features = batch['prompt_features'] all_prompt_outputs = [] for i in range(num_prompts): prompt_output = self.wav2vec2bert( **prompt_features[i] ) all_prompt_outputs.append(self._fuse_layers(prompt_output.hidden_states)) if all_prompt_outputs: fused_prompts = torch.stack([ self.feature_processor(p) for p in all_prompt_outputs ], dim=1) # [batch_size, num_prompts, seq_len, hidden_size] # 获取label embeddings并扩展到对应序列长度 label_embs = self.label_embedding(batch['prompt_labels']) # [batch_size, num_prompts, 512] prompt_embeddings = [] for i in range(batch_size): sequence = [] # 添加示例prompts for j in range(num_prompts): prompt_seq_len = fused_prompts[i, j].size(0) # 获取当前prompt的序列长度 sequence.append(self.special_tokens[1].expand(1, -1)) # [PROMPT] sequence.append(self.special_tokens[2].expand(1, -1)) # [AUDIO] sequence.append(fused_prompts[i, j]) # [seq_len, hidden_size] sequence.append(self.special_tokens[3].expand(1, -1)) # [LABEL] # 扩展label embedding到与音频特征相同的长度 expanded_label = label_embs[i, j].unsqueeze(0).expand(prompt_seq_len, -1) sequence.append(expanded_label) # [seq_len, hidden_size] sequence.append(self.special_tokens[0].expand(1, -1)) # [SEP] # 添加待预测的主特征 main_seq_len = fused_features[i].size(0) # 获取主特征的序列长度 sequence.append(self.special_tokens[1].expand(1, -1)) # [PROMPT] sequence.append(self.special_tokens[2].expand(1, -1)) # [AUDIO] sequence.append(fused_features[i]) # [main_seq_len, hidden_size] sequence.append(self.special_tokens[3].expand(1, -1)) # [LABEL] # 预测位置使用零向量，长度与主特征相同 sequence.append(torch.zeros(main_seq_len, fused_features.size(-1)).to(fused_features.device)) prompt_embeddings.append(torch.cat(sequence, dim=0)) prompt_embeddings = torch.stack(prompt_embeddings, dim=0) else: # 简化无prompt情况的处理 batch_size = fused_features.size(0) main_seq_len = fused_features.size(1) # 直接获取主特征序列长度 # 构建序列 [batch_size, total_len, hidden_size] prompt_embeddings = torch.cat([ self.special_tokens[1].expand(batch_size, 1, -1), # [PROMPT] self.special_tokens[2].expand(batch_size, 1, -1), # [AUDIO] fused_features, # [batch_size, main_seq_len, hidden_size] self.special_tokens[3].expand(batch_size, 1, -1), # [LABEL] torch.zeros(batch_size, main_seq_len, fused_features.size(-1)).to(fused_features.device) # 预测位置 ], dim=1) # 输入到llama_nar output = self.llama_nar(inputs_embeds=prompt_embeddings) # 获取所有预测位置的输出（即最后main_seq_len个位置） pred_pos_embeddings = output[:, -main_seq_len:, :] # [batch_size, main_seq_len, hidden_size] # 对每一帧进行分类 frame_logits = self.classifier(pred_pos_embeddings) # [batch_size, main_seq_len, 2] # 同时返回帧级别的logits和整体的logits（通过平均得到） avg_embedding = torch.mean(pred_pos_embeddings, dim=1) # [batch_size, hidden_size] avg_logits = self.classifier(avg_embedding) # [batch_size, 2] return { 'frame_logits': frame_logits, # 每一帧的预测分数 'avg_logits': avg_logits # 整体的预测分数 } if __name__ == '__main__': import torch from torch.utils.data import DataLoader from dataset.train_MultiDataset import train_MultiDataset, collate_fn from tqdm import tqdm import time # 设置设备 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"\n=== 使用设备: {device} ===") # 初始化模型 print("\n=== 初始化模型 ===") model = Wav2Vec2BERT_Llama().to(device) model.eval() # 设置为评估模式 # 打印wav2vec2bert的参数结构 print("\n=== Wav2Vec2BERT 参数结构 ===") w2v_params_by_layer = {} total_trainable = 0 total_frozen = 0 for name, param in model.wav2vec2bert.named_parameters(): # 获取主要层名称 layer_name = name.split('.')[0] if layer_name not in w2v_params_by_layer: w2v_params_by_layer[layer_name] = { 'trainable_params': 0, 'frozen_params': 0, 'parameter_names': [] } # 统计参数 if param.requires_grad: w2v_params_by_layer[layer_name]['trainable_params'] += param.numel() total_trainable += param.numel() else: w2v_params_by_layer[layer_name]['frozen_params'] += param.numel() total_frozen += param.numel() w2v_params_by_layer[layer_name]['parameter_names'].append(name) # 打印每层的详细信息 print("\n各层参数统计:") for layer_name, info in w2v_params_by_layer.items(): trainable_mb = info['trainable_params'] / 1024 / 1024 frozen_mb = info['frozen_params'] / 1024 / 1024 total_mb = (info['trainable_params'] + info['frozen_params']) / 1024 / 1024 print(f"\n{layer_name}:") print(f" - 总参数量: {total_mb:.2f}MB") print(f" - 可训练参数: {trainable_mb:.2f}MB") print(f" - 冻结参数: {frozen_mb:.2f}MB") print(f" - 参数名称:") for param_name in info['parameter_names']: print(f" * {param_name}") # 打印总体统计 print("\n=== 总体统计 ===") print(f"可训练参数总量: {total_trainable/1024/1024:.2f}MB") print(f"冻结参数总量: {total_frozen/1024/1024:.2f}MB") print(f"参数总量: {(total_trainable + total_frozen)/1024/1024:.2f}MB") print(f"可训练参数占比: {total_trainable/(total_trainable + total_frozen)*100:.2f}%") # 分别统计各个模块的参数量 wav2vec2bert_params = sum(p.numel() for p in model.wav2vec2bert.parameters()) llama_params = sum(p.numel() for p in model.llama_nar.parameters()) other_params = sum(p.numel() for name, p in model.named_parameters() if not name.startswith('wav2vec2bert.') and not name.startswith('llama_nar.')) total_params = wav2vec2bert_params + llama_params + other_params print(f"\n=== 参数量统计 ===") print(f"Wav2Vec2BERT参数量: {wav2vec2bert_params:,} ({wav2vec2bert_params/1024/1024:.2f}MB)") print(f"LlamaNAR参数量: {llama_params:,} ({llama_params/1024/1024:.2f}MB)") print(f"其他模块参数量: {other_params:,} ({other_params/1024/1024:.2f}MB)") print(f"总参数量: {total_params:,} ({total_params/1024/1024:.2f}MB)") # 计算百分比 print(f"\n=== 参数量占比 ===") print(f"Wav2Vec2BERT: {wav2vec2bert_params/total_params*100:.2f}%") print(f"LlamaNAR: {llama_params/total_params*100:.2f}%") print(f"其他模块: {other_params/total_params*100:.2f}%") # 测试运行时间和内存使用 print("\n=== 测试运行时间和内存使用 (batch_size=4) ===") batch_size = 4 total_samples = 600000 # 清空GPU缓存 if torch.cuda.is_available(): torch.cuda.empty_cache() initial_memory = torch.cuda.memory_allocated() / 1024 / 1024 print(f"初始GPU内存使用: {initial_memory:.2f}MB") # 初始化数据集 print("\n初始化数据集...") ds = train_MultiDataset(max_prompts=3) # 创建DataLoader dl = DataLoader(ds, batch_size=batch_size, shuffle=True, collate_fn=collate_fn, num_workers=4) print(f"\n数据集大小: {len(ds)}") print(f"批次数量: {len(dl)}") # 计算一个batch的平均时间 num_test_batches = 10 total_time = 0 max_memory = 0 print(f"\n测试{num_test_batches}个batch的平均运行时间...") with torch.no_grad(): for i, batch in enumerate(tqdm(dl, total=num_test_batches)): if i >= num_test_batches: break # 正确处理字典类型的特征 main_features = { 'input_features': batch['main_features']['input_features'].to(device), 'attention_mask': batch['main_features']['attention_mask'].to(device) } prompt_features = [{ 'input_features': pf['input_features'].to(device), 'attention_mask': pf['attention_mask'].to(device) } for pf in batch['prompt_features']] labels = batch['labels'].to(device) prompt_labels = batch['prompt_labels'].to(device) # 记录开始时间 start_time = time.time() # 前向传播 outputs = model({ 'main_features': main_features, 'prompt_features': prompt_features, 'prompt_labels': prompt_labels }) # 确保GPU运算完成 if torch.cuda.is_available(): torch.cuda.synchronize() # 记录结束时间和内存使用 end_time = time.time() total_time += (end_time - start_time) if torch.cuda.is_available(): current_memory = torch.cuda.memory_allocated() / 1024 / 1024 max_memory = max(max_memory, current_memory) # 打印第一个batch的详细信息 if i == 0: print("\n=== 第一个Batch的详细信息 ===") print(f"主特征形状: {main_features['input_features'].shape}") print(f"主掩码形状: {main_features['attention_mask'].shape}") print(f"Prompt特征形状: {prompt_features[0]['input_features'].shape}") print(f"Prompt掩码形状: {prompt_features[0]['attention_mask'].shape}") print(f"标签形状: {labels.shape}") print(f"Prompt标签形状: {prompt_labels.shape}") print(f"模型输出形状: {outputs.shape}") print(f"输出logits范围: [{outputs.min().item():.3f}, {outputs.max().item():.3f}]") # 计算和打印统计信息 avg_time = total_time / num_test_batches print(f"\n=== 性能统计 ===") print(f"平均每个batch处理时间: {avg_time:.4f}秒") print(f"估计处理{total_samples}个样本需要: {(total_samples/batch_size*avg_time/3600):.2f}小时") if torch.cuda.is_available(): print(f"最大GPU内存使用: {max_memory:.2f}MB") print(f"GPU内存增长: {max_memory - initial_memory:.2f}MB") print("\n测试完成!")