add: filter model based on CNN(V6.0)
This commit is contained in:
parent
1720ff332e
commit
3aa3bc8596
@ -18,4 +18,8 @@ Note
|
||||
0324: V3.5-test3 # 用回3.2的FC层试试
|
||||
0331: V3.6-test3 # 3.5不太行,我试着调下超参
|
||||
0335: V3.7-test3 # 3.6还行,再调超参试试看
|
||||
0414: V3.8-test3 # 3.7不行,从3.6的基础重新调
|
||||
0414: V3.8-test3 # 3.7不行,从3.6的基础重新调
|
||||
1918: V3.9
|
||||
2308: V3.11
|
||||
2243: V3.11 # 256维嵌入
|
||||
2253: V3.11 # 1024维度嵌入(对比)
|
@ -20,7 +20,7 @@ def prepare_batch(batch_data, device="cpu"):
|
||||
"""
|
||||
# 1. 对每个通道的文本分别编码
|
||||
channel_embeddings = []
|
||||
model = StaticModel.from_pretrained("./model/embedding/")
|
||||
model = StaticModel.from_pretrained("./model/embedding_1024/")
|
||||
for channel in ["title", "description", "tags", "author_info"]:
|
||||
texts = batch_data[channel] # 获取当前通道的文本列表
|
||||
embeddings = torch.from_numpy(model.encode(texts)).to(torch.float32).to(device) # 编码为 [batch_size, embedding_dim]
|
||||
@ -28,4 +28,37 @@ def prepare_batch(batch_data, device="cpu"):
|
||||
|
||||
# 2. 将编码结果堆叠为 [batch_size, num_channels, embedding_dim]
|
||||
batch_tensor = torch.stack(channel_embeddings, dim=1) # 在 dim=1 上堆叠
|
||||
return batch_tensor
|
||||
|
||||
def prepare_batch_per_token(batch_data, max_length=1024):
|
||||
"""
|
||||
将输入的 batch_data 转换为模型所需的输入格式 [batch_size, num_channels, seq_length, embedding_dim]。
|
||||
|
||||
参数:
|
||||
batch_data (dict): 输入的 batch 数据,格式为 {
|
||||
"title": [text1, text2, ...],
|
||||
"description": [text1, text2, ...],
|
||||
"tags": [text1, text2, ...],
|
||||
"author_info": [text1, text2, ...]
|
||||
}
|
||||
max_length (int): 最大序列长度。
|
||||
|
||||
返回:
|
||||
torch.Tensor: 形状为 [batch_size, num_channels, seq_length, embedding_dim] 的张量。
|
||||
"""
|
||||
# 1. 对每个通道的文本分别编码
|
||||
channel_embeddings = []
|
||||
model = StaticModel.from_pretrained("./model/embedding_256/")
|
||||
for channel in ["title", "description", "tags", "author_info"]:
|
||||
texts = batch_data[channel] # 获取当前通道的文本列表
|
||||
# 使用tokenizer将文本转换为tokens
|
||||
encoded_input = model.tokenizer(texts, padding=True, truncation=True, max_length=max_length, return_tensors='pt')
|
||||
with torch.no_grad():
|
||||
model_output = model.model(**encoded_input)
|
||||
# 提取最后一个隐藏层的结果
|
||||
embeddings = model_output.last_hidden_state.to(torch.float32) # 将embeddings 放在指定device上
|
||||
channel_embeddings.append(embeddings)
|
||||
|
||||
# 2. 将编码结果堆叠为 [batch_size, num_channels, seq_length, embedding_dim]
|
||||
batch_tensor = torch.stack(channel_embeddings, dim=1)
|
||||
return batch_tensor
|
95
filter/modelV6_0.py
Normal file
95
filter/modelV6_0.py
Normal file
@ -0,0 +1,95 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
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.channel_names = ['title', 'description', 'tags', 'author_info']
|
||||
|
||||
# CNN特征提取层
|
||||
self.conv_layers = nn.Sequential(
|
||||
# 第一层卷积
|
||||
nn.Conv2d(self.num_channels, 64, kernel_size=(3, 3), padding=1),
|
||||
nn.BatchNorm2d(64),
|
||||
nn.GELU(),
|
||||
nn.MaxPool2d(kernel_size=(2, 2)),
|
||||
|
||||
# 第二层卷积
|
||||
nn.Conv2d(64, 128, kernel_size=(3, 3), padding=1),
|
||||
nn.BatchNorm2d(128),
|
||||
nn.GELU(),
|
||||
nn.MaxPool2d(kernel_size=(2, 2)),
|
||||
|
||||
# 第三层卷积
|
||||
nn.Conv2d(128, 256, kernel_size=(3, 3), padding=1),
|
||||
nn.BatchNorm2d(256),
|
||||
nn.GELU(),
|
||||
)
|
||||
|
||||
# 计算卷积后的特征维度
|
||||
self.feature_dim = self._get_conv_output_size(seq_length, embedding_dim)
|
||||
|
||||
# 全连接层
|
||||
self.fc = nn.Sequential(
|
||||
nn.Linear(self.feature_dim, hidden_dim),
|
||||
nn.BatchNorm1d(hidden_dim),
|
||||
nn.Dropout(0.2),
|
||||
nn.GELU(),
|
||||
nn.Linear(hidden_dim, output_dim)
|
||||
)
|
||||
|
||||
self._init_weights()
|
||||
|
||||
def _get_conv_output_size(self, seq_length, embedding_dim):
|
||||
# 用于计算卷积输出尺寸
|
||||
x = torch.zeros(1, self.num_channels, seq_length, embedding_dim)
|
||||
x = self.conv_layers(x)
|
||||
return x.view(1, -1).size(1)
|
||||
|
||||
def _init_weights(self):
|
||||
for module in self.modules():
|
||||
if isinstance(module, nn.Conv2d) or isinstance(module, nn.Linear):
|
||||
nn.init.kaiming_normal_(module.weight, nonlinearity='relu')
|
||||
if module.bias is not None:
|
||||
nn.init.zeros_(module.bias)
|
||||
|
||||
def forward(self, channel_features: torch.Tensor):
|
||||
"""
|
||||
输入格式: [batch_size, num_channels, seq_length, embedding_dim]
|
||||
输出格式: [batch_size, output_dim]
|
||||
"""
|
||||
# CNN特征提取
|
||||
conv_features = self.conv_layers(channel_features)
|
||||
|
||||
# 展平特征
|
||||
flat_features = conv_features.view(conv_features.size(0), -1)
|
||||
|
||||
# 全连接层分类
|
||||
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
|
@ -4,14 +4,13 @@ import numpy as np
|
||||
from torch.utils.data import DataLoader
|
||||
import torch.optim as optim
|
||||
from dataset import MultiChannelDataset
|
||||
from filter.modelV3_10 import VideoClassifierV3_10, AdaptiveRecallLoss
|
||||
from sentence_transformers import SentenceTransformer
|
||||
from filter.modelV6_0 import VideoClassifierV6_0, AdaptiveRecallLoss
|
||||
from sklearn.metrics import f1_score, recall_score, precision_score, accuracy_score, classification_report
|
||||
import os
|
||||
import torch
|
||||
from torch.utils.tensorboard import SummaryWriter # 引入 TensorBoard
|
||||
import time
|
||||
from embedding import prepare_batch
|
||||
from embedding import prepare_batch_per_token
|
||||
|
||||
|
||||
# 动态生成子目录名称
|
||||
@ -52,9 +51,8 @@ class_weights = torch.tensor(
|
||||
)
|
||||
|
||||
# 初始化模型和SentenceTransformer
|
||||
sentence_transformer = SentenceTransformer("Thaweewat/jina-embedding-v3-m2v-1024")
|
||||
model = VideoClassifierV3_10()
|
||||
checkpoint_name = './filter/checkpoints/best_model_V3.11.pt'
|
||||
model = VideoClassifierV6_0()
|
||||
checkpoint_name = './filter/checkpoints/best_model_V6.0.pt'
|
||||
|
||||
# 模型保存路径
|
||||
os.makedirs('./filter/checkpoints', exist_ok=True)
|
||||
@ -78,7 +76,7 @@ def evaluate(model, dataloader):
|
||||
|
||||
with torch.no_grad():
|
||||
for batch in dataloader:
|
||||
batch_tensor = prepare_batch(batch['texts'], device="cpu")
|
||||
batch_tensor = prepare_batch_per_token(batch['texts'], max_length=1024)
|
||||
logits = model(batch_tensor)
|
||||
preds = torch.argmax(logits, dim=1)
|
||||
all_preds.extend(preds.cpu().numpy())
|
||||
@ -111,9 +109,8 @@ for epoch in range(num_epochs):
|
||||
for batch_idx, batch in enumerate(train_loader):
|
||||
optimizer.zero_grad()
|
||||
|
||||
batch_tensor = prepare_batch(batch['texts'], device="cpu")
|
||||
batch_tensor = prepare_batch_per_token(batch['texts'], max_length=1024)
|
||||
|
||||
# 传入文本字典和sentence_transformer
|
||||
logits = model(batch_tensor)
|
||||
|
||||
loss = criterion(logits, batch['label'])
|
||||
|
Loading…
Reference in New Issue
Block a user