Upload Model.py

Model.py

@@ -0,0 +1,184 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from pathlib import Path
+import json
+import pandas as pd
+import numpy as np
+from tqdm import tqdm
+from sklearn.preprocessing import LabelEncoder
+import pickle
+from typing import Dict, List, Optional
+import warnings
+import random
+# ========================= 步骤3: Inter-Task Attention模型 =========================
+
+class InterTaskAttention(nn.Module):
+    """
+    Inter-Task Attention机制
+    学习任务间的相互关系和依赖
+    """
+    
+    def __init__(self, hidden_dim: int, num_tasks: int, num_heads: int = 4):
+        super().__init__()
+        self.hidden_dim = hidden_dim
+        self.num_tasks = num_tasks
+        self.num_heads = num_heads
+        
+        # Multi-head attention for tasks
+        self.task_attention = nn.MultiheadAttention(
+            embed_dim=hidden_dim,
+            num_heads=num_heads,
+            dropout=0.1,
+            batch_first=True
+        )
+        
+        # Task-specific query, key, value projections
+        self.task_query = nn.Linear(hidden_dim, hidden_dim)
+        self.task_key = nn.Linear(hidden_dim, hidden_dim)
+        self.task_value = nn.Linear(hidden_dim, hidden_dim)
+        
+        # Layer normalization
+        self.norm1 = nn.LayerNorm(hidden_dim)
+        self.norm2 = nn.LayerNorm(hidden_dim)
+        
+        # Feed-forward network
+        self.ffn = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim * 2),
+            nn.ReLU(),
+            nn.Dropout(0.1),
+            nn.Linear(hidden_dim * 2, hidden_dim)
+        )
+    
+    def forward(self, task_features):
+        """
+        Args:
+            task_features: [batch_size, num_tasks, hidden_dim]
+        Returns:
+            refined_features: [batch_size, num_tasks, hidden_dim]
+        """
+        # Self-attention across tasks
+        q = self.task_query(task_features)
+        k = self.task_key(task_features)
+        v = self.task_value(task_features)
+        
+        attended_features, attention_weights = self.task_attention(q, k, v)
+        
+        # Residual connection + normalization
+        task_features = self.norm1(task_features + attended_features)
+        
+        # Feed-forward network
+        ffn_output = self.ffn(task_features)
+        task_features = self.norm2(task_features + ffn_output)
+        
+        return task_features, attention_weights
+
+
+class OmniPathWithInterTaskAttention(nn.Module):
+    """
+    OmniPath模型 + Inter-Task Attention
+    从预提取的特征进行多任务学习
+    """
+    
+    def __init__(self,
+                 label_mappings: Dict,
+                 feature_dim: int = 1024,
+                 hidden_dim: int = 256,
+                 dropout: float = 0.3,
+                 use_inter_task_attention: bool = True,
+                 inter_task_heads: int = 4):
+        super().__init__()
+        
+        self.label_mappings = label_mappings
+        self.num_tasks = len(label_mappings)
+        self.use_inter_task_attention = use_inter_task_attention
+        
+        # Tile-level feature encoder
+        self.tile_encoder = nn.Sequential(
+            nn.Linear(feature_dim, hidden_dim),
+            nn.LayerNorm(hidden_dim),
+            nn.ReLU(),
+            nn.Dropout(dropout)
+        )
+        
+        # Tile attention (for aggregating tiles to patient-level)
+        self.tile_attention = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim // 2),
+            nn.Tanh(),
+            nn.Linear(hidden_dim // 2, 1)
+        )
+        
+        # Task-specific encoders (before inter-task attention)
+        self.task_encoders = nn.ModuleDict()
+        for task_name in label_mappings.keys():
+            self.task_encoders[task_name] = nn.Sequential(
+                nn.Linear(hidden_dim, hidden_dim),
+                nn.LayerNorm(hidden_dim),
+                nn.ReLU(),
+                nn.Dropout(dropout)
+            )
+        
+        # Inter-Task Attention
+        if use_inter_task_attention:
+            self.inter_task_attention = InterTaskAttention(
+                hidden_dim=hidden_dim,
+                num_tasks=self.num_tasks,
+                num_heads=inter_task_heads
+            )
+        
+        # Task-specific prediction heads (after inter-task attention)
+        self.task_heads = nn.ModuleDict()
+        for task_name, mapping in label_mappings.items():
+            self.task_heads[task_name] = nn.Sequential(
+                nn.Linear(hidden_dim, hidden_dim // 2),
+                nn.ReLU(),
+                nn.Dropout(dropout),
+                nn.Linear(hidden_dim // 2, mapping['num_classes'])
+            )
+    
+    def forward(self, features, return_attention=False):
+        """
+        Args:
+            features: [batch_size, num_tiles, feature_dim]
+        Returns:
+            outputs: dict of task predictions
+        """
+        batch_size, num_tiles, _ = features.shape
+        
+        # Encode tile features
+        tile_features = self.tile_encoder(features)  # [B, N, H]
+        
+        # Compute tile attention weights
+        attention_logits = self.tile_attention(tile_features)  # [B, N, 1]
+        attention_weights = F.softmax(attention_logits, dim=1)  # [B, N, 1]
+        
+        # Aggregate tiles to patient-level
+        patient_features = torch.sum(tile_features * attention_weights, dim=1)  # [B, H]
+        
+        # Task-specific encoding
+        task_features_list = []
+        task_names = list(self.label_mappings.keys())
+        
+        for task_name in task_names:
+            task_feat = self.task_encoders[task_name](patient_features)  # [B, H]
+            task_features_list.append(task_feat.unsqueeze(1))  # [B, 1, H]
+        
+        task_features = torch.cat(task_features_list, dim=1)  # [B, num_tasks, H]
+        
+        # Inter-Task Attention
+        inter_task_attn_weights = None
+        if self.use_inter_task_attention:
+            task_features, inter_task_attn_weights = self.inter_task_attention(task_features)
+        
+        # Task-specific predictions
+        outputs = {}
+        for i, task_name in enumerate(task_names):
+            task_feat = task_features[:, i, :]  # [B, H]
+            outputs[task_name] = self.task_heads[task_name](task_feat)  # [B, num_classes]
+        
+        if return_attention:
+            outputs['tile_attention'] = attention_weights
+            outputs['inter_task_attention'] = inter_task_attn_weights
+        
+        return outputs