给出一个tensorflow多gpu的demo。用了minst 数据集，可直接运行。（tf1.8）
注：tensorflow gpu之间的操作好像有问题，比如下面的c != a + b。期待解答！
import tensorflow as tf import numpy as np import os import argparse def arg_config(): parser = argparse.ArgumentParser() parser.add_argument('-gpu', type=str, required=False, default='1, 2') args = parser.parse_args() # config log_device_placement = True # 是否打印设备分配日志 allow_soft_placement = True # 如果你指定的设备不存在，允许TF自动分配设备 gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.90, allow_growth=True) os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu # 使用 GPU id config = tf.ConfigProto(log_device_placement=log_device_placement, allow_soft_placement=allow_soft_placement, gpu_options=gpu_options) return args, config def data_loader(): mnist = tf.contrib.learn.datasets.load_dataset("mnist") train_data = mnist.train.images # Returns np.array train_data = np.reshape(train_data, newshape=(-1, 28, 28, 1)) train_labels = np.asarray(mnist.train.labels, dtype=np.int32) # print(train_data.shape) # print(train_labels.shape) # print(train_labels.dtype) ds = tf.data.Dataset.from_tensor_slices((train_data.astype(np.float32), train_labels)) ds = ds.repeat().batch(2000) it = ds.make_one_shot_iterator() dt, lb = it.get_next() return dt, lb # multi-gpu try class Model(object): def __init__(self): pass @staticmethod def var_on_cpu(name, shape, initializer=tf.truncated_normal_initializer(stddev=0.1)): with tf.device('/cpu:0'): var = tf.get_variable(name, shape, tf.float32, initializer) return var def conv2d(self, inputs, filters, kernel_size, strides=(1, 1), name=None): kernel = self.var_on_cpu(name+'/kernel', shape=kernel_size+(inputs.shape[3].value, filters,)) bias = self.var_on_cpu(name+'/bias', shape=(filters,)) conv = tf.nn.conv2d(input=inputs, filter=kernel, strides=(1,)+strides+(1,), padding='SAME', name=name) bias_add = tf.nn.bias_add(conv, bias) return tf.nn.relu(bias_add) def dense(self, inputs, units, name, activation=None): var = self.var_on_cpu(name+'/kernel', shape=(inputs.shape[1].value, units)) bias = self.var_on_cpu(name+'/bias', shape=(units, )) ds = tf.matmul(inputs, var) + bias if activation is not None: return activation(ds) else: return ds def model(self, data, training=True, scope_name='haha'): with tf.variable_scope(name_or_scope=scope_name, reuse=tf.AUTO_REUSE): conv1 = self.conv2d(inputs=data, filters=32, kernel_size=(5, 5), name='conv1') pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=[2, 2], strides=2) conv2 = self.conv2d(inputs=pool1, filters=64, kernel_size=(5, 5), name='conv2') pool2 = tf.layers.max_pooling2d(inputs=conv2, pool_size=[2, 2], strides=2) pool2_flat = tf.reshape(pool2, [-1, 7 * 7 * 64]) dense = self.dense(inputs=pool2_flat, units=1024, name='dense1', activation=tf.nn.relu) dropout = tf.layers.dropout(inputs=dense, rate=0.4, training=training) # Logits Layer logits = self.dense(inputs=dropout, units=10, name='dense2') return logits def get_loss(self, data, labels, training=True): outs = self.model(data, training=training) ls = tf.losses.sparse_softmax_cross_entropy(labels, outs) return tf.reduce_mean(ls) def average_grads(tower): averaged_grads = [] for grads_and_vars in zip(*tower): # print(grads_and_vars) grads = [] for g, _ in grads_and_vars: expanded_grad = tf.expand_dims(g, 0, 'expand_grads') grads.append(expanded_grad) grad = tf.concat(values=grads, axis=0) grad = tf.reduce_mean(input_tensor=grad, axis=0, keepdims=False) g_and_v = (grad, grads_and_vars[0][1]) averaged_grads.append(g_and_v) return averaged_grads def train_multi_gpu(): with tf.device('/cpu:0'): args, config = arg_config() gpu_num = len(args.gpu.split(',')) global_step = tf.Variable(0, dtype=tf.int64, trainable=False) model = Model() data_all, label_all = data_loader() data = tf.split(data_all, gpu_num) label = tf.split(label_all, gpu_num) optimizer = tf.train.MomentumOptimizer(0.01, 0.9) tower = [] with tf.variable_scope('gpu_vars'): for i in range(gpu_num): with tf.device("/gpu:{}".format(i)), tf.name_scope('tower_{}'.format(i)): loss_op = model.get_loss(data[i], label[i]) # print(loss_op) tf.add_to_collection(name='total_loss', value=loss_op) grads_and_vars = optimizer.compute_gradients(loss_op, tf.trainable_variables()) tower.append(grads_and_vars) mean_grads_and_vars = average_grads(tower) total_loss_op = tf.get_collection('total_loss', 'gpu_vars') with tf.control_dependencies([g for g, _ in mean_grads_and_vars]): train_op = optimizer.apply_gradients(mean_grads_and_vars, global_step=global_step, name='optimizer') # print(tf.trainable_variables()) # exit() print('running...') with tf.Session(config=config) as sess: sess.run(tf.global_variables_initializer()) step = 0 while step < 1000: _, loss = sess.run([train_op, total_loss_op]) print(step, loss) step += 1 if __name__ == "__main__": train_multi_gpu() # c != a+b with tf.device('/gpu:1'), tf.variable_scope('haha', reuse=tf.AUTO_REUSE): a = tf.get_variable(name='a', shape=[3], dtype=tf.float32, initializer=tf.truncated_normal_initializer()) # a = tf.Variable([2, 3, 4], name='a', dtype=tf.float32, ) with tf.device('/gpu:2'): b = tf.constant(value=[1, 2, 3], dtype=tf.float32, shape=[3], name='b') with tf.device('/gpu:1'): c = tf.add(a, b) with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as sess: sess.run([tf.global_variables_initializer(), tf.local_variables_initializer()]) print('c: ', sess.run([c])) print('b: ', sess.run([b])) print('a: ', sess.run([a]))