PyTorch GPU / CUDA Aceleración
Las operaciones centrales del aprendizaje profundo son la multiplicación de matrices a gran escala y las operaciones por elementos. La CPU está diseñada para manejar lógica serial compleja, con un número de núcleos que suele ser de 8 a 64; mientras que la GPU cuenta con miles de núcleos paralelos simples, ideal para este tipo de cálculos numéricos altamente paralelos. PyTorch, a través del framework de NVIDIACUDA(Compute Unified Device Architecture) para invocar la GPU, puede aumentar la velocidad de entrenamiento decenas e incluso cientos de veces.
1. Diferencias entre CPU y GPU
La mejora de la velocidad de entrenamiento con GPU proviene principalmente de dos aspectos: primero, la ejecución paralela de una gran cantidad de cálculos idénticos; segundo, la memoria de video de alto ancho de banda permite que la transferencia de datos sea mucho más rápida que en la memoria de la CPU.
Para operaciones intensivas en cómputo, como la multiplicación de matrices, el efecto de aceleración es particularmente notable.
| Elemento de comparación | CPU | GPU（NVIDIA） |
|---|---|---|
| Número de núcleos | 8~64 núcleos grandes | Miles de núcleos pequeños (CUDA Cores) |
| Objetivo de diseño | Procesamiento serial de baja latencia | Computación paralela de alto rendimiento |
| Ancho de banda de memoria | ~50~100 GB/s | ~500~3000 GB/s |
| Velocidad de multiplicación de matrices | Referencia | 10x~100x más rápido |
| Interfaz de PyTorch | "cpu" |
"cuda" |
Si se utiliza Apple Silicon (M1/M2/M3), PyTorch, a través dempsel backend, admite aceleración GPU con Metal; su uso es casi idéntico al de CUDA, solo hay que cambiar el dispositivo a"mps"。
2. Detección del entorno CUDA
Antes de usar la GPU, es necesario confirmar si el entorno actual tiene instalado PyTorch con soporte CUDA.
También hay que comprobar si el sistema tiene algún dispositivo GPU disponible.
Ejemplo
# ¿Es compatible con CUDA?
print("CUDA disponible:", torch.cuda.is_available())
# Número de dispositivos GPU
print("Número de GPU:", torch.cuda.device_count())
# Índice de la GPU predeterminada actual
print("GPU actual:", torch.cuda.current_device())
# Nombre del modelo de GPU
print("Modelo de GPU:", torch.cuda.get_device_name(0))
# Versión de PyTorch y versión de CUDA vinculada en la compilación
print("Versión de PyTorch:", torch.__version__)
print("Versión de CUDA:", torch.version.cuda)
Ejemplo de salida:
CUDA 可用: True GPU 数量: 1 当前 GPU: 0 GPU 型号: NVIDIA GeForce RTX 4090 PyTorch 版本: 2.3.0+cu121 CUDA 版本: 12.1
2.1 Selección dinámica del dispositivo (método recomendado)
Hardcodear en el código"cuda"hará que las máquinas sin GPU den error directamente.
Ejemplo
# Método 1: escritura clásica, mejor compatibilidad
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Método 2: recomendado en PyTorch 2.0+, soporta automáticamente CUDA / MPS / CPU
device = (
"cuda" if torch.cuda.is_available()
else "mps" if torch.backends.mps.is_available()
else "cpu"
)
print(f"Dispositivo en uso: {device}")
3. Mover tensores entre dispositivos
Los tensores en PyTorch se crean en la CPU de forma predeterminada.
Para usar cómputo con GPU, es necesario mover explícitamente el tensor a la GPU, o crearlo directamente en la GPU.
3.1 Métodos básicos de movimiento
Ejemplo
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Crear tensor en CPU
cpu_tensor = torch.tensor([1.0, 2.0, 3.0])
print(cpu_tensor.device) # cpu
# Método 1: .to(device) — recomendado, mejor versatilidad
gpu_tensor = cpu_tensor.to(device)
# Método 2: .cuda() — solo para entorno CUDA
gpu_tensor = cpu_tensor.cuda()
# Método 3: especificar el dispositivo directamente al crear
gpu_tensor = torch.tensor([1.0, 2.0, 3.0], device=device)
gpu_tensor = torch.randn(3, 4, device=device)
# Volver a CPU (para imprimir, convertir a numpy, guardar, etc.)
back_to_cpu = gpu_tensor.cpu()
print(gpu_tensor.device) # cuda:0
print(back_to_cpu.device) # cpu
Al convertir un tensor de GPU a numpy, primero debe trasladarse a la CPU y, si el tensor tiene gradiente, también es necesario hacer detach primero:
Ejemplo
arr = gpu_tensor.cpu().numpy()
# Convertir tensor GPU con gradiente a numpy
arr = gpu_tensor.detach().cpu().numpy()
3.2 Restricciones de consistencia de dispositivos
Los tensores en diferentes dispositivos no pueden participar directamente en la misma operación.
De lo contrario, se lanzaRuntimeError：
Ejemplo
b = torch.randn(3) # En la CPU
# c = a + b # RuntimeError: Expected all tensors to be on the same device
# Correcto: unificar el dispositivo primero
c = a + b.to("cuda")
Ver el dispositivo del tensor:
Ejemplo
print(x.device) # cuda:0
print(x.is_cuda) # True
print(x.get_device()) # 0 (índice de GPU)
3.3 Verificación de comparación de velocidad
Ejemplo
import time
device = torch.device("cuda")
n = 5000
# Multiplicación de matrices en CPU
a_cpu = torch.randn(n, n)
b_cpu = torch.randn(n, n)
start = time.time()
c_cpu = torch.matmul(a_cpu, b_cpu)
print(f"Tiempo CPU: {time.time() - start:.3f}s")
# Multiplicación de matrices en GPU
a_gpu = a_cpu.to(device)
b_gpu = b_cpu.to(device)
torch.cuda.synchronize() # Asegurarse de que los datos se hayan transferido antes de comenzar a cronometrar
start = time.time()
c_gpu = torch.matmul(a_gpu, b_gpu)
torch.cuda.synchronize() # Esperar a que la GPU termine de ejecutar antes de detener el cronómetro
print(f"Tiempo GPU: {time.time() - start:.3f}s")
Ejemplo de salida:
CPU 耗时: 1.847s GPU 耗时: 0.021s
El cálculo con GPU esasincrónico—después de que la llamada de Python regresa, las operaciones en la GPU pueden no haber terminado. Al medir el tiempo, hay que llamar a
torch.cuda.synchronize()para esperar a que la GPU realmente termine; de lo contrario, el resultado de la medición no será preciso.
4. Mover el modelo a la GPU
Todos los parámetros del modelo (weight、bias) también son, en esencia, tensores.
Estos parámetros también deben moverse a la GPU para poder ejecutar la propagación hacia adelante y la propagación hacia atrás en la GPU.
Invocar en todo el modelo.to(device)es suficiente; PyTorch recorrerá automáticamente y moverá todos los parámetros internos.
Ejemplo
import torch.nn as nn
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class SimpleNet(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(784, 256),
nn.ReLU(),
nn.Linear(256, 128),
nn.ReLU(),
nn.Linear(128, 10),
)
def forward(self, x):
return self.net(x)
# Mover el modelo a la GPU, solo se necesita llamar una vez
model = SimpleNet().to(device)
# Verificar si todos los parámetros están en la GPU
for name, param in model.named_parameters():
print(f"{name}: {param.device}")
# net.0.weight: cuda:0
# net.0.bias: cuda:0
# ...
# Los datos de entrada también deben estar en el mismo dispositivo
x = torch.randn(32, 784).to(device)
output = model(x)
print(output.shape) # torch.Size([32, 10])
Si el modelo está en la GPU pero los datos de entrada permanecen en la CPU, la propagación hacia adelante arrojará un error. Asegúrese de que, después de que DataLoader lea cada lote,
inputsylabelsse llame a .to(device) en todos ellos..to(device)。
5. Flujo de entrenamiento completo
A continuación se presenta una plantilla estándar de entrenamiento con GPU que incluye carga de datos, entrenamiento del modelo y evaluación/validación.
Ejemplo
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
# 1. Configuración del dispositivo
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Dispositivo de entrenamiento: {device}")
# 2. Carga de datos
# pin_memory=True: bloquea los datos en la memoria, acelera la transferencia CPU -> GPU
# num_workers: precarga con múltiples procesos, reduce el tiempo de espera de datos
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,)),
])
train_dataset = datasets.MNIST(root="./data", train=True,
download=True, transform=transform)
test_dataset = datasets.MNIST(root="./data", train=False,
download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True,
num_workers=4, pin_memory=True)
test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False,
num_workers=4, pin_memory=True)
# 3. Definir el modelo
class CNN(nn.Module):
def __init__(self):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(64 * 7 * 7, 256), nn.ReLU(), nn.Dropout(0.5),
nn.Linear(256, 10),
)
def forward(self, x):
return self.classifier(self.features(x))
model = CNN().to(device) # Mover el modelo a la GPU
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# 4. Función de entrenamiento
def train_epoch(model, loader, optimizer, criterion):
model.train()
total_loss, correct = 0.0, 0
for inputs, labels in loader:
# non_blocking=True: transferencia asíncrona, la CPU puede seguir preparando el siguiente lote de datos
inputs = inputs.to(device, non_blocking=True)
labels = labels.to(device, non_blocking=True)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
total_loss += loss.item() * inputs.size(0)
correct += (outputs.argmax(1) == labels).sum().item()
n = len(loader.dataset)
return total_loss / n, correct / n
# 5. Función de validación
def eval_epoch(model, loader, criterion):
model.eval()
total_loss, correct = 0.0, 0
with torch.no_grad():
for inputs, labels in loader:
inputs = inputs.to(device, non_blocking=True)
labels = labels.to(device, non_blocking=True)
outputs = model(inputs)
loss = criterion(outputs, labels)
total_loss += loss.item() * inputs.size(0)
correct += (outputs.argmax(1) == labels).sum().item()
n = len(loader.dataset)
return total_loss / n, correct / n
# 6. Bucle principal de entrenamiento
for epoch in range(1, 11):
train_loss, train_acc = train_epoch(model, train_loader, optimizer, criterion)
val_loss, val_acc = eval_epoch(model, test_loader, criterion)
print(f"Epoch {epoch:02d} | "
f"Train Loss: {train_loss:.4f}, Acc: {train_acc:.4f} | "
f"Val Loss: {val_loss:.4f}, Acc: {val_acc:.4f}")
Puntos clave:
deviceLas variables se declaran en la parte superior del código, se referencian globalmente de manera uniforme y no se codifican valores fijos."cuda"- Modelo
.to(device)solo necesita llamarse una vez - de cada batch
inputsylabelstodos deben.to(device) - Configuración de DataLoader
pin_memory=Trueynon_blocking=TrueAcelerar la transferencia de datos - Usar en la fase de validación
torch.no_grad()Desactivar gradientes para ahorrar memoria de GPU
6. Entrenamiento con múltiples GPU
Cuando la memoria de una sola GPU es insuficiente, se puede considerar el uso de entrenamiento paralelo con múltiples GPU.
Además, las múltiples GPU pueden mejorar aún más la velocidad de entrenamiento. PyTorch ofrece dos métodos principales.
6.1 DataParallel
Esta es la forma más sencilla de múltiples GPU.
Usa un solo proceso, divide cada batch uniformemente entre las GPU, la propagación hacia adelante se ejecuta en paralelo y los gradientes se agregan y actualizan en la GPU principal. Es fácil de usar, pero la GPU principal tiene una carga pesada y la utilización de múltiples GPU es desigual; es adecuado para comenzar rápidamente.
Ejemplo
import torch.nn as nn
model = CNN()
if torch.cuda.device_count() > 1:
print(f"Usando {torch.cuda.device_count()} GPU")
model = nn.DataParallel(model)
# También se puede especificar qué GPU usar
# model = nn.DataParallel(model, device_ids=[0, 1])
model = model.to("cuda")
# El código de entrenamiento posterior es exactamente igual que con una sola GPU
# DataParallel divide automáticamente el batch entre las GPU y agrega los resultados
Si necesita acceder a los atributos del modelo original (como métodos personalizados), debe.moduleacceder a:
Ejemplo
print(model.module.classifier)
# Al guardar el modelo, se recomienda guardar model.module para facilitar la carga en una sola GPU
torch.save(model.module.state_dict(), "model.pth")
6.2 DistributedDataParallel
Este es el método de múltiples procesos.
Cada proceso está vinculado a una GPU, cada tarjeta tiene el modelo completo y los gradientes se sincronizan mediante All-Reduce. Es la opción recomendada para entornos de producción y entrenamiento a gran escala, con una eficiencia mucho mayor que DataParallel.
Ejemplo
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
def main(rank, world_size):
# Inicializar el grupo de procesos; el backend nccl está optimizado para GPU NVIDIA
dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
# Cada proceso se vincula a su GPU correspondiente
torch.cuda.set_device(rank)
device = torch.device(f"cuda:{rank}")
# Mover el modelo a la GPU correspondiente y luego envolver con DDP
model = CNN().to(device)
model = DDP(model, device_ids=[rank])
# DataLoader usa DistributedSampler para garantizar que los datos no se superpongan entre tarjetas
sampler = DistributedSampler(train_dataset,
num_replicas=world_size, rank=rank)
loader = DataLoader(train_dataset, batch_size=64,
sampler=sampler, pin_memory=True)
# La lógica de entrenamiento es exactamente igual que con una sola GPU
# ...
dist.destroy_process_group()
# Forma de inicio (se recomienda usar torchrun):
# torchrun --nproc_per_node=4 train_ddp.py
| Elemento de comparación | DataParallel | DistributedDataParallel |
|---|---|---|
| Número de procesos | Un solo proceso | Múltiples procesos (uno por tarjeta) |
| Backend de comunicación | Limitación de GIL de Python | NCCL (eficiente) |
| Carga de la GPU principal | Alta (agregar gradientes) | Equilibrada (All-Reduce) |
| Cambios en el código | Mínimos | Moderados |
| Escenarios de aplicación | Experimentos rápidos | Entrenamiento en producción |
7. Entrenamiento de precisión mixta AMP
El entrenamiento predeterminado usa precisión float32 (FP32).
Precisión mixta automática (Automatic Mixed Precision, AMP)Hace que algunos cálculos usen float16 (FP16) o bfloat16 (BF16), logrando beneficios significativos casi sin perder precisión:
- El uso de memoria de GPU se reduce aproximadamente un 50%
- La velocidad de entrenamiento aumenta de 2x a 3x (aceleración por hardware Tensor Core)
- Los cambios en el código son mínimos, solo se necesitan tres modificaciones
Ejemplo
model = CNN().to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler() # Escalador de gradientes FP16, evita que los gradientes se desborden por debajo a cero
for epoch in range(num_epochs):
model.train()
for inputs, labels in train_loader:
inputs = inputs.to(device, non_blocking=True)
labels = labels.to(device, non_blocking=True)
optimizer.zero_grad()
# Modificación 1: dentro de la región autocast se selecciona automáticamente FP16/FP32
with autocast(device_type="cuda"):
outputs = model(inputs)
loss = criterion(outputs, labels)
# Modificación 2: usar scaler para escalar la pérdida antes de la propagación hacia atrás
scaler.scale(loss).backward()
# Modificación 3: scaler actualiza los parámetros y maneja internamente el escalado de gradientes automáticamente
scaler.step(optimizer)
scaler.update()
7.1 Elección entre FP16 y BF16
| Formato | Bits de precisión | Bits de rango numérico | Hardware compatible | Características |
|---|---|---|---|---|
| float16 | 10 bits | 5 bits | Serie RTX 20/30/40 | Requiere GradScaler para evitar desbordamiento |
| bfloat16 | 7 bits | 8 bits | A100 / H100 / RTX 4090 | El rango es el mismo que FP32, entrenamiento más estable |
Si la GPU admite BF16, priorícelo, sin necesidad de GradScaler:
Ejemplo
with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
8. Consejos de optimización del rendimiento
Esta sección presenta varias técnicas para optimizar la memoria de GPU y la velocidad de entrenamiento.
8.1 Gestión de la memoria de video
Ejemplo
print(f"Asignado: {torch.cuda.memory_allocated() / 1024**2:.1f} MB")
print(f"En caché: {torch.cuda.memory_reserved() / 1024**2:.1f} MB")
# Imprimir informe detallado de memoria de GPU
print(torch.cuda.memory_summary())
# Liberar el grupo de caché (no libera la memoria de GPU asignada)
torch.cuda.empty_cache()
# Usar inference_mode durante la inferencia (más rápido que no_grad, desactiva completamente el motor de gradientes)
with torch.inference_mode():
output = model(x)
# Checkpoint de gradientes: intercambiar memoria de GPU por recálculo (común en entrenamiento de modelos grandes)
from torch.utils.checkpoint import checkpoint
output = checkpoint(model_block, x)
8.2 Optimización de DataLoader
Ejemplo
dataset,
batch_size=256,
num_workers=4, # Precarga multiproceso, se recomienda establecer al 50% de los núcleos de CPU
pin_memory=True, # Memoria fija (pinned memory), acelera la transferencia CPU -> GPU
persistent_workers=True, # Procesos persistentes, evita recrear procesos en cada epoch
prefetch_factor=2, # Número de batches precargados por cada worker
)
8.3 Compilación de modelos con torch.compile (PyTorch 2.0+)
Con una línea de código, compila y optimiza el grafo computacional, la velocidad aumenta entre 30% y 200% (dependiendo de la estructura del modelo):
Ejemplo
model = torch.compile(model)
# Compensaciones entre diferentes modos de compilación
model = torch.compile(model, mode="default") # Equilibrado, adecuado para la mayoría de los escenarios
model = torch.compile(model, mode="reduce-overhead") # Reduce la sobrecarga de programación
model = torch.compile(model, mode="max-autotune") # Optimización máxima, tiempo de compilación más largo
8.4 Acumulación de gradientes (simular batch grande)
Cuando la memoria de GPU es insuficiente, se puede simular un tamaño de batch mayor mediante la acumulación de gradientes, sin necesidad de aumentar realmente el uso de memoria:
Ejemplo
for i, (inputs, labels) in enumerate(train_loader):
inputs = inputs.to(device)
labels = labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels) / accumulation_steps # Dividir la pérdida de manera uniforme
loss.backward() # Acumular gradientes, no poner a cero
if (i + 1) % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad() # Solo poner a cero después de actualizar los parámetros
8.5 Estrategias ante memoria de video insuficiente
| Estrategia | Método | Efecto |
|---|---|---|
| Reducir el batch size | 256 -> 64 | Reduce la memoria de GPU linealmente |
| Precisión mixta | AMP + FP16/BF16 | Reduce la memoria de GPU aproximadamente un 50% |
| Acumulación de gradientes | Actualizar cada N pasos | Simula batch grande, no aumenta la memoria de GPU |
| Checkpoint de gradientes | torch.utils.checkpoint |
Reduce significativamente la memoria de GPU, pero reduce la velocidad de entrenamiento |
| Congelar algunos parámetros | Congelar el backbone en transfer learning | Reduce el uso de memoria de GPU en la retropropagación |
9. Errores comunes y solución de problemas
Esta sección resume los errores más comunes en el entrenamiento con GPU y sus soluciones.
9.1 RuntimeError: Expected all tensors to be on the same device
Causa: los tensores que participan en la operación están distribuidos en diferentes dispositivos (uno en CPU, otro en GPU).
Ejemplo
print(inputs.device, labels.device, next(model.parameters()).device)
# Solución: asegurarse de que cada batch haya ejecutado .to(device)
inputs = inputs.to(device)
labels = labels.to(device)
9.2 RuntimeError: CUDA out of memory
Causa: memoria de GPU insuficiente, comúnmente por batch size demasiado grande, modelo demasiado grande, o acumulación inesperada del grafo computacional.
Ejemplo
# Escritura incorrecta
total_loss += loss # loss es un tensor y mantiene todo el grafo computacional
# Escritura correcta
total_loss += loss.item() # .item() extrae un escalar de Python y libera el grafo computacional
# Otros pasos de solución de problemas:
# 1. Reducir batch_size
# 2. Verificar que el bucle de validación use torch.no_grad()
# 3. Llamar a torch.cuda.empty_cache() para limpiar la caché
# 4. Usar torch.cuda.memory_summary() para localizar el origen del uso de memoria de video
9.3 Can't call numpy() on Tensor that requires grad
Causa: un tensor con gradiente no se puede convertir directamente a un array de numpy.
Ejemplo
arr = gpu_tensor.numpy()
# Escritura correcta
arr = gpu_tensor.detach().cpu().numpy()
9.4 CUDA error: device-side assert triggered
Causa: normalmente el valor de la etiqueta está fuera de rango (por ejemplo, clasificación de 10 clases pero el valor de label es igual a 10), o el índice del array está fuera de rango. El mensaje de error se genera en la GPU y la ubicación mostrada por defecto no es precisa.
Ejemplo
import os
os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
9.5 La velocidad de entrenamiento no mejora (baja utilización de GPU)
Causa: la carga de datos se convierte en el cuello de botella, y la GPU pasa la mayor parte del tiempo esperando a que la CPU prepare los datos.
Pasos de solución de problemas:
- Ejecutar
nvidia-smiowatch -n 1 nvidia-smiObservar la utilización de la GPU- Una utilización continuamente inferior al 80% indica un cuello de botella en los datos
- Aumentar num_workers del DataLoader
- Activar pin_memory=True
- Mover el preprocesamiento de datos a la GPU (torchvision.transforms admite operaciones en GPU)
- Precargar conjuntos de datos de archivos pequeños en memoria
9.6 Usar PyTorch Profiler para localizar con precisión los cuellos de botella
Ejemplo
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
record_shapes=True,
) as prof:
for i, (inputs, labels) in enumerate(train_loader):
if i >= 10:
break
inputs = inputs.to(device)
labels = labels.to(device)
loss = criterion(model(inputs), labels)
loss.backward()
optimizer.step()
optimizer.zero_grad()
# Ordenar por tiempo de CUDA e imprimir los primeros 15 elementos
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=15))
10. Referencia rápida de API
La siguiente es una tabla de referencia rápida de las API comunes relacionadas con GPU en PyTorch.
10.1 Gestión de dispositivos
| Operación | Código |
|---|---|
| Detectar disponibilidad de CUDA | torch.cuda.is_available() |
| Seleccionar dispositivo | torch.device("cuda" if ... else "cpu") |
| Número de GPU | torch.cuda.device_count() |
| Nombre de la GPU | torch.cuda.get_device_name(0) |
| Establecer GPU actual | torch.cuda.set_device(0) |
| Esperar a que la GPU termine | torch.cuda.synchronize() |
10.2 Operaciones con tensores
| Operación | Código |
|---|---|
| Mover tensor a GPU | tensor.to(device)otensor.cuda() |
| Mover tensor a CPU | tensor.cpu() |
| Ver dispositivo actual | tensor.device |
| Si está en GPU | tensor.is_cuda |
| Convertir tensor de GPU a numpy | tensor.detach().cpu().numpy() |
| Transferencia asíncrona | tensor.to(device, non_blocking=True) |
10.3 Operaciones con modelos
| Operación | Código |
|---|---|
| Mover modelo a GPU | model.to(device) |
| Activar modo de entrenamiento | model.train() |
| Activar modo de inferencia | model.eval() |
| Aceleración de compilación | torch.compile(model) |
| Desactivar gradientes | with torch.no_grad(): |
| Modo de inferencia | with torch.inference_mode(): |
10.4 Gestión de la memoria de video
| Operación | Código |
|---|---|
| Memoria de video asignada | torch.cuda.memory_allocated() |
| Memoria de video en caché | torch.cuda.memory_reserved() |
| Informe de memoria de video | torch.cuda.memory_summary() |
| Limpiar caché | torch.cuda.empty_cache() |
10.5 Notas rápidas sobre principios fundamentales
1. 用 device 变量统一管理，不要硬编码 "cuda" 2. 模型和数据必须在同一设备上，每个 batch 都要 .to(device) 3. 验证和推理时一定使用 torch.no_grad() o torch.inference_mode() 4. 生产训练推荐开启 AMP 混合精度，几乎免费获得 2x 加速 5. DataLoader 设置 pin_memory=True 和 num_workers >= 4 减少数据瓶颈 6. PyTorch 2.0+ 可以用 torch.compile(model) 一行提速 30%+ 7. GPU 计算是异步的，精确计时需要调用 torch.cuda.synchronize()Otras extensiones