Metal
A continuación se muestra un script para probar el rendimiento de Metal (GPU) utilizando PyTorch.
import torch
import argparse
import time
parser = argparse.ArgumentParser(description="Probar torch con MPS o CPU.")
parser.add_argument("--device", type=str, default="mps", choices=["mps", "cpu"], help="Dispositivo a usar (mps o cpu)")
args = parser.parse_args()
if args.device == "mps":
if torch.backends.mps.is_available():
print("Metal está disponible")
device = torch.device("mps")
else:
print("Metal no está disponible, usando CPU en su lugar")
device = torch.device("cpu")
elif args.device == "cpu":
device = torch.device("cpu")
print("Usando CPU")
else:
print("Dispositivo no válido especificado, usando CPU en su lugar")
device = torch.device("cpu")
# Crear un tensor y moverlo al dispositivo especificado
x = torch.randn(5000, 5000, device=device)
y = torch.randn(5000, 5000, device=device)
# Realizar un cálculo más complejo
start_time = time.time()
result = torch.matmul(x, y)
for _ in range(10):
result = torch.matmul(result, y)
end_time = time.time()
# Imprimir el resultado
print(result)
print(f"Tiempo transcurrido: {end_time - start_time:.4f} segundos")
Los resultados demuestran que MPS es significativamente más rápido que CPU. El tiempo de ejecución en MPS es solo aproximadamente el 0.2% del tiempo de CPU.
% python scripts/test_metal.py --device cpu
Tiempo transcurrido: 2.8784 segundos
% python scripts/test_metal.py --device mps
Tiempo transcurrido: 0.0061 segundos