Alle prese con PyTorch - Parte 5: JIT compiler

PyTorch è un framework di deep learning, sviluppato principalmente dal Facebook AI Research (FAIR) group, che ha guadagnato una enorme popolarità fra gli sviluppatori grazie alla combinazione di semplicità ed efficienza. Questi tutorial sono dedicati ad esplorare la libreria, partendo dai concetti più semplici fino alla definizione di modelli estremamente sofisticati.

Nella quinta parte, introduciamo una delle novità più attese di PyTorch 1.0: il just-in-time compiler per ottimizzare i modelli e portarli in produzione!

Dalla prototipazione alla messa in produzione

L'aspetto più apprezzato di PyTorch da chi lo utilizza è sicuramente l'estrema semplicità nello scrivere codice e nel prototipare modelli, combinazione di una API ben strutturata ed un meccanismo dinamico di esecuzione del codice che rende lo sviluppo molto intuitivo. Ma questo vantaggio diventa rapidamente un problema quando il codice passa dalla sperimentazione ad un ambiente di produzione, dove la velocità di esecuzione e l'ottimizzazione dei modelli diventano essenziali.

Prima della versione 1.0, in PyTorch questo si poteva gestire esportando il modello in un formato intermedio, ONNX, per poi reimportarlo in Caffe2, un'altra libreria sviluppata da Facebook con molta più attenzione agli aspetti di produzione e di deployment:

Con l'uscita di PyTorch 1.0, Facebook ha deciso di semplificare questo processo muovendosi verso l'unificazione delle due librerie, cominciando 'dal basso' combinando i due repository (leggi anche il comunicato ufficiale dal blog di Caffe2).

Come risultato, la versione 1.0 di PyTorch contiene un nuovissimo JIT compiler per ottimizzare i modelli in fase di produzione, e permettere di eseguirli anche in ambienti non Python (tra cui la nuovissima interfaccia C++ in beta), sfruttando molte delle ottimizzazioni sviluppate per Caffe2.

Il risultato è un nuovo frontend ibrido dove possiamo usare la modalità dinamica per la prototipazione, ed il compiler per portare progressivamente il modello in produzione:

Il compiler analizza l'intero modello, sfruttandone la conoscenza completa per creare un grafo estremamente ottimizzato (fino al 20-30% di miglioramento su modelli particolarmente complessi). Se usate anche TensorFlow, il procedimento è molto simile all'utilizzo della eager execution per la prototipazione, e dei grafi computazionali per la produzione. Un vantaggio notevole dell'approccio è che qualsiasi modello sviluppato con versioni precedenti di PyTorch non richiederà nessuna modifica.

Quindi, senza altri indugi, passiamo al codice!

JIT Compiler 1: tracer

Il cuore del compiler è Torch Script, un insieme di istruzioni di basso livello altamente ottimizzate che descrivono le operazioni di un modello, permettendo tra le altre cose di astrarle dall'interprete di Python ed esportarle in altri ambienti più dinamici (es., C++).

Il compiler ha due modalità principali, tutte e due di facile utilizzo:

1. Un tracer, che analizza una funzione od un modulo già istanziato, ne registra tutte le operazioni, e converte il tutto nelle relative funzioni in Torch Script.
2. Le annotazioni, che permettono di creare il grafo in Torch Script usando una serie di oggetti ed annotazioni aggiuntive sul codice che definisce i moduli stessi.

Cominciamo dal tracer, che è la modalità più semplice. Creiamo una rete neurale come abbiamo visto nei nostri tutorial:

import torch
from torch import nn

class CustomModel(nn.Module):
    def __init__(self):
        super(CustomModel, self).__init__()
        self.hidden = nn.Linear(4, 10)
        self.relu = nn.ReLU()
        self.out = nn.Linear(10, 3)

    def forward(self, x):
        x = self.relu(self.hidden(x))
        return self.out(x)

Quando instanziamo il modello e lo facciamo girare, le istruzioni vengono eseguite sequenzialmente, senza nessuna conoscenza delle istruzioni che seguiranno. Per rimediare, possiamo ora creare una versione statica del modello in Torch Script richiamando il tracer di PyTorch:

model = CustomModel()
traced_model = torch.jit.trace(model, torch.rand(1, 4))

Si noti come il tracer necessita di un esempio di input per permettergli di eseguire il modello al suo interno. Possiamo eseguire il modello compilato come qualsiasi altro modello:

traced_model(torch.rand(1, 4, requires_grad=False))
# tensor([[ 0.1572, -0.3836, -0.1341]], grad_fn=<AddBackward0>)

A differenza del modello originario, traced_model ha una conoscenza di tutte le operazioni da eseguire, codificata nella forma di un grafo. Per vedere l'effetto del compiler, possiamo stampare questo grafo:

traced_model.graph
# graph(%0 : Float(1, 4)
# ...
#  %14 : int = prim::Constant[value=0](), scope: CustomModel/ReLU[relu]
#  %15 : int = prim::Constant[value=0](), scope: CustomModel/ReLU[relu]
#  %16 : Float(1, 100) = aten::threshold(%13, %14, %15), scope: CustomModel/ReLU[relu]
# ...
#  return (%25);
# }

Le istruzioni del grafo corrispondono alle operazioni a basso livello nella libreria C++. A differenza di un modulo classico, ovviamente, non possiamo più modificarlo:

# Non eseguire, genera un errore!
traced_model.hidden = nn.Linear(4, 50)

L'aspetto essenziale è che il modello ora è slegato dall'utilizzo di un interprete Python: possiamo salvarlo su disco e ricaricarlo in altri ambienti, tra cui la nuova API C++ di PyTorch (ancora in beta). Salvare il modello è semplicissimo:

torch.jit.save(traced_model, 'model.pt')

Caricarlo ed eseguirlo da C++ (ripreso dalla guida ufficiale):

#include <torch/script.h>
#include <iostream>
#include <memory>

int main(int argc, const char* argv[]) {

  // Ricarica il modello dal file .pt
  std::shared_ptr<torch::jit::script::Module> module = torch::jit::load('model.pt');

  // Crea un vettore di input (solo 1)
  std::vector<torch::jit::IValue> inputs;
  inputs.push_back(torch::ones({1, 4}));

  // Esegue il modello
  auto output = module->forward(inputs).toTensor();

}

Niente di più facile! Come detto prima, però, il tracer è limitato nelle operazioni che può gestire, tra cui i controlli di flusso. Vediamo quindi un esempio più avanzato in cui andremo direttamente a modificare la definizione del modello.

JIT Compiler 2: Annotazioni

Supponiamo di voler aggiungere del dropout al nostro modello, il cui comportamento varia a seconda se stiamo allenando o meno il modello. Questo genere di comportamento condizionale non è possibile con il tracer, in quanto verrebbe salvato nel grafo solo uno dei due possibili comportamenti.

Per risolvere possiamo ridefinire direttamente il nostro modello in Torch Script con pochi cambiamenti:

from torch.nn import functional as F
class CustomModel(torch.jit.ScriptModule):

    def __init__(self):
        super(CustomModel, self).__init__()
        self.hidden = torch.jit.trace(nn.Linear(4, 100), torch.rand(1, 4))
        self.relu = torch.jit.trace(nn.ReLU(), torch.rand(1, 100))
        self.out = torch.jit.trace(nn.Linear(100, 3), torch.rand(1, 100))

    @torch.jit.script_method
    def forward(self, x, train : bool=False):
        x = self.relu(self.hidden(x))
        x = F.dropout(x, 0.2, train)
        return self.out(x)

Qualche commento in ordine:

1. Il modulo adesso estende torch.jit.ScriptModule invece che torch.nn.Module.
2. Tutti i moduli definiti nel costruttore devono a loro volta essere ScriptModule. Questo si ottiene lanciando il tracer su ciascuno di essi in fase di costruzione.
3. Il metodo forward del modulo viene annotato con @torch.jit.script_method. Questo lancia automaticamente la traduzione a Torch Script.

Il punto (2) è particolarmente importante, perché permette di combinare le due modalità di compilazione del modello in maniera completamente trasparente. Nel caso in cui la logica del modello sia definita in una funzione Python separata, è sufficiente annotarla con @torch.jit.script invece che @torch.jit.script_method.

Torch Script supporta nativamente buona parte delle istruzioni Python, oltre ovviamente a quasi tutte le istruzioni definite in PyTorch (la lista completa di istruzioni la trovate sul sito).

La parte più interessante è l'utilizzo di F.dropout, con un parametro esplicito che permette di controllare se stiamo allenando il modello (si noti anche l'utilizzo dei type hint per dichiarlo esplicitamente come booleano). Ispezionando il grafico, l'istruzione di dropout è correttamente tradotta nella rispettiva operazione C++:

%x : Dynamic = aten::dropout(%x.2, %16, %train)

Ed è tutto per questo tutorial!

Come abbiamo visto, il compiler di PyTorch permette di ottimizzare modelli per passarli in ambienti di produzione in maniera molto semplice, ottenendo un frontend in cui porzioni compilate e porzioni dinamiche si combinano in maniera ibrida. Nei prossimi tutorial passeremo invece ad una nuova classe di modelli: le reti ricorrenti!

Se questo articolo ti è piaciuto e vuoi tenerti aggiornato sulle nostre attività, ricordati che l'iscrizione all'Italian Association for Machine Learning è gratuita! Puoi seguirci anche su Facebook e su LinkedIn.

Previous Post Next Post