116 lines
3.4 KiB
Python
116 lines
3.4 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
from models.image_models import VisionTransformer
|
||
|
||
class MultiModalFusionModel(nn.Module):
|
||
"""多模态融合模型，融合DIFF和WNB散点图的特征"""
|
||
def __init__(self, img_size=224, patch_size=16, in_channels=3,
|
||
embed_dim=768, depth=6, num_heads=12, dropout=0.1, num_classes=2):
|
||
super().__init__()
|
||
|
||
# DIFF散点图特征提取器
|
||
self.diff_encoder = VisionTransformer(
|
||
img_size=img_size,
|
||
patch_size=patch_size,
|
||
in_channels=in_channels,
|
||
embed_dim=embed_dim,
|
||
depth=depth,
|
||
num_heads=num_heads,
|
||
dropout=dropout
|
||
)
|
||
|
||
# WNB散点图特征提取器
|
||
self.wnb_encoder = VisionTransformer(
|
||
img_size=img_size,
|
||
patch_size=patch_size,
|
||
in_channels=in_channels,
|
||
embed_dim=embed_dim,
|
||
depth=depth,
|
||
num_heads=num_heads,
|
||
dropout=dropout
|
||
)
|
||
|
||
# 特征融合层
|
||
self.fusion = nn.Sequential(
|
||
nn.Linear(embed_dim * 2, embed_dim),
|
||
nn.LayerNorm(embed_dim),
|
||
nn.GELU(),
|
||
nn.Dropout(dropout)
|
||
)
|
||
|
||
# 分类头
|
||
self.classifier = nn.Linear(embed_dim, num_classes)
|
||
|
||
def forward(self, diff_img, wnb_img):
|
||
"""
|
||
前向传播
|
||
|
||
Args:
|
||
diff_img: DIFF散点图 [B, C, H, W]
|
||
wnb_img: WNB散点图 [B, C, H, W]
|
||
|
||
Returns:
|
||
logits: 分类logits [B, num_classes]
|
||
"""
|
||
# 提取特征
|
||
diff_features = self.diff_encoder(diff_img) # [B, E]
|
||
wnb_features = self.wnb_encoder(wnb_img) # [B, E]
|
||
|
||
# 特征融合
|
||
combined_features = torch.cat([diff_features, wnb_features], dim=1) # [B, 2*E]
|
||
fused_features = self.fusion(combined_features) # [B, E]
|
||
|
||
# 分类
|
||
logits = self.classifier(fused_features) # [B, num_classes]
|
||
|
||
return logits
|
||
|
||
def extract_features(self, diff_img, wnb_img):
|
||
"""提取各个模态的特征，用于分析"""
|
||
diff_features = self.diff_encoder(diff_img)
|
||
wnb_features = self.wnb_encoder(wnb_img)
|
||
|
||
return {
|
||
'diff': diff_features,
|
||
'wnb': wnb_features
|
||
}
|
||
|
||
|
||
class SingleModalModel(nn.Module):
|
||
"""单模态模型，用于对比实验"""
|
||
def __init__(self, img_size=224, patch_size=16, in_channels=3,
|
||
embed_dim=768, depth=6, num_heads=12, dropout=0.1, num_classes=2):
|
||
super().__init__()
|
||
|
||
# 图像特征提取器
|
||
self.encoder = VisionTransformer(
|
||
img_size=img_size,
|
||
patch_size=patch_size,
|
||
in_channels=in_channels,
|
||
embed_dim=embed_dim,
|
||
depth=depth,
|
||
num_heads=num_heads,
|
||
dropout=dropout
|
||
)
|
||
|
||
# 分类头
|
||
self.classifier = nn.Linear(embed_dim, num_classes)
|
||
|
||
def forward(self, img):
|
||
"""
|
||
前向传播
|
||
|
||
Args:
|
||
img: 输入图像 [B, C, H, W]
|
||
|
||
Returns:
|
||
logits: 分类logits [B, num_classes]
|
||
"""
|
||
# 提取特征
|
||
features = self.encoder(img) # [B, E]
|
||
|
||
# 分类
|
||
logits = self.classifier(features) # [B, num_classes]
|
||
|
||
return logits |