# modified from https://github.com/tensorflow/models/blob/master/research/slim/nets/s3dg.py 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 ) # Learnable Queries self.dropout_feas_anatomy = nn.Dropout(config["dropout"]) self.dropout_feas_pathology = nn.Dropout(config["dropout"]) # Attribute classifier 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): # patch features """ 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) # batch_size,feature_size,patch_num,patch_num 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) # batch_size,patch_num,dim features = x.transpose(0, 1) # patch_num b dim 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) # B query Atributes 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_()