Merge pull request #60 from jmisilo/57-better-clip

Change CLIP version
jmisilo · Nov 14, 2022 · e561112 · e561112
2 parents a9100e3 + c0605e3
commit e561112
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/src/dataset_generation.py b/src/dataset_generation.py
@@ -28,8 +28,8 @@
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
  # Load CLIP model and processor
- preprocessor = CLIPProcessor.from_pretrained('openai/clip-vit-base-patch32')
- model = CLIPModel.from_pretrained('openai/clip-vit-base-patch32').vision_model.to(device)
+ preprocessor = CLIPProcessor.from_pretrained('openai/clip-vit-base-patch14')
+ model = CLIPModel.from_pretrained('openai/clip-vit-base-patch14').vision_model.to(device)
 
  # Load dataset
  df = pd.read_csv(os.path.join(DATA_PATH, 'raw', 'results.csv'), sep='|')

diff --git a/src/model/model.py b/src/model/model.py
@@ -16,8 +16,8 @@ def __init__(self, device='cpu'):
 
  self.device = device
 
- self.preprocessor = CLIPProcessor.from_pretrained('openai/clip-vit-base-patch32')
- self.model = CLIPModel.from_pretrained('openai/clip-vit-base-patch32').vision_model.to(self.device)
+ self.preprocessor = CLIPProcessor.from_pretrained('openai/clip-vit-base-patch14')
+ self.model = CLIPModel.from_pretrained('openai/clip-vit-base-patch14').vision_model.to(self.device)
 
  def forward(self, image):
  # only one image at a time