| ''' | |
| * Adapted from BLIP (https://github.com/salesforce/BLIP) | |
| ''' | |
| import transformers | |
| transformers.logging.set_verbosity_error() | |
| from torch import nn | |
| import os | |
| from .med import BertConfig, BertModel | |
| from .blip import create_vit, init_tokenizer | |
| class BLIP_Pretrain(nn.Module): | |
| def __init__(self, | |
| med_config = "med_config.json", | |
| image_size = 224, | |
| vit = 'base', | |
| vit_grad_ckpt = False, | |
| vit_ckpt_layer = 0, | |
| embed_dim = 256, | |
| queue_size = 57600, | |
| momentum = 0.995, | |
| ): | |
| """ | |
| Args: | |
| med_config (str): path for the mixture of encoder-decoder model's configuration file | |
| image_size (int): input image size | |
| vit (str): model size of vision transformer | |
| """ | |
| super().__init__() | |
| self.visual_encoder, vision_width = create_vit(vit,image_size, vit_grad_ckpt, vit_ckpt_layer, 0) | |
| self.tokenizer = init_tokenizer() | |
| encoder_config = BertConfig.from_json_file(med_config) | |
| encoder_config.encoder_width = vision_width | |
| self.text_encoder = BertModel(config=encoder_config, add_pooling_layer=False) | |
| text_width = self.text_encoder.config.hidden_size | |
| self.vision_proj = nn.Linear(vision_width, embed_dim) | |
| self.text_proj = nn.Linear(text_width, embed_dim) | |