Using Ray with TensorFlow

本文（英版）档介绍了Ray中结合TensorFlow使用的最佳实践。

要查看使用TensorFlow的更多相关示例，请查看A3C，ResNet，Policy Gradients，LBFGS

 

如果在分布式环境中训练深层网络，则可能需要在进程（或计算机）之间运送深层网络。 例如，你可能在一台计算机上更新模型，然后使用该模型在另一台计算机上计算梯度。 但是，运输模型并不总是顺利的。

 

例如，直接尝试序列化TensorFlow图计算会产生不同的结果。 一些示例失败，一些成功（但产生非常大的字符串）。

 

此外，创建TensorFlow计算图可能需要几十秒，因此序列化图并在另一个进程中重新创建它将是低效的。 更好的解决方案是在开始时在每个worker创建相同的TensorFlow图，然后仅在worker之间传递权重。

 

假设我们有一个简单的网络定义（这个定义是从TensorFlow文档中修改的）。

import tensorflow as tf
import numpy as np
x_data = tf.placeholder(tf.float32, shape=[100])
y_data = tf.placeholder(tf.float32, shape=[100])
w = tf.Variable(tf.random_uniform([1], -1.0, 1.0))
b = tf.Variable(tf.zeros([1]))
y = w * x_data + b
loss = tf.reduce_mean(tf.square(y - y_data))
# 参考API，https://blog.csdn.net/xierhacker/article/details/53174558
optimizer = tf.train.GradientDescentOptimizer(0.5)
grads = optimizer.compute_gradients(loss)
train = optimizer.apply_gradients(grads)
init = tf.global_variables_initializer()
sess = tf.Session()

要提取权重并设置权重，可以使用以下辅助方法。

import ray
variables = ray.experimental.TensorFlowVariables(loss, sess)

TensorFlowVariables对象提供了获取和设置权重以及收集模型中所有变量的方法。

 

现在我们可以使用这些方法来提取权重，并将它们设置回网络中，如下所示。

# First initialize the weights.
sess.run(init)
# Get the weights
weights = variables.get_weights()  # Returns a dictionary of numpy arrays
# Set the weights
variables.set_weights(weights)

注意：如果我们使用如下所示的assign方法设置权重，则每次调用assign都会向图中添加一个节点，并且图形会随着时间的推移而变得难以管理。

w.assign(np.zeros(1))  # This adds a node to the graph every time you call it.
b.assig