mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-09-16 05:10:42 -07:00
78 lines
2.5 KiB
Python
78 lines
2.5 KiB
Python
"""Test CUDA detection in voicebox backend"""
|
|
import sys
|
|
import torch
|
|
|
|
print("=" * 60)
|
|
print("PyTorch CUDA Detection Test")
|
|
print("=" * 60)
|
|
|
|
# Basic torch info
|
|
print(f"\nPyTorch version: {torch.__version__}")
|
|
print(f"CUDA available: {torch.cuda.is_available()}")
|
|
|
|
if torch.cuda.is_available():
|
|
print(f"CUDA version: {torch.version.cuda}")
|
|
print(f"GPU count: {torch.cuda.device_count()}")
|
|
print(f"Current GPU: {torch.cuda.current_device()}")
|
|
print(f"GPU name: {torch.cuda.get_device_name(0)}")
|
|
print(f"GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")
|
|
else:
|
|
print("\nNo CUDA available - would run on CPU")
|
|
|
|
# Test backend device selection
|
|
print("\n" + "=" * 60)
|
|
print("Backend Device Selection")
|
|
print("=" * 60)
|
|
|
|
# Simulate the _get_device method from pytorch_backend.py
|
|
def _get_device() -> str:
|
|
"""Get the best available device."""
|
|
if torch.cuda.is_available():
|
|
return "cuda"
|
|
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
|
|
# MPS can have issues, use CPU for stability
|
|
return "cpu"
|
|
return "cpu"
|
|
|
|
selected_device = _get_device()
|
|
print(f"\nSelected device: {selected_device}")
|
|
print(f"Would use dtype: {'torch.bfloat16' if selected_device != 'cpu' else 'torch.float32'}")
|
|
|
|
# Test actual tensor creation on device
|
|
print("\n" + "=" * 60)
|
|
print("Testing Tensor Creation on Device")
|
|
print("=" * 60)
|
|
|
|
try:
|
|
test_tensor = torch.randn(1000, 1000).to(selected_device)
|
|
print(f"\n[OK] Successfully created tensor on {selected_device}")
|
|
print(f" Tensor device: {test_tensor.device}")
|
|
print(f" Tensor dtype: {test_tensor.dtype}")
|
|
|
|
# Test computation
|
|
result = test_tensor @ test_tensor.T
|
|
print(f"[OK] Successfully performed computation on {selected_device}")
|
|
|
|
if selected_device == "cuda":
|
|
print(f"\nCUDA memory allocated: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
|
|
print(f"CUDA memory reserved: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")
|
|
|
|
except Exception as e:
|
|
print(f"\n[ERROR] {e}")
|
|
|
|
print("\n" + "=" * 60)
|
|
print("Summary")
|
|
print("=" * 60)
|
|
|
|
if selected_device == "cuda":
|
|
print("\n[SUCCESS] CUDA IS WORKING!")
|
|
print(" The backend will use your NVIDIA GPU for inference")
|
|
print(f" GPU: {torch.cuda.get_device_name(0)}")
|
|
print(f" This will be significantly faster than CPU")
|
|
else:
|
|
print("\n[FAIL] CUDA is not available")
|
|
print(" The backend will use CPU for inference")
|
|
print(" This will be slower than GPU")
|
|
|
|
print("\n" + "=" * 60)
|