| import torch | |
| import torch.nn as nn | |
| from transformers import CLIPModel | |
| from peft import LoraConfig, get_peft_model | |
| class MLP(nn.Module): | |
| def __init__(self, input_dim=768, hidden_dim1=512, hidden_dim2=256, output_dim=8,dropout_rate=0.5): | |
| super(MLP, self).__init__() | |
| self.fc1 = nn.Linear(input_dim, hidden_dim1) | |
| self.relu1 = nn.ReLU() | |
| self.dropout = nn.Dropout(dropout_rate) | |
| self.fc2 = nn.Linear(hidden_dim1, hidden_dim2) | |
| self.relu2 = nn.ReLU() | |
| self.fc3 = nn.Linear(hidden_dim2, output_dim) | |
| def forward(self, x): | |
| x = self.fc1(x) | |
| x = self.relu1(x) | |
| x = self.dropout(x) | |
| x = self.fc2(x) | |
| x = self.relu2(x) | |
| x = self.dropout(x) | |
| x = self.fc3(x) | |
| return x | |
| class clip_lora_model(nn.Module): | |
| def __init__(self, input_dim=768, hidden_dim1=512, hidden_dim2=256, output_dim=8,dropout_rate=0.5,r=16,lora_alpha=8): | |
| super(clip_lora_model, self).__init__() | |
| self.output_dim=output_dim | |
| self.mlp = MLP(input_dim, hidden_dim1, hidden_dim2, output_dim,dropout_rate) | |
| model_name = 'openai/clip-vit-large-patch14' | |
| model = CLIPModel.from_pretrained(model_name) | |
| self.proj = model.visual_projection | |
| for param in self.proj.parameters(): | |
| param.requires_grad = False | |
| encoder = model.vision_model | |
| target_modules = ["k_proj", "v_proj", "q_proj"] | |
| config = LoraConfig( | |
| r=int(r), | |
| lora_alpha=lora_alpha, | |
| target_modules=target_modules, | |
| lora_dropout=0.1, | |
| bias="none", | |
| ) | |
| self.model = get_peft_model(encoder, config) | |
| def forward(self, x): | |
| model_outputs = self.model(x) | |
| image_embeds = model_outputs[1] | |
| model_outputs = self.proj(image_embeds) | |
| outputs = self.mlp(model_outputs) | |
| return outputs | |