स्प्लिट आरएनएन मेमोरी खपत टेंसरफ्लो

Question

में जीपीयू के बीच समान रूप से दो जीपीयू के बीच एक seq2seq नेटवर्क के मेमोरी लोड को समान रूप से विभाजित करने का सबसे रणनीतिक तरीका जानने की कोशिश कर रहा हूं।

संक्रामक नेटवर्क के साथ, कार्य बहुत आसान है। हालांकि, मैं यह पता लगाने की कोशिश कर रहा हूं कि 2 टाइटन एक्स के मेमोरी उपयोग को अधिकतम कैसे करें। लक्ष्य सबसे बड़ा नेटवर्क बनाना है जो संयुक्त 24 जीबी मेमोरी की अनुमति देगा।

एक विचार प्रत्येक आरएनएन परत को एक अलग जीपीयू में रखना था।

GPU1 --> RNN Layer 1 & Backward Pass 
GPU2 --> RNN Layer 2,3,4 

हालांकि, backprop संगणना स्मृति का एक महत्वपूर्ण राशि की आवश्यकता है। इसलिए, एक और विचार है कि एक जीपीयू पर पूरे आगे के पास और अलग GPU पर पिछड़ा पास करना है।

GPU1 --> Forward Pass 
GPU2 --> Backward Pass 

वहाँ GPU स्मृति का कितना मापने के लिए कोई तरीका है (स्मृति लोड के सबसे हालांकि, GPU2 अभी भी लेता है) का इस्तेमाल किया जा रहा है? यह हमें यह समझने की अनुमति देगा कि प्रत्येक जीपीयू को "भरने" से पहले अधिकतम कैसे किया जाए।

एक बार 2 जीपीयू का उपयोग किया जाता है, तो मैं अंततः चार का उपयोग करना चाहता हूं। हालांकि, मुझे लगता है कि 2 जीपीयू अधिकतम करना पहला कदम है।

Answer 1

सत्य के रूप में "colocate_gradients_with_ops" को सेट करना शायद काम कर सकता है। यह जीपीयू मेमोरी को समान रूप से आवंटित करने की अनुमति देता है।

optimizer = tf.train.AdamOptimizer(learning_rate) 
gvs = optimizer.compute_gradients(loss, colocate_gradients_with_ops=True) 
train_op = optimizer.apply_gradients(gvs, global_step=self.global_step)