PyTorch 1.0 preview: JIT compiler

In pochissimo tempo PyTorch è diventato una delle librerie più usate per il deep learning, e si avvicina ormai il rilascio della prima versione stabile. Il codice per la versione 1.0 è in realtà già disponibile da qualche giorno in developer's preview, e noi ne abbiamo approfittato per testare una delle novità più attese: 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.

Finora, 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 conterrà 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 di PyTorch precedenti 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. Per farlo, installiamo la versione in developer's preview di PyTorch dal repository ufficiale:

!pip install torch_nightly -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html

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 nuovo compiler di PyTorch permette di ottimizzare modelli per passarli in ambienti di produzione in maniera molto semplice. Mentre aspettiamo che le ultime funzionalità vengano aggiunte al JIT compiler per la release definitiva, ci sembra un buon momento per cominciare ad aggiornare i vostri modelli e prepararne la compilazione!

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