import torch import torch.nn as nn import torch.nn.functional as F from transformers import AutoModel, AutoModelForZeroShotImageClassification, AutoProcessor, AutoConfig import random # Probability of keeping the text branch for a sample during training (text dropout = 1 - this). TEXT_KEEP_PROB = 0.5 class Attn_Net_Gated(nn.Module): def __init__(self, L=1024, D=256, dropout=0.25, n_classes=1): super(Attn_Net_Gated, self).__init__() self.attention_a = nn.Sequential( nn.Linear(L, D), # L is the visual_projection_dim nn.Tanh() ) self.attention_b = nn.Sequential( nn.Linear(L, D), nn.Sigmoid() ) if dropout: self.attention_a.add_module("Dropout", nn.Dropout(dropout)) self.attention_b.add_module("Dropout", nn.Dropout(dropout)) self.attention_c = nn.Linear(D, n_classes) def forward(self, x): a = self.attention_a(x) b = self.attention_b(x) A = a * b # Element-wise multiplication for gated attention A = self.attention_c(A) return A, x class Attn_Net(nn.Module): def __init__(self, L=1024, D=256, dropout=0.25, n_classes=1): super(Attn_Net, self).__init__() self.attention = nn.Sequential( nn.Linear(L, D), nn.Tanh(), # Just a single non-linearity for attention scores ) if dropout: self.attention.add_module("Dropout", nn.Dropout(dropout)) self.attention_c = nn.Linear(D, n_classes) def forward(self, x): A = self.attention(x) # Single attention branch A = self.attention_c(A) return A, x class CLAM_ViT(nn.Module): def __init__(self, base_model_name, training_config, gate=True, size_arg="small", dropout=0.25, k_sample=8, device="cuda", freeze_vit=False, freeze_query_features_encoder=False, use_cell_type_embedding=False): super(CLAM_ViT, self).__init__() # Load pre-trained ViT for patch feature extraction config = AutoConfig.from_pretrained(base_model_name) self.patch_encoder = AutoModel.from_pretrained(base_model_name, config=config, use_safetensors=True) self.patch_processor = AutoProcessor.from_pretrained(base_model_name) self.freeze_vit = freeze_vit self.config = training_config self.freeze_query_features_encoder = freeze_query_features_encoder self.use_cell_type_embedding = use_cell_type_embedding self.device = device self.patch_encoder.to(self.device) # Unfreeze the patch encoder (ViT) if self.freeze_vit: for param in self.patch_encoder.vision_model.parameters(): param.requires_grad = False else: for param in self.patch_encoder.vision_model.parameters(): param.requires_grad = True if self.freeze_query_features_encoder: for param in self.patch_encoder.text_model.parameters(): param.requires_grad = False else: for param in self.patch_encoder.text_model.parameters(): param.requires_grad = True self.visual_projection_dim = self.patch_encoder.vision_model.encoder.layers[-1].mlp.fc2.out_features self.text_projection_dim = self.patch_encoder.text_model.encoder.layers[-1].mlp.fc2.out_features size_dict = {"tiny": [self.visual_projection_dim, 128], "small": [self.visual_projection_dim, 256], "big": [self.visual_projection_dim, 384]} size = size_dict[size_arg] if self.use_cell_type_embedding: # Replace the Embedding layer with a Linear layer self.cell_type_projection = nn.Linear(39, self.visual_projection_dim) # self.cell_type_projection.to(self.device) # self.wsi_vision_final_projection = nn.Linear(patch_dim, size[1]) # self.wsi_vision_final_projection.to(self.device) # Attention network (gated or non-gated) if gate: self.attention_net = Attn_Net_Gated(L=size[0], D=size[1], dropout=dropout, n_classes=1) else: self.attention_net = Attn_Net(L=size[0], D=size[1], dropout=dropout, n_classes=1) self.k_sample = k_sample # Replace average pooling with a linear layer self.visual_token_projection = nn.Linear(self.visual_projection_dim, self.visual_projection_dim) # Add a projection layer from 512 to 768 self.text_projection_to_visual_dim = nn.Linear(self.text_projection_dim, self.visual_projection_dim) # Task-specific heads: staining intensity, location, quantity self.classifiers = nn.ModuleList([ nn.Linear(size[0], 4), # staining intensity nn.Linear(size[0], 4), # staining location nn.Linear(size[0], 4), # staining quantity nn.Linear(size[0], 58), # tissue type nn.Linear(size[0], 2) # tumor vs non-tumor ]) # Add L2 regularization self.l2_reg = 1e-5 def forward(self, patches, query_input, cell_type_one_hot, phase="test"): # Process patches in smaller chunks to save memory chunk_size = 32 # Adjust this value based on your GPU memory patch_features_list = [] patch_counts = [] for ix, patch_tensor in enumerate(patches): # Process patches in chunks num_patches = patch_tensor.size(0) chunk_features = [] for i in range(0, num_patches, chunk_size): chunk = patch_tensor[i:i + chunk_size] with torch.cuda.amp.autocast(): # Use mixed precision patch_outputs = self.patch_encoder.vision_model(chunk, output_hidden_states=True) patch_tokens = patch_outputs.hidden_states[-1] chunk_features.append(self.visual_token_projection(patch_tokens)) # Concatenate chunks patch_features = torch.cat(chunk_features, dim=0) patch_features_list.append(patch_features) patch_counts.append(patch_features.shape[0]) # Attention mechanism to weigh patch features A_list = [] h_list = [] for features in patch_features_list: # features: [N patches x N tokens x feature dim], for example, [115, 50, 768] A, h = self.attention_net(features) # here `h` is the same as `features` A_list.append(A) h_list.append(h) # A_list: [torch.Size([115, 50, 1]), torch.Size([111, 50, 1]), ..., torch.Size([104, 50,1])] # h_list: [torch.Size([115, 50, 768]), torch.Size([111, 50, 768]), ..., torch.Size([104, 50, 768])] # Process each sample separately M_list = [] A_raw_list = [] for A, h_value, count in zip(A_list, h_list, patch_counts): A = A[:count] # Only keep attention weights for real patches A_raw = A.clone() A = F.softmax(A, dim=0) # Softmax over patches for this sample # (115, 50, 1) # M = torch.mm(A.t(), h[:count]) # Weighted sum of features (1, 115) x (115, 512) M = (A * h_value).sum(dim=0) # [50, 768] M_list.append(M) A_raw_list.append(A_raw) # Stack results A_raw = torch.cat(A_raw_list, dim=0) M = torch.stack(M_list) # 16 x 50 x 768 # All-token aggregation: mean over the token axis. M_single_token = torch.mean(M, dim=1) # (B, D) # The text (context) branch is used during training; `force_query` (default False) # switches it on at inference as well. _force_q = getattr(self, "force_query", False) if phase == "train" or _force_q: text_inputs = self.patch_processor.tokenizer( query_input, return_tensors="pt", padding=True, truncation=True, max_length=80) text_inputs.to(self.device) query_features = self.patch_encoder.get_text_features(**text_inputs) # (B, 512) query_features = self.text_projection_to_visual_dim(query_features) # (B, D) # Text dropout, per sample: each image keeps its text features with probability # TEXT_KEEP_PROB. `force_query` keeps the branch on for every sample. if _force_q: keep = torch.ones(query_features.size(0), 1, device=query_features.device) else: keep = (torch.rand(query_features.size(0), 1, device=query_features.device) < TEXT_KEEP_PROB).float() M_single_token = M_single_token + keep * query_features if self.use_cell_type_embedding: cell_type_one_hot = cell_type_one_hot.to(self.device) # Use the linear projection instead of embedding cell_type_embed = self.cell_type_projection(cell_type_one_hot) # one embedding row per sample assert cell_type_embed.size(0) == M_single_token.size(0), ( f"cell_type_one_hot has {cell_type_embed.size(0)} rows but the batch has " f"{M_single_token.size(0)}") M_single_token = M_single_token + cell_type_embed else: pass # Task-specific outputs outputs = [classifier(M_single_token) for classifier in self.classifiers] """ Just use raw outputs, do not apply sigmoid or softmax! Because we will use CrossEntropy Loss, which already contains sigmoid If you apply sigmoid here, it will be a double sigmoid, and will destroy the gradient flow. Since the loss function expects raw logits and not sigmoid outputs, the calculated loss will be incorrect, and this will negatively impact the learning process. """ intensity_out = outputs[0] location_out = outputs[1] quantity_out = outputs[2] tissue_out = outputs[3] malignancy_out = outputs[4] return intensity_out, location_out, quantity_out, tissue_out, malignancy_out, A_raw def print_last_layer_weights(self): print("Vision model last layer weights:") vision_last_layer = self.patch_encoder.vision_model.encoder.layers[-1].mlp.fc2 print(vision_last_layer.weight.data) print("\nText model last layer weights:") text_last_layer = self.patch_encoder.text_model.encoder.layers[-1].mlp.fc2 print(text_last_layer.weight.data) def get_l2_regularization(self): l2_loss = 0 for param in self.parameters(): l2_loss += torch.norm(param, 2) return self.l2_reg * l2_loss def get_model_state(self): return { 'attention_net': self.attention_net.state_dict(), 'wsi_vision_final_projection': self.wsi_vision_final_projection.state_dict() if hasattr(self, 'wsi_vision_final_projection') else None, 'patch_encoder': self.patch_encoder.state_dict(), 'classifiers': [clf.state_dict() for clf in self.classifiers] } # Add a function to create an optimizer with learning rate scheduler def create_optimizer_and_scheduler(model, lr=1e-4, weight_decay=1e-5): optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=500) return optimizer, scheduler def create_clam_vit(base_model_name, training_config, gate=True, size_arg="small", dropout=0.25, k_sample=8, n_classes=4, device="cuda", lr=1e-4, weight_decay=1e-5, freeze_vit=False, freeze_query_features_encoder=False, use_cell_type_embedding=False): model = CLAM_ViT(base_model_name=base_model_name, training_config=training_config, gate=gate, size_arg=size_arg, dropout=dropout, k_sample=k_sample, device=device, freeze_vit=freeze_vit, freeze_query_features_encoder=freeze_query_features_encoder, use_cell_type_embedding=use_cell_type_embedding) optimizer, scheduler = create_optimizer_and_scheduler(model, lr=lr, weight_decay=weight_decay) return model, optimizer, scheduler