109 lines
2.9 KiB
Python
109 lines
2.9 KiB
Python
"""
|
|
Simple test to verify OpenCLIP GPU functionality
|
|
"""
|
|
import subprocess
|
|
import os
|
|
|
|
def test_openclip_gpu():
|
|
"""Test OpenCLIP GPU functionality with a simple script"""
|
|
venv_python = "openclip_gpu_env\\Scripts\\python.exe"
|
|
|
|
if not os.path.exists(venv_python):
|
|
print("ERROR: Virtual environment not found")
|
|
return False
|
|
|
|
# Simple test script
|
|
test_script = """
|
|
import torch
|
|
import open_clip
|
|
from PIL import Image
|
|
import tempfile
|
|
import os
|
|
|
|
print("Testing OpenCLIP GPU...")
|
|
print(f"PyTorch version: {torch.__version__}")
|
|
print(f"CUDA available: {torch.cuda.is_available()}")
|
|
|
|
if torch.cuda.is_available():
|
|
print(f"GPU: {torch.cuda.get_device_name(0)}")
|
|
else:
|
|
print("ERROR: CUDA not available")
|
|
exit(1)
|
|
|
|
# Load model
|
|
print("Loading OpenCLIP model...")
|
|
model, _, preprocess = open_clip.create_model_and_transforms(
|
|
model_name="ViT-B-32",
|
|
pretrained="laion2b_s34b_b79k"
|
|
)
|
|
|
|
print("Model loaded successfully")
|
|
|
|
# Move to GPU
|
|
model = model.cuda()
|
|
print("Model moved to GPU")
|
|
|
|
# Create a test image (simple red square)
|
|
with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as f:
|
|
img_path = f.name
|
|
|
|
img = Image.new('RGB', (224, 224), color='red')
|
|
img.save(img_path)
|
|
print(f"Created test image: {img_path}")
|
|
|
|
# Process image
|
|
image = preprocess(img).unsqueeze(0).cuda()
|
|
print("Image processed and moved to GPU")
|
|
|
|
# Test classification
|
|
text_labels = ["a photo of a bee", "a photo of a flower", "a photo of a document"]
|
|
text_tokens = open_clip.tokenize(text_labels).cuda()
|
|
|
|
with torch.no_grad():
|
|
image_features = model.encode_image(image)
|
|
text_features = model.encode_text(text_tokens)
|
|
|
|
image_features /= image_features.norm(dim=-1, keepdim=True)
|
|
text_features /= text_features.norm(dim=-1, keepdim=True)
|
|
|
|
similarity = (100.0 * image_features @ text_features.T).softmax(dim=-1)
|
|
values, indices = similarity[0].topk(3)
|
|
|
|
print("Classification results:")
|
|
for i, (val, idx) in enumerate(zip(values, indices)):
|
|
label = text_labels[idx]
|
|
conf = float(val)
|
|
print(f" {i+1}. {label}: {conf:.4f}")
|
|
|
|
# Clean up
|
|
os.unlink(img_path)
|
|
print("Test completed successfully")
|
|
"""
|
|
|
|
print("Running OpenCLIP GPU test...")
|
|
result = subprocess.run([venv_python, "-c", test_script], capture_output=True, text=True, timeout=60)
|
|
|
|
print("STDOUT:")
|
|
print(result.stdout)
|
|
|
|
if result.stderr:
|
|
print("STDERR:")
|
|
print(result.stderr)
|
|
|
|
if result.returncode == 0:
|
|
print("SUCCESS: OpenCLIP GPU test passed!")
|
|
return True
|
|
else:
|
|
print("ERROR: OpenCLIP GPU test failed")
|
|
return False
|
|
|
|
if __name__ == "__main__":
|
|
print("SIMPLE OPENCLIP GPU TEST")
|
|
print("=" * 50)
|
|
|
|
success = test_openclip_gpu()
|
|
|
|
if success:
|
|
print("\nOpenCLIP GPU is working correctly!")
|
|
else:
|
|
print("\nOpenCLIP GPU test failed. Please check dependencies.") |