github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

fixed batched clip computation (#6292)

Browse Source
This commit is contained in:
Tobias Fischer
2024-08-26 20:48:15 -04:00
committed by GitHub
parent 3918f6eea0
commit 211bfb6d8a
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
extra/models/clip.py
View File

@@ -451,9 +451,9 @@ class OpenClipEncoder:
def get_clip_score(self, tokens:Tensor, image:Tensor) -> Tensor:
image_features: Tensor = self.visual(image)
image_features /= image_features.square().sum([-1,-2], keepdim=True).sqrt() # Frobenius Norm
image_features /= image_features.square().sum(-1, keepdim=True).sqrt() # Frobenius Norm
text_features = self.encode_tokens(tokens).squeeze(0)
text_features /= text_features.square().sum([-1,-2], keepdim=True).sqrt() # Frobenius Norm
text_features = self.encode_tokens(tokens)
text_features /= text_features.square().sum(-1, keepdim=True).sqrt() # Frobenius Norm
return image_features @ text_features.T
return (image_features * text_features).sum(axis=-1)