Cómo optimizar tu tubería de entrada de datos de DL con un operador personalizado de PyTorch
Optimizando la tubería de entrada de datos de DL con PyTorch
Análisis de rendimiento y optimización de modelos PyTorch – Parte 5
Esta publicación es la quinta de una serie de publicaciones sobre el tema del análisis de rendimiento y optimización de cargas de trabajo de PyTorch basadas en GPU, y una continuación directa de la parte cuatro. En la parte cuatro, demostramos cómo se puede utilizar el Perfilador de PyTorch y TensorBoard para identificar, analizar y solucionar cuellos de botella de rendimiento en la tubería de preprocesamiento de datos de una carga de trabajo de entrenamiento de DL. En esta publicación, discutimos el soporte de PyTorch para la creación de operadores personalizados y demostramos cómo nos permite resolver cuellos de botella de rendimiento en la tubería de entrada de datos, acelerar las cargas de trabajo de DL y reducir el costo de entrenamiento. Agradecemos a Yitzhak Levi y a Gilad Wasserman por sus contribuciones a esta publicación. El código asociado a esta publicación se puede encontrar en este repositorio de GitHub.
Creando extensiones de PyTorch
PyTorch ofrece varias formas de crear operaciones personalizadas, incluyendo la extensión de torch.nn con módulos y/o funciones personalizadas. En esta publicación, nos interesa el soporte de PyTorch para integrar código C++ personalizado. Esta capacidad es importante debido al hecho de que algunas operaciones se pueden implementar de manera (mucho) más eficiente y/o fácil en C++ que en Python. Utilizando utilidades designadas de PyTorch, como CppExtension, estas operaciones se pueden incluir fácilmente como “extensiones” a PyTorch sin necesidad de extraer y recompilar todo el código base de PyTorch. Para obtener más información sobre la motivación detrás de esta función y los detalles de cómo usarla, consulte el tutorial oficial de PyTorch sobre extensiones personalizadas de C++ y CUDA. Dado que nuestro interés en esta publicación es acelerar la tubería de preprocesamiento de datos basada en CPU, nos conformaremos con una extensión de C++ y no requeriremos código CUDA. En una publicación futura, esperamos demostrar cómo utilizar esta funcionalidad para implementar una extensión personalizada de CUDA con el fin de acelerar el código de entrenamiento que se ejecuta en la GPU.
Ejemplo de juguete
En nuestra publicación anterior, definimos una tubería de entrada de datos que comenzaba con la decodificación de una imagen JPEG de 533×800 píxeles y luego extrayendo un recorte aleatorio de 256×256 píxeles que, después de algunas transformaciones adicionales, se alimenta al bucle de entrenamiento. Utilizamos el Perfilador de PyTorch y TensorBoard para medir el tiempo asociado con la carga de la imagen desde el archivo y reconocimos la ineficiencia de la decodificación. Por razones de integridad, copiamos el código a continuación:
import numpy as npfrom PIL import Imagefrom torchvision.datasets.vision import VisionDatasetinput_img_size = [533, 800]img_size = 256class FakeDataset(VisionDataset): def __init__(self, transform): super().__init__(root=None, transform=transform) size = 10000 self.img_files = [f'{i}.jpg' for i in range(size)] self.targets = np.random.randint(low=0,high=num_classes, size=(size),dtype=np.uint8).tolist() def __getitem__(self, index): img_file, target = self.img_files[index], self.targets[index] img = Image.open(img_file) if self.transform is not None: img = self.transform(img) return img, target def __len__(self): return len(self.img_files)transform = T.Compose( [T.PILToTensor(), T.RandomCrop(img_size), RandomMask(), ConvertColor(), Scale()])Recuerde de nuestra publicación anterior que el tiempo promedio optimizado que alcanzamos fue de 0.72 segundos. Presumiblemente, si pudiéramos decodificar solo el recorte en el que estamos interesados, nuestra tubería se ejecutaría más rápido. Desafortunadamente, hasta el momento de escribir esto, PyTorch no incluye una función que admita esto. Sin embargo, utilizando las herramientas para la creación de operaciones personalizadas, ¡podemos definir e implementar nuestra propia función!
Función personalizada para decodificar y recortar imágenes JPEG
La biblioteca libjpeg-turbo es un códec de imágenes JPEG que incluye varias mejoras y optimizaciones en comparación con libjpeg. En particular, libjpeg-turbo incluye varias funciones que nos permiten decodificar solo un recorte predefinido dentro de una imagen, como jpeg_skip_scanlines y jpeg_crop_scanline. Si está ejecutando en un entorno conda, puede instalarlo con el siguiente comando:
- Usa Python para descargar múltiples archivos (o URLs) en paralelo
- Construyendo un Pipeline de Aprendizaje Automático para la Clasific...
- Herramientas de IA principales para analistas de datos 2023
conda install -c conda-forge libjpeg-turboTen en cuenta que libjpeg-turbo viene preinstalado en la imagen oficial de Docker de AWS PyTorch 2.0 Deep Learning que utilizaremos en nuestros experimentos a continuación.
En el bloque de código a continuación, modificamos la función decode_jpeg de torchvision 0.15 para decodificar y devolver un recorte solicitado de una imagen codificada en formato JPEG.
torch::Tensor decode_and_crop_jpeg(const torch::Tensor& data, unsigned int crop_y, unsigned int crop_x, unsigned int crop_height, unsigned int crop_width) { struct jpeg_decompress_struct cinfo; struct torch_jpeg_error_mgr jerr; auto datap = data.data_ptr<uint8_t>(); // Configurar estructura de descompresión cinfo.err = jpeg_std_error(&jerr.pub); jerr.pub.error_exit = torch_jpeg_error_exit; /* Establecer el contexto de retorno setjmp para que my_error_exit lo use. */ setjmp(jerr.setjmp_buffer); jpeg_create_decompress(&cinfo); torch_jpeg_set_source_mgr(&cinfo, datap, data.numel()); // leer información del encabezado. jpeg_read_header(&cinfo, TRUE); int channels = cinfo.num_components; jpeg_start_decompress(&cinfo); int stride = crop_width * channels; auto tensor = torch::empty({int64_t(crop_height), int64_t(crop_width), channels}, torch::kU8); auto ptr = tensor.data_ptr<uint8_t>(); unsigned int update_width = crop_width; jpeg_crop_scanline(&cinfo, &crop_x, &update_width); jpeg_skip_scanlines(&cinfo, crop_y); const int offset = (cinfo.output_width - crop_width) * channels; uint8_t* temp = nullptr; if(offset > 0) temp = new uint8_t[cinfo.output_width * channels]; while (cinfo.output_scanline < crop_y + crop_height) { /* jpeg_read_scanlines espera una matriz de punteros a líneas de escaneo. * Aquí la matriz tiene solo un elemento, pero se puede solicitar * más de una línea de escaneo a la vez si es más conveniente. */ if(offset>0){ jpeg_read_scanlines(&cinfo, &temp, 1); memcpy(ptr, temp + offset, stride); } else jpeg_read_scanlines(&cinfo, &ptr, 1); ptr += stride; } if(offset > 0){ delete[] temp; temp = nullptr; } if (cinfo.output_scanline < cinfo.output_height) { // Saltar el resto de líneas de escaneo, necesario para jpeg_destroy_decompress. jpeg_skip_scanlines(&cinfo, cinfo.output_height - crop_y - crop_height); } jpeg_finish_decompress(&cinfo); jpeg_destroy_decompress(&cinfo); return tensor.permute({2, 0, 1});}PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("decode_and_crop_jpeg",&decode_and_crop_jpeg,"decode_and_crop_jpeg");}El archivo completo de C++ se puede encontrar aquí.
En la siguiente sección, seguiremos los pasos del tutorial de PyTorch para convertir esto en un operador de PyTorch que podamos utilizar en nuestra tubería de preprocesamiento.
Implementación de una extensión de PyTorch
Como se describe en el tutorial de PyTorch, hay diferentes formas de implementar un operador personalizado. Hay varias consideraciones que podrían influir en el diseño de tu implementación. Aquí hay algunos ejemplos de lo que consideramos importante:
- Compilación en tiempo de ejecución: Para asegurarnos de que nuestra extensión de C++ se compile con la misma versión de PyTorch con la que entrenamos, programamos nuestro script de implementación para compilar el código justo antes del entrenamiento dentro del entorno de entrenamiento.
- Soporte para múltiples procesos: El script de implementación debe admitir la posibilidad de que nuestra extensión de C++ se cargue desde varios procesos (por ejemplo, varios trabajadores de DataLoader).
- Soporte para entrenamiento administrado: Dado que a menudo entrenamos en entornos de entrenamiento administrados (como Amazon SageMaker), requerimos que el script de implementación admita esta opción. (Consulta aquí para obtener más información sobre cómo personalizar un entorno de entrenamiento administrado.)
En el bloque de código a continuación, definimos un script setup.py simple que compila e instala nuestra función personalizada, como se describe aquí.
from setuptools import setupfrom torch.utils import cpp_extensionsetup(name='decode_and_crop_jpeg', ext_modules=[cpp_extension.CppExtension('decode_and_crop_jpeg', ['decode_and_crop_jpeg.cpp'], libraries=['jpeg'])], cmdclass={'build_ext': cpp_extension.BuildExtension})Colocamos nuestro archivo C++ y el script setup.py en una carpeta llamada custom_op y definimos un __init__.py que asegura que el script de configuración se ejecute una sola vez y por un solo proceso:
import osimport sysimport subprocessimport shleximport filelockp_dir = os.path.dirname(__file__)with filelock.FileLock(os.path.join(pkg_dir, f".lock")): try: from custom_op.decode_and_crop_jpeg import decode_and_crop_jpeg except ImportError: install_cmd = f"{sys.executable} setup.py build_ext --inplace" subprocess.run(shlex.split(install_cmd), capture_output=True, cwd=p_dir) from custom_op.decode_and_crop_jpeg import decode_and_crop_jpegPor último, revisamos nuestro pipeline de entrada de datos para usar nuestra función personalizada recién creada:
from torchvision.datasets.vision import VisionDatasetinput_img_size = [533, 800]class FakeDataset(VisionDataset): def __init__(self, transform): super().__init__(root=None, transform=transform) size = 10000 self.img_files = [f'{i}.jpg' for i in range(size)] self.targets = np.random.randint(low=0,high=num_classes, size=(size),dtype=np.uint8).tolist() def __getitem__(self, index): img_file, target = self.img_files[index], self.targets[index] with torch.profiler.record_function('decode_and_crop_jpeg'): import random from custom_op.decode_and_crop_jpeg import decode_and_crop_jpeg with open(img_file, 'rb') as f: x = torch.frombuffer(f.read(), dtype=torch.uint8) h_offset = random.randint(0, input_img_size[0] - img_size) w_offset = random.randint(0, input_img_size[1] - img_size) img = decode_and_crop_jpeg(x, h_offset, w_offset, img_size, img_size) if self.transform is not None: img = self.transform(img) return img, target def __len__(self): return len(self.img_files)transform = T.Compose( [RandomMask(), ConvertColor(), Scale()])Resultados
Después de la optimización que hemos descrito, nuestro tiempo de paso disminuye a 0.48 segundos (desde 0.72) para un aumento del rendimiento del 50%. Naturalmente, el impacto de nuestra optimización está directamente relacionado con el tamaño de las imágenes JPEG sin procesar y nuestra elección del tamaño de recorte.
Resumen
Los cuellos de botella en el pipeline de preprocesamiento de datos son incidentes comunes que pueden causar escasez de GPU y ralentizar el entrenamiento. Dadas las posibles implicaciones de costos, es imperativo que tengas una variedad de herramientas y técnicas para analizar y resolverlos. En esta publicación hemos revisado la opción de optimizar el pipeline de entrada de datos mediante la creación de una extensión personalizada de PyTorch en C++, hemos demostrado su facilidad de uso y mostrado su impacto potencial. Por supuesto, las ganancias potenciales de este tipo de mecanismo de optimización variarán considerablemente según el proyecto y los detalles del cuello de botella de rendimiento.
¿Qué sigue? La técnica de optimización discutida aquí se une a una amplia gama de métodos de optimización de pipeline de entrada que hemos discutido en muchas de nuestras publicaciones de blog. Te animamos a que los consultes (por ejemplo, comenzando aquí).
