only use clip-vision-model

pull/45/head
chongkuiqi 2025-03-03 14:51:54 +08:00
parent c1a069b532
commit b696a7f084
1 changed files with 3 additions and 3 deletions

View File

@ -3,7 +3,7 @@ from .model import *
from typing import Optional, Tuple, List from typing import Optional, Tuple, List
from torch import nn from torch import nn
import warnings import warnings
from transformers import CLIPProcessor, CLIPModel from transformers import CLIPProcessor, CLIPVisionModel
import torch import torch
warnings.filterwarnings('ignore') warnings.filterwarnings('ignore')
@ -36,7 +36,7 @@ class MiniMindVLM(MiniMindLM):
@staticmethod @staticmethod
def get_vision_model(model_path="./model/vision_model/clip-vit-base-patch16"): def get_vision_model(model_path="./model/vision_model/clip-vit-base-patch16"):
model = CLIPModel.from_pretrained(model_path) model = CLIPVisionModel.from_pretrained(model_path)
processor = CLIPProcessor.from_pretrained(model_path) processor = CLIPProcessor.from_pretrained(model_path)
# 冻结 vision_encoder 的所有参数 # 冻结 vision_encoder 的所有参数
for param in model.parameters(): for param in model.parameters():
@ -52,7 +52,7 @@ class MiniMindVLM(MiniMindLM):
@staticmethod @staticmethod
def get_image_embeddings(image_tensors, vision_model): def get_image_embeddings(image_tensors, vision_model):
with torch.no_grad(): with torch.no_grad():
outputs = vision_model.vision_model(pixel_values=image_tensors) outputs = vision_model(pixel_values=image_tensors)
img_embedding = outputs.last_hidden_state[:, 1:, :].squeeze() img_embedding = outputs.last_hidden_state[:, 1:, :].squeeze()
return img_embedding return img_embedding