Análisis de rendimiento y optimización de modelos PyTorch – Parte 2
'Performance analysis and optimization of PyTorch models - Part 2'
Cómo Identificar y Reducir la Computación de la CPU en su Paso de Entrenamiento con PyTorch Profiler y TensorBoard
Esta es la segunda parte de una serie de publicaciones sobre el tema de analizar y optimizar un modelo PyTorch que se ejecuta en una GPU. En nuestra primera publicación, demostramos el proceso, y el potencial significativo, de analizar y optimizar iterativamente un modelo PyTorch utilizando PyTorch Profiler y TensorBoard. En esta publicación nos enfocaremos en un tipo específico de problema de rendimiento que es particularmente prevalente en PyTorch debido a su uso de una ejecución ágil: La dependencia de la CPU para porciones de la ejecución del modelo. Identificar la presencia y la fuente de estos tipos de problemas puede ser bastante difícil y a menudo requiere el uso de un analizador de rendimiento dedicado. En esta publicación, compartiremos algunos consejos para identificar estos problemas de rendimiento al usar PyTorch Profiler y el complemento PyTorch Profiler TensorBoard.
Los Pros y Contras de la Ejecución Ágil
Uno de los principales atractivos de PyTorch es su modo de ejecución ágil. En el modo ágil, cada operación de PyTorch que forma el modelo se ejecuta de manera independiente tan pronto como se alcanza. Esto se contrasta con el modo de gráficos en el que todo el modelo es precompilado en un solo gráfico de manera óptima para ejecutarse en la GPU y se ejecuta como un todo. Por lo general, esta precompilación resulta en un mejor rendimiento (por ejemplo, ver aquí). En el modo ágil, el contexto de programación regresa a la aplicación después de cada operación, lo que nos permite acceder y evaluar tensores arbitrarios. Esto facilita la construcción, el análisis y la depuración de modelos de ML. Por otro lado, también hace que nuestro modelo sea más susceptible a la inserción (a veces accidental) de bloques de código subóptimos. Como demostraremos, saber cómo identificar y solucionar estos bloques de código puede tener un impacto significativo en la velocidad de su modelo.
Ejemplo de Juguete
En los siguientes bloques presentamos el ejemplo de juguete que usaremos para nuestra demostración. El código se basa muy vagamente en el ejemplo de nuestra publicación anterior y en la función de pérdida definida en este tutorial de PyTorch.
Comenzamos definiendo un modelo de clasificación simple. Su arquitectura no es significativa para esta publicación.
import torchimport torch.nn as nnimport torch.nn.functional as Fimport torch.optimimport torch.profilerimport torch.utils.dataimport torchvision.modelsimport torchvision.transforms as Tfrom torchvision.datasets.vision import VisionDatasetimport numpy as npfrom PIL import Image# muestra del modeloclass Net(nn.Module): def __init__(self): super().__init__() self.conv1 = nn.Conv2d(3, 8, 3, padding=1) self.conv2 = nn.Conv2d(8, 12, 3, padding=1) self.conv3 = nn.Conv2d(12, 16, 3, padding=1) self.conv4 = nn.Conv2d(16, 20, 3, padding=1) self.conv5 = nn.Conv2d(20, 24, 3, padding=1) self.conv6 = nn.Conv2d(24, 28, 3, padding=1) self.conv7 = nn.Conv2d(28, 32, 3, padding=1) self.conv8 = nn.Conv2d(32, 10, 3, padding=1) self.pool = nn.MaxPool2d(2, 2) def forward(self, x): x = self.pool(F.relu(self.conv1(x))) x = self.pool(F.relu(self.conv2(x))) x = self.pool(F.relu(self.conv3(x))) x = self.pool(F.relu(self.conv4(x))) x = self.pool(F.relu(self.conv5(x))) x = self.pool(F.relu(self.conv6(x))) x = self.pool(F.relu(self.conv7(x))) x = self.pool(F.relu(self.conv8(x))) x = torch.flatten(x, 1) # aplanar todas las dimensiones excepto el batch return x
Luego, definimos una función de pérdida de entropía cruzada bastante estándar. Esta función de pérdida será el foco principal de nuestra discusión.
- Una Guía Práctica para el Aprendizaje por Transferencia utilizando ...
- Word2Vec, GloVe y FastText, explicados.
- WAYVE presenta GAIA-1 Un nuevo modelo de inteligencia artificial ge...
def log_softmax(x): return x - x.exp().sum(-1).log().unsqueeze(-1)def weighted_nll(pred, target, weight): assert target.max() < 10 nll = -pred[range(target.shape[0]), target] nll = nll * weight[target] nll = nll / weight[target].sum() sum_nll = nll.sum() return sum_nll# definición de pérdida personalizadaclass CrossEntropyLoss(nn.Module): def forward(self, input, target): pred = log_softmax(input) loss = weighted_nll(pred, target, torch.Tensor([0.1]*10).cuda()) return loss
Por último, definimos el conjunto de datos y el bucle de entrenamiento:
# conjunto de datos con imágenes aleatorias que imitan las propiedades de CIFAR10class FakeCIFAR(VisionDataset): def __init__(self, transform): super().__init__(root=None, transform=transform) self.data = np.random.randint(low=0,high=256,size=(10000,32,32,3),dtype=np.uint8) self.targets = np.random.randint(low=0,high=10,size=(10000),dtype=np.uint8).tolist() def __getitem__(self, index): img, target = self.data[index], self.targets[index] img = Image.fromarray(img) if self.transform is not None: img = self.transform(img) return img, target def __len__(self) -> int: return len(self.data)transform = T.Compose( [T.Resize(256), T.PILToTensor()])train_set = FakeCIFAR(transform=transform)train_loader = torch.utils.data.DataLoader(train_set, batch_size=1024, shuffle=True, num_workers=8, pin_memory=True)device = torch.device("cuda:0")model = Net().cuda(device)criterion = CrossEntropyLoss().cuda(device)optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)model.train()# bucle de entrenamiento envuelto con objeto profilerwith torch.profiler.profile( schedule=torch.profiler.schedule(wait=1, warmup=4, active=3, repeat=1), on_trace_ready=torch.profiler.tensorboard_trace_handler(’./log/example’), record_shapes=True, profile_memory=True, with_stack=True) as prof: for step, data in enumerate(train_loader): inputs = data[0].to(device=device, non_blocking=True) labels = data[1].to(device=device, non_blocking=True) inputs = (inputs.to(torch.float32) / 255. - 0.5) / 0.5 if step >= (1 + 4 + 3) * 1: break outputs = model(inputs) loss = criterion(outputs, labels) optimizer.zero_grad(set_to_none=True) loss.backward() optimizer.step() prof.step()
Un desarrollador experimentado de PyTorch ya puede haber notado que nuestro ejemplo contiene varias líneas de código ineficientes en la función de pérdida. Al mismo tiempo, no hay nada obviamente incorrecto con ella y este tipo de ineficiencias no son infrecuentes. Si desea probar su habilidad en PyTorch, intente encontrar tres problemas con nuestra implementación de la pérdida de entropía cruzada antes de continuar leyendo. En las siguientes secciones asumiremos que no pudimos encontrar estos problemas por nuestra cuenta y mostraremos cómo podemos usar PyTorch Profiler y su plugin asociado de TensorBoard para identificarlos.
Al igual que en nuestro post anterior, ejecutaremos iterativamente un experimento, identificaremos problemas de rendimiento e intentaremos solucionarlos. Ejecutaremos nuestros experimentos en una instancia de Amazon EC2 g5.2xlarge (que contiene una GPU NVIDIA A10G y 8 vCPUs) y utilizando la imagen Docker oficial de AWS PyTorch 2.0. Nuestra elección de entorno de entrenamiento fue algo arbitraria y no debe ser vista como un respaldo para ninguno de sus componentes.
Resultados de rendimiento iniciales
En la imagen a continuación mostramos la pestaña Resumen del informe de rendimiento del script anterior.

Como podemos ver, nuestra utilización de la GPU está en un 92,04% relativamente alto y nuestro tiempo de paso es de 216 milisegundos. (Al igual que en nuestro post anterior, la pestaña Resumen en torch-tb-profiler versión 0.4.1 suma el tiempo de paso de los tres pasos de entrenamiento.) A partir de este informe, es posible que no piense que haya algo mal con nuestro modelo. Sin embargo, la Vista de trazas del informe de rendimiento cuenta una historia completamente diferente:

Como se destacó anteriormente, ¡el pase hacia adelante de nuestra función de pérdida de entropía cruzada sola ocupa 211 de los 216 milisegundos del paso de entrenamiento! Esto es una clara indicación de que algo está mal. Nuestra función de pérdida contiene un pequeño número de cálculos en comparación con el modelo y ciertamente no debería representar el 98% del tiempo del paso. Al examinar de cerca la pila de llamadas, podemos ver algunas llamadas de funciones que refuerzan nuestras sospechas, incluyendo “to”, “copy_” y “cudaStreamSynchronize”. Esta combinación generalmente indica que se está copiando datos desde la CPU a la GPU, lo cual no es algo que queramos que suceda en medio de nuestro cálculo de pérdida. En este caso, nuestro problema de rendimiento también se alinea con una breve caída en la utilización de la GPU, como se destaca en la imagen. Sin embargo, esto no siempre es el caso. A menudo, las caídas en la utilización de la GPU no estarán alineadas con el problema de rendimiento o es posible que no se vean en absoluto.
Ahora sabemos que tenemos un problema de rendimiento en nuestra función de pérdida y que es probable que esté relacionado con la copia de tensores desde el host a la GPU. Sin embargo, esto podría no ser suficiente para identificar la línea de código precisa que está causando el problema. Para facilitar nuestra búsqueda, envolveremos cada línea de código con un administrador de contexto torch.profiler.record_function etiquetado y volveremos a ejecutar el análisis de perfilado.
# definición de pérdida personalizadaclass CrossEntropyLoss(nn.Module): def forward(self, input, target): with torch.profiler.record_function('log_softmax'): pred = log_softmax(input) with torch.profiler.record_function('define_weights'): weights = torch.Tensor([0.1]*10).cuda() with torch.profiler.record_function('weighted_nll'): loss = weighted_nll(pred, target, torch.Tensor([0.1]*10).cuda()) return loss
La adición de las etiquetas nos ayuda a identificar la definición de peso, o más precisamente, la copia de los pesos en la GPU, como la línea de código problemática.

Optimización #1: Eliminar las copias redundantes de host a GPU del paso de entrenamiento
Una vez que hemos identificado nuestro primer problema, solucionarlo es bastante trivial. En el bloque de código a continuación, copiamos nuestro vector de peso a la GPU una sola vez en la función de inicialización de pérdida:
class CrossEntropyLoss(nn.Module): def __init__(self): super().__init__() self.weight = torch.Tensor([0.1]*10).cuda() def forward(self, input, target): with torch.profiler.record_function('log_softmax'): pred = log_softmax(input) with torch.profiler.record_function('weighted_nll'): loss = weighted_nll(pred, target, self.weight) return loss
La imagen a continuación muestra los resultados del análisis de rendimiento después de esta solución:

Decepcionantemente, nuestra primera optimización tuvo un impacto muy marginal en el tiempo del paso. Si examinamos el informe de Vista de traza, podemos ver que tenemos un nuevo problema de rendimiento severo que debemos abordar.

Nuestro nuevo informe indica un problema proveniente de nuestra función de weighted_nll. Como antes, usamos torch.profiler.record_function para identificar la línea de código problemática. En este caso, es la llamada assert.
def weighted_nll(pred, target, weight): with torch.profiler.record_function('assert'): assert target.max() < 10 with torch.profiler.record_function('range'): r = range(target.shape[0]) with torch.profiler.record_function('index'): nll = -pred[r, target] with torch.profiler.record_function('nll_calc'): nll = nll * weight[target] nll = nll/ weight[target].sum() sum_nll = nll.sum() return sum_nll
Es importante tener en cuenta que este problema ya existía en el experimento base, pero estaba oculto por nuestro problema de rendimiento anterior. En el curso de la optimización del rendimiento, no es raro que aparezcan problemas graves que estaban ocultos por otros problemas de repente de esta manera.
Un análisis más detallado de la pila de llamadas muestra llamadas a “item”, “_local_scalar_dense” y “cudaMemcpyAsync”. Esto es a menudo una indicación de que los datos se están copiando desde la GPU al host. De hecho, nuestra llamada assert, que se realiza en la CPU, requiere acceso al tensor objetivo que reside en la GPU, lo que provoca una copia de datos altamente ineficiente.
Optimización #2: Eliminar las copias redundantes de GPU a host en el paso de entrenamiento
Si bien puede ser justificado verificar la legalidad de las etiquetas de entrada, debe hacerse de tal manera que no afecte negativamente el rendimiento de nuestro entrenamiento. En nuestro caso, arreglar el problema es simplemente cuestión de mover la assert a la tubería de entrada de datos, antes de que las etiquetas sean copiadas en la GPU. Tras la eliminación de la assert, nuestro rendimiento sigue siendo en su mayoría inalterado:

Nota importante: Aunque nuestro objetivo suele ser intentar reducir las copias entre el host y la GPU en el paso hacia adelante, hay momentos en los que esto no es posible (por ejemplo, si requerimos un kernel que no es compatible con la GPU) o no es deseable (por ejemplo, si ejecutar un kernel particular en la CPU aumentará el rendimiento).
Analizar la Vista de trazas nos presenta nuestro siguiente problema de rendimiento:

Una vez más, vemos que nuestra optimización anterior ha revelado un nuevo problema de rendimiento grave, esta vez al indexar nuestro tensor pred. Los índices son definidos por los tensores r y target. Si bien el tensor target ya reside en la GPU, el tensor r, que fue definido en la línea anterior, no lo hace. Esto, una vez más, desencadena una ineficiente copia de datos del host a la GPU.
Optimización #3: Reemplazar range por torch.arange
La función range de Python produce una lista en la CPU. La presencia de cualquier lista en su paso de entrenamiento debería ser una señal de advertencia. En el bloque de código a continuación, reemplazamos el uso de range con torch.arange y lo configuramos para crear el tensor de salida directamente en la GPU:
def weighted_nll(pred, target, weight): with torch.profiler.record_function('range'): r = torch.arange(target.shape[0], device="cuda:0") with torch.profiler.record_function('index'): nll = -pred[r, target] with torch.profiler.record_function('nll_calc'): nll = nll * weight[target] nll = nll/ weight[target].sum() sum_nll = nll.sum() return sum_nll
Los resultados de esta optimización se muestran a continuación:

¡Ahora estamos hablando! Nuestro tiempo de paso ha bajado a 5.8 milisegundos, un aumento de rendimiento del 3700%.
La Vista de trazas actualizada muestra que la función de pérdida ha disminuido a un valor muy razonable de 0.5 milisegundos.

Pero todavía hay margen de mejora. Veamos de cerca la vista de traza de la función weighted_nll que ocupa la mayoría del cálculo de pérdida.

Podemos ver en la traza que la función está formada por múltiples bloques pequeños, cada uno de los cuales se asigna finalmente a un kernel individual de CUDA que se carga en la GPU mediante la llamada CudaLaunchKernel. Idealmente, nos gustaría reducir el número total de kernels de GPU para reducir la cantidad de interacción entre la CPU y la GPU. Una forma de hacerlo es preferir, siempre que sea posible, operadores de PyTorch de nivel superior, como torch.nn.NLLLoss. Se presume que estas funciones “fusionan” las operaciones subyacentes, lo que requiere un menor número total de kernels.
Optimización #4: Reemplazar NLL personalizado con torch.nn.NLLLoss
El bloque de código a continuación contiene nuestra definición de pérdida actualizada, que ahora utiliza torch.nn.NLLLoss.
class CrossEntropyLoss(nn.Module): def __init__(self): super().__init__() self.weight = torch.Tensor([0.1]*10).cuda() def forward(self, input, target): pred = log_softmax(input) nll = torch.nn.NLLLoss(self.weight) loss = nll(pred, target) return loss
Aquí hemos tomado la libertad de introducir otro error común que procederemos a demostrar.
El uso de la función de nivel superior reduce aún más nuestro tiempo de paso a 5.3 milisegundos (frente a 5.8).

Sin embargo, si observamos de cerca la Vista de traza, podemos ver que una parte significativa de la función de pérdida se gasta ahora en inicializar el objeto torch.nn.NLLLoss.

Al mirar hacia atrás en nuestra función de pérdida, podemos ver que estamos inicializando un nuevo objeto NLLLoss en cada iteración del paso de entrenamiento. Naturalmente, la inicialización del objeto ocurre en la CPU, y aunque (en nuestro caso) es relativamente rápido, es algo que nos gustaría evitar hacer durante nuestro paso de entrenamiento.
Optimización #5: Evite inicializar objetos en el paso de entrenamiento
En el bloque de código a continuación hemos modificado nuestra implementación de pérdida para que se cree una única instancia de torch.nn.NLLLoss en la función init.
class CrossEntropyLoss(nn.Module): def __init__(self): super().__init__() self.weight = torch.Tensor([0.1]*10).cuda() self.nll = torch.nn.NLLLoss(self.weight) def forward(self, input, target): pred = log_softmax(input) loss = self.nll(pred, target) return loss
Los resultados muestran una mejora aún mayor en el tiempo de paso, que ahora es de 5.2 milisegundos.
Optimización #6: Use torch.nn.CrossEntropyLoss en lugar de pérdida personalizada
PyTorch incluye una torch.nn.CrossEntropyLoss integrada que ahora evaluamos y comparamos con nuestra implementación de pérdida personalizada.
criterion = torch.nn.CrossEntropyLoss().cuda(device)
El tiempo resultante de este paso es un nuevo mínimo de 5 milisegundos, lo que representa un aumento total del rendimiento del 4200% (en comparación con los 216 milisegundos con los que comenzamos).
La mejora de rendimiento en el proceso de cálculo de pérdida es aún más dramática: de un punto de partida de 211 milisegundos, hemos bajado hasta los 79 microsegundos (¡¡!), como se muestra a continuación:
Optimización #7: Compilar la función de pérdida
En nuestro intento final de optimización, configuraremos la función de pérdida para que se ejecute en modo gráfico utilizando la API torch.compile. Como discutimos en detalle en esta publicación y demostramos en la precuela de esta publicación, torch.compile utilizará técnicas como la fusión de kernel y la ejecución fuera de orden para asignar la función de pérdida en núcleos de bajo nivel de cómputo de una manera que sea óptima para el acelerador de entrenamiento subyacente.
criterion = torch.compile(torch.nn.CrossEntropyLoss().cuda(device))
La imagen a continuación muestra el resultado de la vista de trazado de este experimento.
Lo primero que podemos ver es la aparición de términos que contienen “OptimizedModule” y “dynamo”, que son indicativos del uso de torch.compile. También podemos ver que, en la práctica, la compilación del modelo no redujo el número de núcleos cargados por la función de pérdida, lo que significa que no identificó oportunidades para una fusión adicional de núcleos. De hecho, en nuestro caso, la compilación de pérdida hizo que el tiempo del pase hacia adelante de la función de pérdida aumentara de 79 a 154 microsegundos. Parece que CrossEntropyLoss no es lo suficientemente sustancial como para beneficiarse de esta optimización.
Es posible que se pregunte por qué no podemos simplemente aplicar la compilación de antorcha a nuestra función de pérdida inicial y confiar en ella para compilar nuestro código de manera óptima. Esto podría ahorrarnos todo el problema de la optimización paso a paso que describimos anteriormente. El problema con este enfoque es que, aunque la compilación de PyTorch 2.0 (en el momento de escribir esto) optimiza ciertos tipos de cruces de GPU a CPU, algunos tipos provocarán el bloqueo de la compilación de gráficos, y otros resultarán en la creación de múltiples gráficos pequeños en lugar de uno grande. La última categoría causa interrupciones en el gráfico que limitan la capacidad de la función de compilación de antorcha para mejorar el rendimiento. (Una forma de abordar esto es llamar a torch.compile con la bandera de fullgraph establecida en True). Consulte nuestra publicación anterior para obtener más detalles sobre el uso de esta opción.
Resultados
En la tabla a continuación, resumimos los resultados de los experimentos que hemos realizado:

Nuestras sucesivas optimizaciones han llevado a un aumento de rendimiento impresionante del 4143%. Recuerde que comenzamos con una función de pérdida aparentemente inocente. Sin un análisis detallado del comportamiento de nuestra aplicación, es posible que nunca hayamos sabido que algo estaba mal y habríamos continuado con nuestras vidas pagando 41 veces(!!) más de lo necesario.
Puede haber notado que la utilización de la GPU disminuyó significativamente en nuestras pruebas finales. Esto indica un potencial importante para una mayor optimización de rendimiento. Aunque nuestra demostración ha llegado a su fin, nuestro trabajo no ha terminado. Consulte nuestra publicación anterior para obtener algunas ideas sobre cómo proceder a partir de aquí.
Conclusiones
Resumamos algunas de las cosas que hemos aprendido. Dividimos el resumen en dos partes. En la primera, describimos algunos hábitos de codificación que pueden afectar el rendimiento de entrenamiento. En la segunda, recomendamos algunos consejos para la realización de un perfil de rendimiento. Tenga en cuenta que estas conclusiones se basan en el ejemplo que hemos compartido en esta publicación y es posible que no se apliquen a su propio caso de uso. Los modelos de aprendizaje automático varían mucho en propiedad y comportamiento. Por lo tanto, se recomienda encarecidamente que evalúe estas conclusiones en función de los detalles de su propio proyecto.
Consejos de programación
La forma en que implementas el pase hacia adelante de tu modelo puede tener un impacto significativo en su rendimiento. Aquí enumeramos solo algunas recomendaciones basadas en el ejemplo que cubrimos en esta publicación.
- Evita inicializar tensores constantes en el pase hacia adelante. Hazlo en el constructor en su lugar.
- Evita el uso de aserciones en tensores que residen en la GPU en el pase hacia adelante. Muévelos a la canalización de entrada de datos y/o verifica si PyTorch tiene algún método incorporado para realizar la verificación de datos que necesitas.
- Evita el uso de listas. Verifica si usar torch.arange para crear un tensor directamente en el dispositivo puede ser una mejor alternativa.
- Usa operadores de PyTorch como torch.nn.NLLLoss y torch.nn.CrossEntropyLoss en lugar de crear tus propias implementaciones de pérdida.
- Evita la inicialización de objetos en el pase hacia adelante. Hazlo en el constructor en su lugar.
- Considera usar torch.compile cuando sea relevante.
Consejos de análisis de rendimiento
Como demostramos, la Vista de trazas del complemento PyTorch Profiler de Tensorboard fue fundamental para identificar los problemas de rendimiento en nuestro modelo. A continuación, resumimos algunas de las principales conclusiones de nuestro ejemplo:
- La alta utilización de la GPU NO es necesariamente una señal de que tu código se esté ejecutando de manera óptima.
- Busca porciones del código que tardan más de lo esperado.
- Usa torch.profiler.record_function para localizar problemas de rendimiento.
- Los descensos en la utilización de la GPU no están necesariamente alineados con la fuente del problema de rendimiento.
- Busca copias de datos no intencionales del host a la GPU. Estos se identifican típicamente por las llamadas a “to”, “copy_” y “cudaStreamSynchronize”, que puedes buscar en la Vista de trazas.
- Busca copias de datos no intencionales de la GPU al host. Estos se identifican típicamente por las llamadas a “item” y “cudaStreamSynchronize”, que puedes buscar en la Vista de trazas.
Resumen
En esta publicación nos hemos centrado en problemas de rendimiento en aplicaciones de entrenamiento resultantes de la interacción redundante entre la CPU y la GPU durante el pase hacia adelante del paso de entrenamiento. Demostramos cómo los analizadores de rendimiento como PyTorch Profiler y su complemento asociado TensorBoard pueden usarse para identificar tales problemas y facilitar una mejora significativa del rendimiento.
Al igual que en nuestra publicación anterior, enfatizamos que el camino hacia una optimización exitosa variará considerablemente según los detalles del proyecto de entrenamiento, incluida la arquitectura del modelo y el entorno de entrenamiento. En la práctica, alcanzar tus objetivos puede ser más difícil que en el ejemplo que presentamos aquí. Algunas de las técnicas que describimos pueden tener poco impacto en tu rendimiento o incluso empeorarlo. También señalamos que las optimizaciones precisas que elegimos, y el orden en que elegimos aplicarlas, fueron algo arbitrarias. Te recomendamos encarecidamente que desarrolles tus propias herramientas y técnicas para alcanzar tus objetivos de optimización en función de los detalles específicos de tu proyecto.