Noveno en nuestra serie en Perfil de rendimiento y optimización en Pytorch Dirigido a enfatizar el papel crítico del análisis y la optimización del rendimiento en el desarrollo del aprendizaje automático. A lo largo de la serie, hemos revisado una amplia variedad de herramientas y técnicas prácticas para analizar y aumentar el rendimiento del tiempo de ejecución de los modelos AI/ML basados en Pytorch. Nuestro objetivo ha sido doble:
- Para enfatizar la importancia de la evaluación de rutina y la optimización de las cargas de trabajo AI/ML.
- Para demostrar la accesibilidad de una amplia variedad de herramientas y técnicas para analizar y optimizar el rendimiento de tiempo de ejecución de AI/ML. No necesita ser un experto en CUDA para mejorar significativamente el rendimiento de su modelo y reducir los costos de cómputo.
En esta publicación, exploraremos el uso de transmisiones CUDA, una característica poderosa del modelo de programación CUDA de NVIDIA que ofrece un método sofisticado para superponer las operaciones de GPU y ejecutarlas simultáneamente. Aunque normalmente asociamos nuestra carga de trabajo de capacitación de modelos AI/ML con un solo gráfico de computación monolítico (también conocido como “Unbreakable”) GRAMO Ejecutando en la GPU, hay algunos escenarios en los que el gráfico se puede descomponer en dos subgrafías distintas G1 y G2dónde G = G2*G1. En tales casos, las secuencias de CUDA permiten “tuberías” el gráfico de cálculo, es decir, programar nuestro paso de entrenamiento para ejecutar G1 (En la entrada por lotes N+1) en paralelo a G2 (en el enésimo de salida de G1). Esta técnica es especialmente útil cuando:
- Ninguno de los subgrafías utiliza completamente la GPU cuando se ejecuta solo, y
- Los dos subgraphs son de costo computacional similar (es decir, ninguno de los dos domina el tiempo de ejecución).
Exploraremos dos escenarios comunes en los que es factible “tuberías”:
- Entrenamiento de modelo parcial o finecir:
Es común congelar un modelo previamente capacitado columna vertebral (por ejemplo, extractor o codificador de características) y entrenar solo un modelo cabeza (por ejemplo, decodificador). Desde el congelado columna vertebral no confía en los gradientes del cabezalos dos se pueden ejecutar simultáneamente. - Descargar el preprocesamiento de datos a la GPU:
Un método común para abordar los cuellos de botella en la tubería de entrada (también conocido como inanición de GPU), el preprocesamiento de datos se puede mover a la GPU. Si bien la preparación de las operaciones de preprocesamiento al gráfico del modelo mejora el rendimiento, se pueden lograr ganancias adicionales ejecutando el preprocesamiento en una corriente CUDA separada en paralelo con la ejecución del modelo, lo que el preprocesamiento no es trivial en comparación con el cómputo modelo.
Para facilitar nuestra discusión, definiremos dos guiones de entrenamiento de juguetes y mediremos el rendimiento de la capacitación en diferentes escenarios. Los experimentos se ejecutaron en un Amazon EC2 G5.2xLarge instancia (que contiene una GPU NVIDIA A10G y 8 VCPU) que ejecuta un Pytorch (2.6) Aprendizaje profundo ami (Dlami).
Tenga en cuenta: los fragmentos de código que compartimos son solo para fines de demostración: por favor no confíe en su corrección u optimización. El impacto del uso de transmisiones CUDA variará según la arquitectura del modelo y la configuración del sistema. Le recomendamos que realice su propio perfil y experimentación antes de integrar las transmisiones CUDA (o cualquier otra técnica de herramienta a la que nos referimos) en su flujo de trabajo.
Parte 1: Libelación de un modelo de codificador codificador
El primer caso de uso que exploramos implica un modelo de segmentación de imágenes basado en CNN que consiste en un codificador fijo (previamente entrenado) y un decodificador capacitable. En este escenario, dado que los pesos del codificador se congelan y no se ven afectados por la backpropagation, el codificador puede ejecutarse independientemente de la capacitación del decodificador. En esta sección, evaluamos el impacto de la canalización del proceso de capacitación utilizando transmisiones CUDA.
Un experimento de entrenamiento de segmentación de imágenes de juguete
Comenzamos definiendo un codificador de imagen basado en CNN simple junto con su decodificador correspondiente.
undefinedA continuación, construimos un conjunto de datos sintético de imágenes aleatorias y mapas de segmentación.
from torch.utils.data import DataLoader
from torchvision.datasets.vision import VisionDataset
# A dataset with random images and per-pixel labels
class FakeDataset(VisionDataset):
def __init__(self):
super().__init__(root=None)
self.size = 1000000
def __getitem__(self, index):
# create a random image
img = torch.randint(0, 256, (3, img_size, img_size),
dtype=torch.uint8)
# create a random label map
target = torch.randint(0, num_classes, (img_size, img_size))
return img, target
def __len__(self):
return self.size
train_set = FakeDataset()
train_loader = DataLoader(
dataset=train_set,
batch_size=8,
num_workers=8
)Finalmente, definimos la función de pérdida, el optimizador y el bucle de entrenamiento. Tenga en cuenta que congelamos los pesos del codificador y entramos solo al decodificador.
import time
device = torch.device("cuda")
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(decoder.parameters())
# Freeze the encoder weights
encoder.requires_grad_(False)
encoder.eval().to(device)
decoder.train().to(device)
warmup = 10
active_batches = 100
total_iters = warmup + active_batches
for idx, data in enumerate(train_loader):
inputs = data[0].to(device=device, non_blocking=True).float()
labels = data[1].to(device=device, non_blocking=True)
optimizer.zero_grad()
with torch.no_grad():
features = encoder(inputs)
output = decoder(features)
loss = criterion(output, labels)
loss.backward()
optimizer.step()
if idx == warmup:
# sync the GPU and start the timer
torch.cuda.synchronize()
t0 = time.perf_counter()
if idx == total_iters:
break
# wait for the GPU to finnish and then stop the timer
torch.cuda.synchronize()
total_time = time.perf_counter() - t0
print(f'throughput: {active_batches / total_time}')Nuestro script de entrenamiento de línea de base logra un rendimiento promedio de 83 pasos por segundo, con una utilización promedio de GPU del 85%.
Canalización de la ejecución del modelo con corrientes CUDA
En la versión revisada del bucle de entrenamiento que se muestra a continuación, presentamos dos transmisiones CUDA: una para ejecutar el codificador y otro para capacitar al decodificador. En cada iteración, realizamos dos operaciones simultáneamente:
- Entrena el decodificador utilizando las características de la imagen y las etiquetas desde el lote norte.
- Ejecutar el codificador en el lote de entrada N+1 para generar sus características de imagen.
encoder_stream = torch.cuda.Stream()
decoder_stream = torch.cuda.Stream()
# initialize the features to None
features = None
for idx, data in enumerate(train_loader):
inputs = data[0].to(device, non_blocking=True).float()
labels_next = data[1].to(device, non_blocking=True)
if features is not None:
with torch.cuda.stream(decoder_stream):
decoder_stream.wait_stream(encoder_stream)
optimizer.zero_grad()
output = decoder(features)
loss = criterion(output, labels)
loss.backward()
optimizer.step()
with torch.cuda.stream(encoder_stream):
with torch.no_grad():
features = encoder(inputs)
# Record that features was produced on s1_backbone
features.record_stream(encoder_stream)
labels = labels_next
if idx == warmup:
# sync the GPU and start the timer
torch.cuda.synchronize()
t0 = time.perf_counter()
if idx == total_iters:
break
# wait for the GPU to finish and then stop the timer
torch.cuda.synchronize()
total_time = time.perf_counter() - t0
print(f'throughput: {active_batches / total_time}')Esta modificación produce un rendimiento promedio de 91 pasos por segundo, lo que representa una aceleración del 9.6%. Esta es una mejora significativa, especialmente teniendo en cuenta que nuestra línea de base ya tenía una alta utilización de GPU (85%).
Sensibilidad de las propiedades de la carga de trabajo
La efectividad del canalización con las corrientes CUDA depende en gran medida de los detalles de la carga de trabajo de capacitación y el entorno de tiempo de ejecución. Si el codificador es significativamente mayor que el decodificador (o viceversa), la tuberías puede ofrecer poco beneficio o incluso obstaculizar el rendimiento. Por el contrario, cuando la GPU está subutilizada, la canalización tiende a producir ganancias más sustanciales.
Para ilustrar esta dependencia, volvemos al experimento con diferentes tamaños de lotes. Los resultados se resumen a continuación:
A medida que aumenta el tamaño del lote, disminuye el beneficio de la tubería. Esto es probable porque los tamaños de lotes más grandes conducen naturalmente a una mayor utilización de la GPU (y más eficiente), dejando menos margen de mejora a través de la ejecución concurrente.
Parte 2: descarga de aumentos en la GPU
En esta sección, aplicaremos el uso de transmisiones CUDA a la aceleración del aumento de datos. En publicaciones de blog anteriores (por ejemplo, aquí y aquí), hemos estudiado el problema de los cuellos de botella en la tubería de entrada de datos desde diferentes perspectivas y revisamos varias técnicas para diagnosticarlos y abordarlos. Una causa común de estos cuellos de botella es el agotamiento de los recursos de la CPU, donde la CPU no puede satisfacer las demandas computacionales de la tubería de preprocesamiento. El resultado es la inanición de la GPU, un escenario en el que la costosa GPU está inactiva, esperando que lleguen datos.
Una solución efectiva es descargar un preprocesamiento de datos pesados en la GPU. Demostraremos esta técnica y lo llevaremos un paso más al ejecutar los aumentos en una transmisión CUDA dedicada, lo que permite la ejecución concurrente con la capacitación del modelo.
Un experimento de entrenamiento de clasificación de imagen de juguete
Comenzamos definiendo un modelo simple de clasificación de imágenes basado en CNN:
import torch
import torch.nn as nn
import torch
import torch.nn as nn
img_size = 256
num_classes = 10
model = nn.Sequential(
# Start with 256x256 image
nn.Conv2d(3, 16, kernel_size=1),
nn.ReLU(inplace=True),
nn.Conv2d(16, 32, kernel_size=2, stride=2), # 2x downsample
nn.ReLU(inplace=True),
nn.Conv2d(32, 64, kernel_size=2, stride=2), # 4x downsample
nn.ReLU(inplace=True),
nn.Conv2d(64, 128, kernel_size=2, stride=2), # 8x downsample
nn.ReLU(inplace=True),
nn.Conv2d(128, 256, kernel_size=2, stride=2), # 16x downsample
nn.ReLU(inplace=True),
nn.Conv2d(256, 512, kernel_size=2, stride=2), # 32x downsample
nn.ReLU(inplace=True),
nn.Conv2d(512, 1024, kernel_size=2, stride=2), # 64x downsample
nn.ReLU(inplace=True),
nn.Conv2d(1024, 2048, kernel_size=2, stride=2), # 128X downsample
nn.ReLU(inplace=True),
nn.Conv2d(2048, 4096, kernel_size=2, stride=2), # 256X
nn.Flatten(),
nn.Linear(4096, num_classes)
)A continuación, creamos un conjunto de datos sintético con una tubería de aumento diseñada intencionalmente para causar un cuello de botella de rendimiento severo:
import random
from torch.utils.data import DataLoader
import torchvision.transforms.v2 as T
from torchvision.datasets.vision import VisionDataset
import torchvision.transforms.v2.functional as F
import torchvision.ops as ops
# A dataset with random images and labels
class FakeDataset(VisionDataset):
def __init__(self, transform = None):
super().__init__(root=None, transform=transform)
self.size = 1000000
def __getitem__(self, index):
# create a random image
img = torch.randint(0, 256, (3, img_size, img_size),
dtype=torch.uint8)
# create a random label
target = torch.randint(0, num_classes, (1, ))
if self.transform:
# Apply tranformations
img = self.transform(img)
return img, target
def __len__(self):
return self.size
augmentations = T.Compose([
T.ToDtype(torch.float32),
T.RandomCrop(img_size//2),
T.Resize(img_size),
T.RandomRotation(degrees=45.0),
T.GaussianBlur(kernel_size=7),
T.Normalize(mean=[0, 0, 0], std=[1, 1, 1])
])
train_set = FakeDataset(transform=augmentations)
train_loader = DataLoader(
dataset=train_set,
batch_size=32,
num_workers=8
)Finalmente, definimos la función de pérdida, el optimizador y el bucle de entrenamiento:
import time
device = torch.device("cuda")
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters())
model.train().to(device)
warmup = 10
active_batches = 100
total_iters = warmup + active_batches
for idx, data in enumerate(train_loader):
inputs = data[0].to(device=device, non_blocking=True)
labels = data[1].to(device=device, non_blocking=True).squeeze()
optimizer.zero_grad()
output = model(inputs)
loss = criterion(output, labels)
loss.backward()
optimizer.step()
if idx == warmup:
# sync the GPU and start the timer
torch.cuda.synchronize()
t0 = time.perf_counter()
if idx == total_iters:
break
# wait for the GPU to finnish and then stop the timer
torch.cuda.synchronize()
total_time = time.perf_counter() - t0
print(f'throughput: {active_batches / total_time}')Ejecutar este script de línea de base da como resultado un rendimiento promedio de 20.41 pasos por segundo y una utilización de GPU de solo 42%. Los pesados aumentos de datos están ahogando la CPU que conduce a la inanición de la GPU. Ver nuestro publicación anterior Para obtener más información sobre la detección de cuellos de botella en la tubería de entrada de datos.
Descarga de aumentos de datos a la GPU
Para abordar el cuello de botella de rendimiento en la tubería de entrada de datos, movemos los aumentos a la GPU.
El primer paso es definir Transformaciones de datos personalizadas que aplican rotaciones aleatorias y cultivos por muestra en un lote. Esto es importante porque el incorporado vía antorcha Las transformaciones aplican el mismo aumento en todo el lote: perder la aleatoriedad por muestra que se observa en la CPU.
Implementamos el Batchrandomcrop transformar usando el ROI_ALIGN operador.
class BatchRandomCrop(T.Transform):
def __init__(self, output_size):
super().__init__()
self.output_size = output_size
def transform(self, img: torch.Tensor, params: dict):
batch_size, _, original_height, original_width = img.shape
device = img.device
max_top = original_height - self.output_size
max_left = original_width - self.output_size
# Generate random top and left coords for each image in the batch
random_top = torch.randint(0, max_top + 1, (batch_size,),
device=device, dtype=torch.float32)
random_left = torch.randint(0, max_left + 1, (batch_size,),
device=device, dtype=torch.float32)
image_indices = torch.arange(batch_size, device=device,
dtype=torch.float32)
boxes = torch.stack([
image_indices,
random_left,
random_top,
random_left + self.output_size,
random_top + self.output_size
], dim=1)
cropped_batch = ops.roi_align(
img,
boxes,
output_size=self.output_size
)
return cropped_batch Implementamos el Batchrandomrotate Transfour al iterando sobre todas las imágenes en el lote y aplicando una rotación aleatoria a cada una. Tenga en cuenta que esta versión no está vectorizada; Una implementación totalmente vectorizada sería más requeriría un mayor esfuerzo.
class BatchRandomRotation(T.Transform):
def __init__(self, degrees):
super().__init__()
self .degrees = degrees
def transform(self, inpt: torch.Tensor, params: dict):
# split the batch into a list of individual images
images = list(torch.unbind(inpt, dim=0))
augmented_images = []
for img_tensor in images:
# generate a random angle
angle = random.uniform(-self.degrees, self.degrees)
# apply the rotation to the single image
transformed_img = F.rotate(
img_tensor,
angle=angle
)
augmented_images.append(transformed_img)
# stack the transformed images
return torch.stack(augmented_images, dim=0)Ahora definimos lote_transform Eso imita la tubería de aumento basada en CPU definida anteriormente:
batch_transform = T.Compose([
T.ToDtype(torch.float32),
BatchRandomCrop(img_size//2),
T.Resize(img_size),
BatchRandomRotation(degrees=45.0),
T.GaussianBlur(kernel_size=7),
T.Normalize(mean=[0, 0, 0], std=[1, 1, 1])
]) Finalmente, restablecemos el conjunto de datos y actualizamos el bucle de capacitación para aplicar el nuevo lote_transform:
train_set = FakeDataset(transform=None)
train_loader = DataLoader(
dataset=train_set,
batch_size=32,
num_workers=8
)
for idx, data in enumerate(train_loader):
inputs = data[0].to(device=device, non_blocking=True)
labels = data[1].to(device=device, non_blocking=True).squeeze()
# apply augmentations
inputs = batch_transform(inputs)
optimizer.zero_grad()
output = model(inputs)
loss = criterion(output, labels)
loss.backward()
optimizer.step()
if idx == warmup:
torch.cuda.synchronize()
t0 = time.perf_counter()
if idx == total_iters:
break
torch.cuda.synchronize()
total_time = time.perf_counter() - t0
print(f'throughput: {active_batches / total_time}')Este script de entrenamiento actualizado mejora el rendimiento a 35.22 pasos por segundo: una aceleración del 72.57% sobre el resultado de la línea de base.
Aumentos de tuberías con corrientes CUDA
A continuación, cubrimos los pasos de aumento y capacitación utilizando dos transmisiones CUDA separadas: una para ejecutar la transformación de datos uno para capacitar al modelo. En cada iteración del bucle realizamos dos operaciones concurrentes:
- Entrenamos el modelo en el lote aumentado norte.
- Realizar aumentos de datos basados en GPU en lotes N+1
transform_stream = torch.cuda.Stream()
model_stream = torch.cuda.Stream()
# initialize the transformed value to None
transformed = None
for idx, data in enumerate(train_loader):
inputs = data[0]
labels_next = data[1]
if transformed is not None:
with torch.cuda.stream(model_stream):
labels = labels.to(device, non_blocking=True).squeeze()
model_stream.wait_stream(transform_stream)
optimizer.zero_grad()
output = model(transformed)
loss = criterion(output, labels)
loss.backward()
optimizer.step()
with torch.cuda.stream(transform_stream):
inputs = inputs.to(device, non_blocking=True)
transformed = batch_transform(inputs)
# Record that the tensor was produced on transform_stream
transformed.record_stream(transform_stream)
labels = labels_next
if idx == warmup:
torch.cuda.synchronize()
t0 = time.perf_counter()
if idx == total_iters:
break
torch.cuda.synchronize()
total_time = time.perf_counter() - t0
print(f'throughput: {active_batches / total_time}')Esto mejora aún más el rendimiento a 38.82 pasos por segundo: un aumento del 10.2% sobre la solución serializada y 90.20% más rápido que la línea de base original
Sensibilidad de las propiedades de la carga de trabajo
Como vimos en la Parte 1, el beneficio del canalización del uso de transmisiones CUDA varía según los detalles de la carga de trabajo. En la tabla a continuación, capturamos los resultados para varios tamaños de lotes diferentes:
A medida que aumenta el tamaño del lote, la descarga de GPU se vuelve más efectiva, aumentando significativamente el rendimiento. Al mismo tiempo, las ganancias del canalización disminuyen. Es probable que esto sea para el hecho de que los tamaños de lotes más grandes aumentan la eficiencia de la GPU, reduciendo las oportunidades de superposición.
Resumen
Cuando se trata de ejecutar cargas de trabajo AI/ML, cada milisegundos cuenta. En esta publicación exploramos el impacto de la tubería de un paso de entrenamiento AI/ML usando la transmisión CUDA en dos escenarios comunes: entrenamiento de modelos parcial y descarga de aumentos de datos a la GPU. En ambos casos, la solución canalizada superó la implementación serializada, aunque la extensión de la mejora varió significativamente en función del valor del tamaño del lote.
Como hemos enfatizado a lo largo de la publicación, el impacto esperado del uso de transmisiones CUDA puede variar mucho en función de la carga de trabajo AI/ML. Por ejemplo, en los casos en que la GPU ya se está utilizando de manera eficiente, la sobrecarga del uso de transmisiones CUDA en realidad puede conducir a una degradación en el rendimiento del tiempo de ejecución. Recomendamos encarecidamente probar esta técnica en sus propias cargas de trabajo antes de adoptar este enfoque.
Esperamos que encuentre útil la técnica descrita en esta publicación. Para obtener más consejos, trucos y técnicas para perfilar y optimizar los flujos de trabajo AI/ML, consulte las otras publicaciones en este serie.