Analisi e ottimizzazione delle prestazioni del modello 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 e un seguito diretto di parte quarta. Nella quarta parte abbiamo dimostrato come Profilo PyTorch E TensorBoard può essere utilizzato per identificare, analizzare e risolvere i colli di bottiglia delle prestazioni nella pipeline di pre-elaborazione dei dati di un carico di lavoro di formazione DL. In questo post discutiamo del supporto di PyTorch per creazione di operatori personalizzati e dimostrare come ci consente di risolvere i colli di bottiglia delle prestazioni sulla pipeline di input dei dati, accelerare i carichi di lavoro DL e ridurre i costi di formazione. Grazie a Yitzhak Levi E Gilad Wassermann per il loro contributo a questo post. Il codice associato a questo post può essere trovato in questo repository GitHub.
PyTorch offre diversi modi per creare operazioni personalizzate, tra cui torcia estensibile.nn con consuetudine Moduli e/o Funzioni. 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 utilità PyTorch designate, come CppExtensionqueste operazioni possono essere facilmente incluse come “estensioni” in PyTorch senza la necessità di estrarre e ricompilare l’intera base di codice PyTorch. Per ulteriori informazioni sulla motivazione alla base di questa funzionalità e dettagli su come utilizzarla, vedere il tutorial ufficiale di PyTorch sulle estensioni C++ e CUDA personalizzate. Poiché il nostro interesse in questo post è accelerare la pipeline di pre-elaborazione dei dati basata su CPU, sarà sufficiente un’estensione C++ e non richiederemo il codice CUDA. In un post futuro speriamo di dimostrare come utilizzare questa funzionalità per implementare un’estensione CUDA personalizzata al fine di accelerare l’esecuzione del codice di training sulla GPU.
Nel nostro messaggio precedente abbiamo definito una pipeline di input dei dati che è iniziata con la decodifica di a 533X800 JPEG e quindi estrarne un file casuale 256X256 coltura che, dopo alcune ulteriori trasformazioni, viene immessa nel ciclo di allevamento. Abbiamo usato Profilo PyTorch E TensorBoard per misurare il tempo associato al caricamento dell’immagine dal file e ha riconosciuto lo spreco della decodificazione. Per completezza riportiamo il codice qui sotto:
import numpy as np
from PIL import Image
from torchvision.datasets.vision import VisionDataset
input_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 di passo medio ottimizzato che abbiamo raggiunto era 0,72 secondi. Presumibilmente, se fossimo stati in grado di decodificare solo il raccolto a cui eravamo interessati, il nostro gasdotto avrebbe funzionato più velocemente. Sfortunatamente, al momento della stesura di questo articolo PyTorch non include una funzione che lo supporti. Tuttavia, utilizzando gli strumenti per la creazione di operazioni personalizzate, possiamo definire e implementare la nostra funzione!
IL libjpeg-turbo è un codec di immagine JPEG che include una serie di miglioramenti e ottimizzazioni rispetto a libjpeg. In particolare, libjpeg-turbo include una serie di funzioni che ci consentono di decodificare solo un ritaglio predefinito all’interno di un’immagine come jpeg_skip_scanlines E jpeg_crop_scanline. Se stai utilizzando un ambiente conda puoi eseguire l’installazione con il seguente comando:
conda install -c conda-forge libjpeg-turbo
Notare che libjpeg-turbo viene preinstallato nel funzionario Immagine Docker AWS PyTorch 2.0 Deep Learning che utilizzeremo nei nostri esperimenti di seguito.
Nel blocco di codice sottostante modifichiamo il file decodifica_jpeg funzione di torciavisione 0.15 per decodificare e restituire un ritaglio richiesto da un’immagine codificata JPEG di input.
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>();
// Setup decompression structure
cinfo.err = jpeg_std_error(&jerr.pub);
jerr.pub.error_exit = torch_jpeg_error_exit;
/* Establish the setjmp return context for my_error_exit to use. */
setjmp(jerr.setjmp_buffer);
jpeg_create_decompress(&cinfo);
torch_jpeg_set_source_mgr(&cinfo, datap, data.numel());
// read info from header.
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 expects an array of pointers to scanlines.
* Here the array is only one element long, but you could ask for
* more than one scanline at a time if that's more convenient.
*/
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) {
// Skip the rest of scanlines, required by 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");
}
È possibile trovare il file C++ completo Qui.
Nella sezione successiva, seguiremo i passaggi del tutorial PyTorch per convertirlo in un operatore PyTorch che possiamo utilizzare nella nostra pipeline di pre-elaborazione.
Come descritto nel Tutorial su PyTorchesistono diversi modi per distribuire un operatore personalizzato. Esistono numerose considerazioni che potrebbero influire sulla progettazione della distribuzione. Ecco alcuni esempi di ciò che riteniamo importante:
- Compilazione just in time: Per garantire che la nostra estensione C++ sia compilata con la stessa versione di PyTorch con cui ci alleniamo, programmiamo il nostro script di distribuzione per compilare il codice subito prima dell’addestramento all’interno dell’ambiente formativo.
- Supporto multiprocesso: Lo script di distribuzione deve supportare la possibilità che la nostra estensione C++ venga caricata da più processi (ad esempio, più lavoratori DataLoader).
- Supporto alla formazione gestita: Poiché spesso ci alleniamo in ambienti di formazione gestiti (come Amazon SageMaker) è necessario che lo script di distribuzione supporti questa opzione. (Vedere Qui per ulteriori informazioni sul tema della personalizzazione di un ambiente di formazione gestito.)
Nel blocco di codice seguente definiamo un semplice setup.py script che compila e installa la nostra funzione personalizzata, come descritto Qui.
from setuptools import setup
from 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})
Inseriamo il nostro file C++ e il file setup.py script in una cartella denominata op.personalizzata e definire un __init__.py ciò garantisce che lo script di installazione venga eseguito una sola volta e da un singolo processo:
import os
import sys
import subprocess
import shlex
import 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 la nostra pipeline di immissione dei dati per utilizzare la nostra funzione personalizzata appena creata:
from torchvision.datasets.vision import VisionDataset
input_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()))
Dopo l’ottimizzazione che abbiamo descritto, il tempo del nostro passo scende a 0,48 secondi (da 0,72) per un aumento delle prestazioni del 50%! Naturalmente, l’impatto della nostra ottimizzazione è direttamente correlato alla dimensione delle immagini JPEG grezze e alla scelta della dimensione del ritaglio.
I colli di bottiglia nella pipeline di pre-elaborazione dei dati sono eventi comuni che possono causare la carenza di GPU e rallentare l’addestramento. Date le potenziali implicazioni in termini di costi, è fondamentale disporre di una varietà di strumenti e tecniche per analizzarli e risolverli. In questo post abbiamo esaminato l’opzione di ottimizzare la pipeline di input dei dati creando un’estensione PyTorch C++ personalizzata, dimostrato la sua facilità d’uso e mostrato il suo potenziale impatto. Naturalmente, i potenziali guadagni derivanti da questo tipo di meccanismo di ottimizzazione varieranno notevolmente in base al progetto e ai dettagli del collo di bottiglia delle prestazioni.
E dopo? La tecnica di ottimizzazione qui discussa si unisce a un’ampia gamma di metodi di ottimizzazione della pipeline di input di cui abbiamo parlato in molti dei post del nostro blog. Ti invitiamo a verificarli (ad esempio, Starting Qui).
