File size: 2,742 Bytes
a499520 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | import gradio as gr
import torch
import time
def get_gpu_info():
"""Get information about available GPUs."""
if torch.cuda.is_available():
gpu_count = torch.cuda.device_count()
gpu_info = []
for i in range(gpu_count):
name = torch.cuda.get_device_name(i)
mem_total = torch.cuda.get_device_properties(i).total_memory / (1024**3)
gpu_info.append(f"GPU {i}: {name} ({mem_total:.1f}GB)")
return "\n".join(gpu_info)
return "No GPU available"
def test_gpu(gpu_name):
"""Test the selected GPU with a simple computation."""
if not torch.cuda.is_available():
return "No GPU available on this system"
# Find the GPU index based on the name
device_idx = 0
for i in range(torch.cuda.device_count()):
if gpu_name in torch.cuda.get_device_name(i):
device_idx = i
break
device = torch.device(f"cuda:{device_idx}")
# Simple GPU computation test
start_time = time.time()
# Create a tensor on GPU
x = torch.randn(1000, 1000, device=device)
y = torch.randn(1000, 1000, device=device)
# Perform matrix multiplication
z = torch.matmul(x, y)
# Move back to CPU
result = z.cpu()
elapsed = time.time() - start_time
return f"GPU Test Successful!\n\nSelected GPU: {gpu_name}\nComputation Time: {elapsed:.4f} seconds\nResult shape: {result.shape}\nResult sum: {result.sum().item():.2f}"
def get_available_gpus():
"""Get list of available GPU names."""
if torch.cuda.is_available():
return [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())]
return ["No GPU available"]
# Get available GPUs
available_gpus = get_available_gpus()
gpu_info = get_gpu_info()
# Create Gradio interface
with gr.Blocks(title="GPU Tester") as demo:
gr.Markdown("# GPU Selector and Tester")
gr.Markdown(f"### Available GPU Information:\n```\n{gpu_info}\n```")
with gr.Row():
with gr.Column():
gpu_dropdown = gr.Dropdown(
choices=available_gpus if available_gpus else ["No GPU available"],
label="Select GPU",
value=available_gpus[0] if available_gpus else None
)
test_button = gr.Button("Test GPU", variant="primary")
with gr.Column():
output = gr.Textbox(label="Test Results", lines=10)
test_button.click(fn=test_gpu, inputs=gpu_dropdown, outputs=output)
gr.Markdown("---")
gr.Markdown("This app allows you to select a GPU and test its computation capability using PyTorch.")
# Launch the app
demo.launch(server_name="0.0.0.0", server_port=7860) |