| | |
| | from sklearn.metrics import log_loss |
| | import torch.nn as nn |
| | import torch |
| | import math |
| | import numpy as np |
| | from torch.nn.utils.rnn import pad_sequence |
| | import torch.nn.functional as F |
| | from .transformer import * |
| | import torchvision.models as models |
| | from einops import rearrange |
| | from transformers import AutoModel |
| |
|
| | """ |
| | args.N |
| | args.d_model |
| | args.res_base_model |
| | args.H |
| | args.num_queries |
| | args.dropout |
| | args.attribute_set_size |
| | """ |
| |
|
| |
|
| | class MeDSLIP(nn.Module): |
| | def __init__(self, config, pathology_book): |
| | super(MeDSLIP, self).__init__() |
| |
|
| | self.d_model = config["d_model"] |
| | with torch.no_grad(): |
| | bert_model = self._get_bert_basemodel( |
| | config["text_encoder"], freeze_layers=None |
| | ).to(pathology_book["input_ids"].device) |
| | self.pathology_book = bert_model( |
| | input_ids=pathology_book["input_ids"], |
| | attention_mask=pathology_book["attention_mask"], |
| | ) |
| | self.pathology_book = self.pathology_book.last_hidden_state[:, 0, :] |
| | self.pathology_embedding_layer = nn.Linear(768, 256) |
| |
|
| | self.cl_fc_pathology = nn.Linear(256, 768) |
| | self.cl_fc_anatomy = nn.Linear(256, 768) |
| |
|
| | """ visual backbone""" |
| | self.resnet_dict = { |
| | "resnet18": models.resnet18(pretrained=False), |
| | "resnet50": models.resnet50(pretrained=False), |
| | } |
| | resnet = self._get_res_basemodel(config["res_base_model"]) |
| | num_ftrs = int(resnet.fc.in_features / 2) |
| | self.res_features = nn.Sequential(*list(resnet.children())[:-3]) |
| | self.res_l1_anatomy = nn.Linear(num_ftrs, num_ftrs) |
| | self.res_l2_anatomy = nn.Linear(num_ftrs, self.d_model) |
| | self.res_l1_pathology = nn.Linear(num_ftrs, num_ftrs) |
| | self.res_l2_pathology = nn.Linear(num_ftrs, self.d_model) |
| |
|
| | self.mask_generator = nn.Linear(num_ftrs, num_ftrs) |
| |
|
| | |
| | """ Query Decoder""" |
| | |
| |
|
| | self.H = config["H"] |
| | decoder_layer = TransformerDecoderLayer( |
| | self.d_model, config["H"], 1024, 0.1, "relu", normalize_before=True |
| | ) |
| | decoder_norm = nn.LayerNorm(self.d_model) |
| | self.decoder_anatomy = TransformerDecoder( |
| | decoder_layer, config["N"], decoder_norm, return_intermediate=False |
| | ) |
| | self.decoder_pathology = TransformerDecoder( |
| | decoder_layer, config["N"], decoder_norm, return_intermediate=False |
| | ) |
| |
|
| | |
| | self.dropout_feas_anatomy = nn.Dropout(config["dropout"]) |
| | self.dropout_feas_pathology = nn.Dropout(config["dropout"]) |
| |
|
| | |
| | self.classifier_anatomy = nn.Linear(self.d_model, config["attribute_set_size"]) |
| | self.classifier_pathology = nn.Linear(self.d_model, config["attribute_set_size"]) |
| |
|
| | self.apply(self._init_weights) |
| |
|
| | def _get_res_basemodel(self, res_model_name): |
| | try: |
| | res_model = self.resnet_dict[res_model_name] |
| | print("Image feature extractor:", res_model_name) |
| | return res_model |
| | except: |
| | raise ( |
| | "Invalid model name. Check the config file and pass one of: resnet18 or resnet50" |
| | ) |
| |
|
| | def _get_bert_basemodel(self, bert_model_name, freeze_layers): |
| | try: |
| | model = AutoModel.from_pretrained(bert_model_name) |
| | print("text feature extractor:", bert_model_name) |
| | except: |
| | raise ( |
| | "Invalid model name. Check the config file and pass a BERT model from transformers lybrary" |
| | ) |
| |
|
| | if freeze_layers is not None: |
| | for layer_idx in freeze_layers: |
| | for param in list(model.encoder.layer[layer_idx].parameters()): |
| | param.requires_grad = False |
| | return model |
| |
|
| | def image_encoder(self, xis): |
| | |
| | """ |
| | 16 torch.Size([16, 1024, 14, 14]) |
| | torch.Size([16, 196, 1024]) |
| | torch.Size([3136, 1024]) |
| | torch.Size([16, 196, 256]) |
| | """ |
| | batch_size = xis.shape[0] |
| | res_fea = self.res_features(xis) |
| | res_fea = rearrange(res_fea, "b d n1 n2 -> b (n1 n2) d") |
| | x = rearrange(res_fea, "b n d -> (b n) d") |
| |
|
| | masks = self.mask_generator(x) |
| | x_pathology = masks * x |
| | x_anatomy = (1 - masks) * x |
| |
|
| | x_pathology = self.res_l1_pathology(x_pathology) |
| | x_anatomy = self.res_l1_anatomy(x_anatomy) |
| | x_pathology = F.relu(x_pathology) |
| | x_anatomy = F.relu(x_anatomy) |
| |
|
| | x_pathology = self.res_l2_pathology(x_pathology) |
| | x_anatomy = self.res_l2_anatomy(x_anatomy) |
| |
|
| | out_emb_pathology = rearrange(x_pathology, "(b n) d -> b n d", b=batch_size) |
| | out_emb_anatomy = rearrange(x_anatomy, "(b n) d -> b n d", b=batch_size) |
| | return out_emb_pathology, out_emb_anatomy |
| |
|
| | def forward(self, images): |
| | B = images.shape[0] |
| |
|
| | device = images.device |
| | """ Visual Backbone """ |
| | x, _ = self.image_encoder(images) |
| | features = x.transpose(0, 1) |
| |
|
| | query_embed = self.pathology_embedding_layer(self.pathology_book) |
| | query_embed = query_embed.unsqueeze(1).repeat(1, B, 1) |
| | features, ws = self.decoder_pathology( |
| | query_embed, |
| | features, |
| | memory_key_padding_mask=None, |
| | pos=None, |
| | query_pos=None, |
| | ) |
| | features = self.dropout_feas_pathology(features) |
| | x = self.classifier_pathology(features).transpose(0, 1) |
| |
|
| | return x |
| |
|
| | @staticmethod |
| | def _init_weights(module): |
| | r"""Initialize weights like BERT - N(0.0, 0.02), bias = 0.""" |
| |
|
| | if isinstance(module, nn.Linear): |
| | module.weight.data.normal_(mean=0.0, std=0.02) |
| |
|
| | elif isinstance(module, nn.MultiheadAttention): |
| | module.in_proj_weight.data.normal_(mean=0.0, std=0.02) |
| | module.out_proj.weight.data.normal_(mean=0.0, std=0.02) |
| |
|
| | elif isinstance(module, nn.Embedding): |
| | module.weight.data.normal_(mean=0.0, std=0.02) |
| | if module.padding_idx is not None: |
| | module.weight.data[module.padding_idx].zero_() |
