add: the V6.1 filter model

filter/embedding.py

@@ -31,10 +31,8 @@ def prepare_batch(batch_data, device="cpu"):
     return batch_tensor
 
 import onnxruntime as ort
-from transformers import AutoTokenizer
-from itertools import accumulate
 
-def prepare_batch_per_token(batch_data, max_length=1024):
+def prepare_batch_per_token(session, tokenizer, batch_data, max_length=1024):
     """
     将输入的 batch_data 转换为模型所需的输入格式 [batch_size, num_channels, seq_length, embedding_dim]。
 
@@ -42,69 +40,33 @@ def prepare_batch_per_token(batch_data, max_length=1024):
         batch_data (dict): 输入的 batch 数据，格式为 {
             "title": [text1, text2, ...],
             "description": [text1, text2, ...],
-            "tags": [text1, text2, ...],
+            "tags": [text1, text2, ...]
-            "author_info": [text1, text2, ...]
         }
         max_length (int): 最大序列长度。
 
     返回:
-        torch.Tensor: 形状为 [batch_size, num_channels, seq_length, embedding_dim] 的张量。
+        torch.Tensor: 形状为 [batch_size, num_channels, max_length, embedding_dim] 的张量。
     """
-    # 初始化 tokenizer 和 ONNX 模型
-    tokenizer = AutoTokenizer.from_pretrained("alikia2x/jina-embedding-v3-m2v-1024")
-    session = ort.InferenceSession("./model/embedding_256/onnx/model.onnx")
 
-    # 1. 对每个通道的文本分别编码
+    batch_size = len(batch_data["title"])
-    channel_embeddings = []
+    batch_tensor = torch.zeros(batch_size, 3, max_length, 256)
-    for channel in ["title", "description", "tags", "author_info"]:
+    for i in range(batch_size):
-        texts = batch_data[channel]  # 获取当前通道的文本列表
+        channel_embeddings = torch.zeros((3, 1024, 256))
+        for j, channel in enumerate(["title", "description", "tags"]):
+            # 获取当前通道的文本
+            text = batch_data[channel][i]  
+            encoded_inputs = tokenizer(text, truncation=True, max_length=max_length, return_tensors='np')
 
-        # Step 1: 生成 input_ids 和 offsets
+            # embeddings: [max_length, embedding_dim]
-        # 对每个文本单独编码，保留原始 token 长度
+            embeddings = torch.zeros((1024, 256))
-        encoded_inputs = [tokenizer(text, truncation=True, max_length=max_length, return_tensors='np') for text in texts]
+            for idx, token in enumerate(encoded_inputs['input_ids'][0]):
+                inputs = {
+                    "input_ids": ort.OrtValue.ortvalue_from_numpy(np.array([token])),
+                    "offsets": ort.OrtValue.ortvalue_from_numpy(np.array([0], dtype=np.int64))
+                }
+                output = session.run(None, inputs)[0]
+                embeddings[idx] =  torch.from_numpy(output)
+            channel_embeddings[j] = embeddings
+        batch_tensor[i] = channel_embeddings
 
-        # 提取每个文本的 input_ids 长度（考虑实际的 token 数量）
-        input_ids_lengths = [len(enc["input_ids"][0]) for enc in encoded_inputs]
-
-        # 生成 offsets: [0, len1, len1+len2, ...]
-        offsets = list(accumulate([0] + input_ids_lengths[:-1]))  # 累积和，排除最后一个长度
-
-        # 将所有 input_ids 展平为一维数组
-        flattened_input_ids = np.concatenate([enc["input_ids"][0] for enc in encoded_inputs], axis=0).astype(np.int64)
-
-        # Step 2: 构建 ONNX 输入
-        inputs = {
-            "input_ids": ort.OrtValue.ortvalue_from_numpy(flattened_input_ids),
-            "offsets": ort.OrtValue.ortvalue_from_numpy(np.array(offsets, dtype=np.int64))
-        }
-
-        # Step 3: 运行 ONNX 模型
-        embeddings = session.run(None, inputs)[0]  # 假设输出名为 "embeddings"
-
-        # Step 4: 将输出重塑为 [batch_size, seq_length, embedding_dim]
-        # 注意：这里假设 ONNX 输出的形状是 [total_tokens, embedding_dim]
-        # 需要根据实际序列长度重新分组
-        batch_size = len(texts)
-        embeddings_split = np.split(embeddings, np.cumsum(input_ids_lengths[:-1]))
-        padded_embeddings = []
-        for emb, seq_len in zip(embeddings_split, input_ids_lengths):
-            # 对每个序列填充到 max_length
-            if seq_len > max_length:
-                # 如果序列长度超过 max_length，截断
-                emb = emb[:max_length]
-                pad_length = 0
-            else:
-                # 否则填充到 max_length
-                pad_length = max_length - seq_len
-
-            # 填充到 [max_length, embedding_dim]
-            padded = np.pad(emb, ((0, pad_length), (0, 0)), mode='constant')
-            padded_embeddings.append(padded)
-
-        # 确保所有填充后的序列形状一致
-        embeddings_tensor = torch.tensor(np.stack(padded_embeddings), dtype=torch.float32)
-        channel_embeddings.append(embeddings_tensor)
-
-    # 2. 将编码结果堆叠为 [batch_size, num_channels, seq_length, embedding_dim]
-    batch_tensor = torch.stack(channel_embeddings, dim=1)
     return batch_tensor

filter/modelV6_0.py

@@ -5,8 +5,8 @@ import torch.nn.functional as F
 class VideoClassifierV6_0(nn.Module):
     def __init__(self, embedding_dim=256, seq_length=1024, hidden_dim=512, output_dim=3):
         super().__init__()
-        self.num_channels = 4
+        self.num_channels = 3
-        self.channel_names = ['title', 'description', 'tags', 'author_info']
+        self.channel_names = ['title', 'description', 'tags']
         
         # CNN特征提取层
         self.conv_layers = nn.Sequential(
@@ -66,28 +66,3 @@ class VideoClassifierV6_0(nn.Module):
         
         # 全连接层分类
         return self.fc(flat_features)
-
-# 损失函数保持不变
-class AdaptiveRecallLoss(nn.Module):
-    def __init__(self, class_weights, alpha=0.8, gamma=2.0, fp_penalty=0.5):
-        super().__init__()
-        self.class_weights = class_weights
-        self.alpha = alpha
-        self.gamma = gamma
-        self.fp_penalty = fp_penalty
-
-    def forward(self, logits, targets):
-        ce_loss = F.cross_entropy(logits, targets, weight=self.class_weights, reduction='none')
-        pt = torch.exp(-ce_loss)
-        focal_loss = ((1 - pt) ** self.gamma) * ce_loss
-        
-        class_mask = F.one_hot(targets, num_classes=len(self.class_weights))
-        class_weights = (self.alpha + (1 - self.alpha) * pt.unsqueeze(-1)) * class_mask
-        recall_loss = (class_weights * focal_loss.unsqueeze(-1)).sum(dim=1)
-        
-        probs = F.softmax(logits, dim=1)
-        fp_mask = (targets != 0) & (torch.argmax(logits, dim=1) == 0)
-        fp_loss = self.fp_penalty * probs[:, 0][fp_mask].pow(2).sum()
-        
-        total_loss = recall_loss.mean() + fp_loss / len(targets)
-        return total_loss

filter/modelV6_1.py

@@ -0,0 +1,108 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class VideoClassifierV6_1(nn.Module):
+    def __init__(self, embedding_dim=256, seq_length=1024, hidden_dim=256, output_dim=3, num_heads=4):
+        super().__init__()
+        self.num_channels = 3
+        self.channel_names = ['title', 'description', 'tags']
+        self.embedding_dim = embedding_dim
+        self.hidden_dim = hidden_dim  # 每个通道处理后的特征维度
+        
+        # 通道独立处理模块（每个通道独立的Transformer编码器）
+        self.channel_processors = nn.ModuleList()
+        for _ in range(self.num_channels):
+            self.channel_processors.append(
+                nn.Sequential(
+                    # 自注意力层
+                    nn.MultiheadAttention(
+                        embed_dim=embedding_dim,
+                        num_heads=num_heads,
+                        dropout=0.1
+                    ),
+                    # 层归一化和前馈网络
+                    nn.LayerNorm(embedding_dim),
+                    nn.Linear(embedding_dim, hidden_dim),
+                    nn.GELU(),
+                    nn.Linear(hidden_dim, hidden_dim),
+                    nn.LayerNorm(hidden_dim)
+                )
+            )
+        
+        # 通道权重（可学习，Sigmoid约束）
+        self.channel_weights = nn.Parameter(torch.ones(self.num_channels))
+        
+        # 全连接层（扩展维度）
+        self.fc = nn.Sequential(
+            nn.Linear(self.num_channels * hidden_dim, 1024),  # 拼接后的特征维度
+            nn.BatchNorm1d(1024),
+            nn.Dropout(0.2),
+            nn.GELU(),
+            nn.Linear(1024, 512),
+            nn.BatchNorm1d(512),
+            nn.Dropout(0.2),
+            nn.GELU(),
+            nn.Linear(512, output_dim)
+        )
+        
+        self._init_weights()
+    
+    def _init_weights(self):
+        """权重初始化（Xavier初始化）"""
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.MultiheadAttention):
+                # 初始化MultiheadAttention的参数（输入投影和输出投影）
+                for name, param in m.named_parameters():
+                    if "in_proj" in name or "out_proj" in name:
+                        if "weight" in name:
+                            nn.init.xavier_uniform_(param)
+                        elif "bias" in name:
+                            nn.init.zeros_(param)
+            elif isinstance(m, nn.LayerNorm):
+                nn.init.ones_(m.weight)
+    
+    def forward(self, channel_features: torch.Tensor):
+        """
+        输入格式: [batch_size, num_channels, seq_length, embedding_dim]
+        输出格式: [batch_size, output_dim]
+        """
+        batch_size = channel_features.size(0)
+        processed_channels = []
+        
+        for c in range(self.num_channels):
+            # 提取当前通道的特征 [B, S, E]
+            c_data = channel_features[:, c]
+            # 转置为 [S, B, E] 以适配MultiheadAttention
+            c_data = c_data.permute(1, 0, 2)
+            
+            # 通道独立处理
+            x = c_data
+            for layer in self.channel_processors[c]:
+                if isinstance(layer, nn.MultiheadAttention):
+                    # 自注意力层需要显式提供键、值
+                    x = layer(x, x, x)[0]
+                else:
+                    x = layer(x)
+            # 转回 [B, S, hidden_dim]
+            x = x.permute(1, 0, 2)
+            # 全局池化（序列维度平均）
+            pooled = x.mean(dim=1)  # [B, hidden_dim]
+            processed_channels.append(pooled)
+        
+        # 堆叠通道特征 [B, C, hidden_dim]
+        processed_channels = torch.stack(processed_channels, dim=1)
+        
+        # 应用通道权重（Sigmoid约束）
+        weights = torch.sigmoid(self.channel_weights).unsqueeze(0).unsqueeze(-1)  # [1, C, 1]
+        weighted_features = processed_channels * weights  # [B, C, hidden_dim]
+        
+        # 拼接所有通道特征
+        combined = weighted_features.view(batch_size, -1)  # [B, C*hidden_dim]
+        
+        # 全连接层分类
+        return self.fc(combined)

filter/train.py

@@ -4,14 +4,16 @@ import numpy as np
 from torch.utils.data import DataLoader
 import torch.optim as optim
 from dataset import MultiChannelDataset
-from filter.modelV3_15 import AdaptiveRecallLoss, VideoClassifierV3_15
+from filter.modelV6_1 import VideoClassifierV6_1
 from sklearn.metrics import f1_score, recall_score, precision_score, accuracy_score, classification_report
 import os
 import torch
 from torch.utils.tensorboard import SummaryWriter
 import time
-from embedding import prepare_batch
+from embedding import prepare_batch_per_token
-import torch.nn as nn
+import onnxruntime as ort
+from transformers import AutoTokenizer
+from torch import nn
 
 
 run_name = f"run_{time.strftime('%Y%m%d_%H%M')}"
@@ -24,6 +26,8 @@ writer = SummaryWriter(log_dir=log_dir)
 train_dataset = MultiChannelDataset('./data/filter/labeled_data.jsonl', mode='train')
 eval_dataset = MultiChannelDataset('./data/filter/labeled_data.jsonl', mode='eval')
 
+samples_count = len(train_dataset)
+
 # 加载test数据集
 test_file = './data/filter/test.jsonl'
 if not os.path.exists(test_file):
@@ -50,21 +54,26 @@ class_weights = torch.tensor(
     device='cpu'
 )
 
-# 初始化模型和SentenceTransformer
+model = VideoClassifierV6_1()
-model = VideoClassifierV3_15()
+checkpoint_name = './filter/checkpoints/best_model_V6.2-test2.pt'
-checkpoint_name = './filter/checkpoints/best_model_V3.17.pt'
+
+# 初始化tokenizer和embedding模型
+tokenizer = AutoTokenizer.from_pretrained("alikia2x/jina-embedding-v3-m2v-1024")
+session = ort.InferenceSession("./model/embedding_256/onnx/model.onnx")
 
 # 模型保存路径
 os.makedirs('./filter/checkpoints', exist_ok=True)
 
 # 优化器
-optimizer = optim.AdamW(model.parameters(), lr=4e-4)
+eval_interval = 20
-criterion = AdaptiveRecallLoss(
+num_epochs = 20
-    class_weights=class_weights,
+total_steps = samples_count * num_epochs / train_loader.batch_size
-    alpha=0.9,     # 召回率权重
+warmup_rate = 0.1
-    gamma=1.6,     # 困难样本聚焦
+optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-5)
-    fp_penalty=0.8 # 假阳性惩罚强度
+cosine_annealing_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=total_steps - int(total_steps * warmup_rate))
-)
+warmup_scheduler = optim.lr_scheduler.LinearLR(optimizer, start_factor=0.1, end_factor=1.0, total_iters=int(total_steps * warmup_rate))
+scheduler = optim.lr_scheduler.SequentialLR(optimizer, schedulers=[warmup_scheduler, cosine_annealing_scheduler], milestones=[int(total_steps * warmup_rate)])
+criterion = nn.CrossEntropyLoss(weight=class_weights)
 
 def count_trainable_parameters(model):
     return sum(p.numel() for p in model.parameters() if p.requires_grad)
@@ -76,7 +85,7 @@ def evaluate(model, dataloader):
     
     with torch.no_grad():
         for batch in dataloader:
-            batch_tensor = prepare_batch(batch['texts'])
+            batch_tensor = prepare_batch_per_token(session, tokenizer, batch['texts'])
             logits = model(batch_tensor)
             preds = torch.argmax(logits, dim=1)
             all_preds.extend(preds.cpu().numpy())
@@ -98,8 +107,6 @@ print(f"Trainable parameters: {count_trainable_parameters(model)}")
 # 训练循环
 best_f1 = 0
 step = 0
-eval_interval = 20
-num_epochs = 8
 
 for epoch in range(num_epochs):
     model.train()
@@ -109,7 +116,8 @@ for epoch in range(num_epochs):
     for batch_idx, batch in enumerate(train_loader):
         optimizer.zero_grad()
 
-        batch_tensor = prepare_batch(batch['texts'])
+        
+        batch_tensor = prepare_batch_per_token(session, tokenizer, batch['texts'])
 
         logits = model(batch_tensor)
         
@@ -142,7 +150,8 @@ for epoch in range(num_epochs):
                 best_f1 = eval_f1
                 torch.save(model.state_dict(), checkpoint_name)
                 print("  Saved best model")
-                print("Channel weights: ", model.get_channel_weights())
+        scheduler.step()
+        writer.add_scalar('Train/LR', scheduler.get_last_lr()[0], step)
     
     # 记录每个 epoch 的平均训练损失
     avg_epoch_loss = epoch_loss / len(train_loader)