Как тренировать огромную GAN на двух (или более) графических процессорах в TensorFlow v.1.15 - PullRequest
Новая
       33

Как тренировать огромную GAN на двух (или более) графических процессорах в TensorFlow v.1.15

0 голосов
/ 24 марта 2020

Существует огромная GAN, которую я хочу тренировать на 2 графических процессорах. На этапе построения графа я связываю различные подграфы с графическими процессорами, используя менеджер контекста tf.device ().

Оказывается, что когда я тренирую GAN на одном графическом процессоре, он потребляет около 10 Гб, но когда я тренирую его на двух графических процессорах он потребляет одинаковые 10 ГБ от обоих графических процессоров. Чтобы убедиться в отсутствии дублирования узлов на каждом устройстве, я поэкспериментировал с простым NN (3 линейных + релейных слоя), где первый слой принадлежит gpu: 0, а остальные - gpu: 1 (см. прикрепление полученного графика, зеленый цвет соответствует gpu: 0, фиолетовый - gpu: 1).

Computational graph

Код, который я использовал для эксперимент: 

import tensorflow as tf
import numpy as np

LOG_DIR = './temp/logs/'

def build_linear_relu(x, n_feats, use_bias=True, scope='linear'):

    with tf.variable_scope(scope):
        w = tf.get_variable(name='weight', shape=[x.shape[1], n_feats], dtype=tf.float32)
        out = tf.matmul(x, w)
        if use_bias:
            b = tf.get_variable(name='bias', shape=[n_feats], dtype=tf.float32)
            out += b
        out = tf.nn.relu(out, name='relu')
        return out

def build_loss(y_pred):

    y_true = tf.constant(5, dtype=tf.float32, name='target')
    loss = tf.reduce_mean(tf.square(y_pred - y_true), axis=0)
    return loss

def build_model():

    x = tf.placeholder(tf.float32, shape=[None, 1024], name='input')
    with tf.device('/device:GPU:0'):
        out = build_linear_relu(x, n_feats=1024, scope='linear_gpu_0')
    with tf.device('/device:GPU:0'):
        out = build_linear_relu(out, n_feats=1024, scope='linear_gpu_1')
        out = build_linear_relu(out, n_feats=1, scope='linear_gpu_2')
        loss = build_loss(out)
    optim = tf.train.GradientDescentOptimizer(learning_rate=1e-3).minimize(loss, colocate_gradients_with_ops=True)

    return x, loss, optim

def run():
    num_steps = 10000
    x_input = np.random.randn(100000, 1024)
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True

    with tf.Session(config=config) as sess:
        x_ph, loss_op, optim_op = build_model()
        var_init = tf.global_variables_initializer()
        tb_writer = tf.summary.FileWriter(logdir=LOG_DIR, graph=sess.graph)
        sess.run(var_init)
        for it in range(num_steps):
            loss_val, _ = sess.run([loss_op, optim_op], feed_dict={x_ph: x_input})
            print('[**] Loss value on iteration {0:d}: {1}'.format(it, loss_val))


if __name__ == '__main__':
    run()
...