Julia se lance maintenant sur les TPU de Google

**dourouc05** · 08/11/2018, 12h38

Les processeurs tensoriels de Google parlent maintenant Julia,
XLA.jl est une manière très simple de lancer ses calculs sur un TPU

Google avait quelque peu étonné en annonçant qu'il développait ses propres processeurs : plus précisément, des accélérateurs pour les réseaux neuronaux, dénommés TPU (tensor processing unit). La première génération s'est limitée à l'inférence, la deuxième s'est étendue vers l'apprentissage de ces réseaux neuronaux, en apportant des avantages en performance brute et, surtout, en efficacité énergétique par rapport aux processeurs graphiques. Google est en train d'ouvrir l'accès à ce matériel par le biais de sa plateforme infonuagique, mais reste la difficulté de programmer ces processeurs très spécifiques. Pour cela, Google a déjà proposé son compilateur XLA, spécifique à TensorFlow. Il se focalise sur les opérations d'algèbre linéaire utilisées dans TensorFlow et cible plusieurs types de processeurs (les x86 traditionnels, les cartes graphiques NVIDIA par le truchement de CUDA, mais aussi les TPU).

Julia, au contraire, est un langage de programmation orienté performance, d'abord prévu pour le calcul scientifique (et donc l'algèbre linéaire), mais qui s'est vite étendu notamment à l'apprentissage automatique et à l'apprentissage profond en particulier. Le langage est aussi construit autour du choix précis de la fonction à appeler : les types de tous les arguments sont pris en compte pour spécialiser un maximum le code (et donc augmenter sa performance). Ce système permet, notamment, d'interagir avec les TPU sans difficulté : il suffit de définir des tableaux d'un type spécifique (qui indique que ces données doivent être hébergées sur un TPU et les calculs doivent y être effectués), le reste se passe automatiquement. Ainsi, le moteur d'exécution transfère les données au bon endroit, les rapatrie en cas de besoin. L'utilisateur ne voit pas les détails d'implémentation, ni le lien avec le compilateur XLA et son langage HLO.

En guise de tests, les développeurs de XLA.jl (le lien entre Julia et le compilateur XLA) se sont bien évidemment lancés dans des réseaux neuronaux, l'architecture VGG19 en particulier (utilisée dans la reconnaissance d'images). Pour ce faire, ils ont utilisé les paquets Julia Flux.jl (réseaux neuronaux) et Zygote.jl (différentiation automatique). En termes de performance, le résultat est intéressant. Ils ont comparé leur solution à l'utilisation de Flux.jl pure sur vingt cœurs (Intel Xeon Silver 4114, 2,20 GHz), à une implémentation équivalente avec PyTorch sur le même processeur, à leur solution simulée (avec un simulateur de TPU fonctionnant sur le même processeur), ainsi qu'à leur solution sur un vrai TPU (en séparant les temps en calcul pur et en temps total — entre l'émission du premier octet vers le TPU à travers Google Cloud et la réception des résultats).

Et dans le code ? Les changements sont minimes : il suffit d'ouvrir une session avec Google Cloud et d'annoter le code avec une macro (@tpu) et les calculs s'y lancent.

Code :

Sélectionner tout - Visualiser dans une fenêtre à part

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
using XLA
using Metalhead
using Flux
using TensorFlow
 
# Création du réseau neuronal.
ic = ImmutableChain(filter(x->!isa(x, Flux.Dropout), VGG19().layers.layers)...)
 
# Connexion à un TPU distant.
sess = Session(Graph(); target="grpc://localhost:8470")
run(sess, TensorFlow.Ops.configure_distributed_tpu())
 
# Génération des données. 
x = rand(Float32, 224, 224, 3, 1)
 
# Lancement du calcul. 
@tpu ic(x)