Come ottimizzare il tuo flusso di dati di input DL con un operatore PyTorch personalizzato

Ottimizzazione flusso dati input DL con operatore PyTorch personalizzato

Analisi delle prestazioni e ottimizzazione dei modelli PyTorch – Parte 5

Questo post è il quinto di una serie di post sul tema dell’analisi delle prestazioni e dell’ottimizzazione dei carichi di lavoro PyTorch basati su GPU ed è un seguito diretto della parte quattro. Nella parte quattro, abbiamo dimostrato come PyTorch Profiler e TensorBoard possono essere utilizzati per identificare, analizzare e affrontare i colli di bottiglia delle prestazioni nel pipeline di pre-elaborazione dei dati di un carico di lavoro di addestramento DL. In questo post discutiamo il supporto di PyTorch per la creazione di operatori personalizzati e dimostriamo come ci permette di risolvere i colli di bottiglia delle prestazioni nel pipeline di input dei dati, accelerare i carichi di lavoro DL e ridurre i costi di addestramento. Un ringraziamento va a Yitzhak Levi e Gilad Wasserman per il loro contributo a questo post. Il codice associato a questo post può essere trovato in questo repository GitHub.

Creazione di estensioni PyTorch

PyTorch offre diversi modi per creare operazioni personalizzate, tra cui l’estensione di torch.nn con moduli e/o funzioni personalizzate. In questo post siamo interessati al supporto di PyTorch per l’integrazione di codice C++ personalizzato. Questa capacità è importante perché alcune operazioni possono essere implementate (molto) più efficientemente e/o facilmente in C++ che in Python. Utilizzando le utility designate di PyTorch, come CppExtension, queste operazioni possono essere facilmente incluse come “estensioni” in PyTorch senza dover estrarre e ricompilare l’intero codice di PyTorch. Per ulteriori informazioni sulla motivazione di questa funzionalità e sui dettagli su come usarla, si prega di consultare il tutorial ufficiale di PyTorch sulle estensioni C++ e CUDA personalizzate. Dal momento che il nostro interesse in questo post è quello di accelerare il pipeline di pre-elaborazione dei dati basato su CPU, ci limiteremo a un’estensione C++ e non richiederemo codice CUDA. In un futuro post speriamo di dimostrare come utilizzare questa funzionalità per implementare un’estensione CUDA personalizzata al fine di accelerare il codice di addestramento in esecuzione sulla GPU.

Esempio di gioco

Nel nostro post precedente abbiamo definito un pipeline di input dei dati che iniziava con la decodifica di un’immagine JPEG di dimensioni 533×800 e quindi estraeva un ritaglio casuale di dimensioni 256×256 che, dopo alcune trasformazioni aggiuntive, veniva inserito nel ciclo di addestramento. Abbiamo utilizzato PyTorch Profiler e TensorBoard per misurare il tempo associato al caricamento dell’immagine da file e abbiamo riconosciuto l’inefficienza della decodifica. Per completezza, copiamo di seguito il codice:

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()])

Ricordiamo dal nostro post precedente che il tempo medio ottimizzato che abbiamo raggiunto era di 0,72 secondi. Presumibilmente, se fossimo stati in grado di decodificare solo il ritaglio di nostro interesse, il nostro pipeline sarebbe stato più veloce. Purtroppo, al momento della stesura di questo testo, PyTorch non include una funzione che supporta questo. Tuttavia, utilizzando gli strumenti per la creazione di operatori personalizzati, possiamo definire e implementare la nostra stessa funzione!

• L’IA potrebbe contrabbandare messaggi segreti nei meme
• Come costruire LLM per il codice?
• AR e AI Il ruolo dell’IA nella Realtà Aumentata

Funzione personalizzata per la decodifica e il ritaglio delle immagini JPEG

La libreria libjpeg-turbo è un codec di immagini JPEG che include una serie di miglioramenti e ottimizzazioni rispetto a libjpeg. In particolare, libjpeg-turbo include una serie di funzioni che ci permettono di decodificare solo un ritaglio predefinito all’interno di un’immagine, come jpeg_skip_scanlines e jpeg_crop_scanline. Se si sta utilizzando un ambiente conda, è possibile installare con il seguente comando:

conda install -c conda-forge libjpeg-turbo

Nota che libjpeg-turbo è preinstallato nell’immagine Docker di PyTorch 2.0 Deep Learning ufficiale di AWS che useremo nei nostri esperimenti di seguito.

Nel blocco di codice qui sotto, modifichiamo la funzione decode_jpeg di torchvision 0.15 per decodificare e restituire un ritaglio richiesto da un’immagine codificata in 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>();  // Imposta la struttura di decompressione  cinfo.err = jpeg_std_error(&jerr.pub);  jerr.pub.error_exit = torch_jpeg_error_exit;  /* Stabilisce il contesto di ritorno setjmp per my_error_exit. */  setjmp(jerr.setjmp_buffer);  jpeg_create_decompress(&cinfo);  torch_jpeg_set_source_mgr(&cinfo, datap, data.numel());  // Leggi le informazioni dall'intestazione.  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 si aspetta un array di puntatori a linee di scansione.     * Qui l'array ha lunghezza di un solo elemento, ma potresti chiedere     * più di una linea di scansione alla volta se è più 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) {    // Salta il resto delle linee di scansione, richiesto da 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");}

Il file C++ completo può essere trovato qui.

Nella sezione successiva, seguiremo i passaggi del tutorial di PyTorch per convertirlo in un operatore PyTorch che possiamo utilizzare nel nostro processo di pre-elaborazione.

Deploying a PyTorch Extension

Come descritto nel tutorial di PyTorch, ci sono diversi modi per distribuire un operatore personalizzato. Ci sono diverse considerazioni che potrebbero influire sul design di distribuzione. Di seguito sono riportati alcuni esempi di ciò che riteniamo importante:

1. Compilazione just-in-time: Per assicurarsi che la nostra estensione C++ sia compilata con la stessa versione di PyTorch con cui stiamo addestrando, il nostro script di distribuzione programma la compilazione del codice proprio prima dell’addestramento all’interno dell’ambiente di addestramento.
2. Suupporto multi-processo: Lo script di distribuzione deve supportare la possibilità che la nostra estensione C++ venga caricata da più processi (ad esempio, più worker DataLoader).
3. Suupporto per l’addestramento gestito: Poiché spesso ci alleniamo in ambienti di addestramento gestiti (come Amazon SageMaker), richiediamo che lo script di distribuzione supporti questa opzione. (Vedere qui per ulteriori informazioni sull’argomento della personalizzazione di un ambiente di addestramento gestito.)

Nel blocco di codice sottostante definiamo uno script setup.py semplice che compila e installa la nostra funzione personalizzata, come descritto qui.

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})

Poniamo il nostro file C++ e lo script setup.py in una cartella chiamata custom_op e definiamo un __init__.py che garantisce che lo script di configurazione venga eseguito una sola volta e da un singolo processo:

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_jpeg

Infine, rivediamo il nostro flusso di input dati per utilizzare la nostra funzione personalizzata appena creata:

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()])

Risultati

Seguendo l’ottimizzazione che abbiamo descritto, il tempo di esecuzione scende a 0,48 secondi (da 0,72) per un aumento delle prestazioni del 50%! Naturalmente, l’impatto della nostra ottimizzazione è direttamente correlato alle dimensioni delle immagini JPEG grezze e alla scelta delle dimensioni del ritaglio.

Sommario

I colli di bottiglia nel flusso di pre-elaborazione dei dati sono situazioni comuni che possono causare una carenza di GPU e rallentare l’addestramento. Date le possibili implicazioni finanziarie, è fondamentale disporre di una varietà di strumenti e tecniche per analizzarli e risolverli. In questo post abbiamo esaminato l’opzione di ottimizzare il flusso di input dati creando un’estensione personalizzata di PyTorch in C++, dimostrandone la facilità d’uso e mostrandone l’impatto potenziale. Naturalmente, i vantaggi potenziali di questo tipo di meccanismo di ottimizzazione varieranno notevolmente in base al progetto e ai dettagli del collo di bottiglia delle prestazioni.

Cosa fare dopo? La tecnica di ottimizzazione discussa qui si aggiunge a una vasta gamma di metodi di ottimizzazione del flusso di input dati che abbiamo discusso in molti dei nostri post del blog. Ti invitiamo a darci un’occhiata (ad esempio, partendo da qui).