Comment diviser un graphique TensorFlow (modèle) sur plusieurs GPU pour éviter les MOO?

Question

J'ai donc ce modèle très grand et profond que j'ai implémenté avec TensorFlow r1.2, fonctionnant sur un NVIDIA Tesla k40 avec 12 Go de mémoire. Le modèle comprend plusieurs RNN, un ensemble de matrices de poids et d'inclusion ainsi que des vecteurs de biais. Lorsque j'ai lancé le programme de formation, il a fallu environ 2 à 3 heures pour construire le modèle, puis il s'est écrasé en raison de problèmes de MOO. J'ai essayé de réduire la taille du lot à 1 échantillon de données par lot, mais j'ai toujours rencontré le même problème.

Si je google tensorflow muitlple gpu, les exemples que je trouve principalement concentré sur l'utilisation de plusieurs processeurs graphiques par parallèle la conception du modèle, ce qui signifie avoir chaque GPU exécuter le même graphique et ont la CPU calcule le gradient totale ainsi se propager retour à chaque paramètre.

Je sais qu'une solution possible pourrait être l'exécution du modèle sur un GPU avec plus de mémoire. Mais je me demande si il ya un moyen de divisé mon graphique (modèle) en différentes parties séquentiellement et les attribuer à différents GPU?

Answer 1

Le guide officiel sur using GPUs vous montre cet exemple dans "Utilisation de plusieurs GPU". Vous avez juste besoin de créer les opérations dans différents contextes tf.device; les nœuds seront toujours ajoutés au même graphique, mais ils seront annotés avec des directives de périphérique indiquant où ils doivent être exécutés. Par exemple:

with tf.device("/gpu:0"): 
    net0 = make_subnet0() 
with tf.device("/gpu:1"): 
    net1 = make_subnet1() 
result = combine_subnets(net0, net1)