liwenyun
/
CLIP
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add test - `test_hpu_support`

This commit is contained in:
PiotrBLL 2024-12-19 02:05:11 +01:00
parent 5393fdde9c
commit deb2964c26
1 changed files with 20 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
tests/test_consistency.py
View File

@ -6,7 +6,7 @@ from PIL import Image
import clip
@pytest.mark.parametrize('model_name', clip.available_models())
@pytest.mark.parametrize("model_name", clip.available_models())
def test_consistency(model_name):
device = "cpu"
jit_model, transform = clip.load(model_name, device=device, jit=True)
@ -23,3 +23,22 @@ def test_consistency(model_name):
py_probs = logits_per_image.softmax(dim=-1).cpu().numpy()
assert np.allclose(jit_probs, py_probs, atol=0.01, rtol=0.1)
@pytest.mark.parametrize("model_name", clip.available_models())
def test_hpu_support(model_name):
device = "hpu"
jit_model, transform = clip.load(model_name, device=device, jit=True)
py_model, _ = clip.load(model_name, device=device, jit=False)
image = transform(Image.open("CLIP.png")).unsqueeze(0).to(device)
text = clip.tokenize(["a diagram", "a dog", "a cat"]).to(device)
with torch.no_grad():
logits_per_image, _ = jit_model(image, text)
jit_probs = logits_per_image.softmax(dim=-1).cpu().numpy()
logits_per_image, _ = py_model(image, text)
py_probs = logits_per_image.softmax(dim=-1).cpu().numpy()
assert np.allclose(jit_probs, py_probs, atol=0.01, rtol=0.1)