基本思路：
1、对数据分块，使用多个worker分别处理一个数据块，每个worker暴露两个接口，分别是损失计算的接口loss和梯度计算的接口grad；
2、同时定义full_loss和full_grad接口对每个worker的loss和grad进行聚合；
3、使用bfgs算法进行参数优化，分别使用full_loss和full_grad作为bfgs的损失函数和梯度函数，即可进行网络参数优化；
注意：在此实现中，每个worker内部每次均计算一个数据块上的损失和梯度，而非一个batch。

#0、导入依赖

import numpy as np

import os

import scipy.optimize

import tensorflow as tf

from tensorflow.examples.tutorials.mnist import input_data

import ray

import ray.experimental.tf_utils

#1、定义模型

class LinearModel(object):

    def __init__(self, shape):

        """Creates a LinearModel object."""

        x = tf.placeholder(tf.float32, [None, shape[0]])

        w = tf.Variable(tf.zeros(shape))

        b = tf.Variable(tf.zeros(shape[1]))

        self.x = x

        self.w = w

        self.b = b

        y = tf.nn.softmax(tf.matmul(x, w) + b)

        y_ = tf.placeholder(tf.float32, [None, shape[1]])

        self.y_ = y_

        cross_entropy = tf.reduce_mean(

            -tf.reduce_sum(y_ * tf.log(y), reduction_indices=[1]))

        self.cross_entropy = cross_entropy

        self.cross_entropy_grads = tf.gradients(cross_entropy, [w, b])

        self.sess = tf.Session()

        self.variables = ray.experimental.tf_utils.TensorFlowVariables(

            cross_entropy, self.sess)

    def loss(self, xs, ys):

        """计算loss"""

        return float(

            self.sess.run(

                self.cross_entropy, feed_dict={

                    self.x: xs,

                    self.y_: ys

                }))

    def grad(self, xs, ys):

        """计算梯度"""

        return self.sess.run(

            self.cross_entropy_grads, feed_dict={

                self.x: xs,

                self.y_: ys

            })

#2、定义远程worker，用于计算模型loss、grads

@ray.remote

class NetActor(object):

    def __init__(self, xs, ys):

        os.environ["CUDA_VISIBLE_DEVICES"] = ""

        with tf.device("/cpu:0"):

            self.net = LinearModel([784, 10])

            self.xs = xs

            self.ys = ys

    # 计算一个数据块的loss

    def loss(self, theta):

        net = self.net

        net.variables.set_flat(theta)

        return net.loss(self.xs, self.ys)

    # 计算一个数据块的梯度

    def grad(self, theta):

        net = self.net

        net.variables.set_flat(theta)

        gradients = net.grad(self.xs, self.ys)

        return np.concatenate([g.flatten() for g in gradients])

    def get_flat_size(self):

        return self.net.variables.get_flat_size()

#3、获取远程worker损失的函数

def full_loss(theta):

    theta_id = ray.put(theta)

    loss_ids = [actor.loss.remote(theta_id) for actor in actors]

    return sum(ray.get(loss_ids))

#4、获取远程worker梯度的函数

def full_grad(theta):

    theta_id = ray.put(theta)

    grad_ids = [actor.grad.remote(theta_id) for actor in actors]

    # 使用fmin_l_bfgs_b须转换为float64数据类型

    return sum(ray.get(grad_ids)).astype("float64")

#5、使用lbfgs进行训练

if __name__ == "__main__":

    ray.init()

    mnist = input_data.read_data_sets("MNIST_data", one_hot=True)

　　# 数据分块，每个worker跑一个数据块

    num_batches = 10

    batch_size = mnist.train.num_examples // num_batches

    batches = [mnist.train.next_batch(batch_size) for _ in range(num_batches)]

    actors = [NetActor.remote(xs, ys) for (xs, ys) in batches]

　　# 参数初始化

    dim = ray.get(actors[0].get_flat_size.remote())

    theta_init = 1e-2 * np.random.normal(size=dim)

　　# 优化

    result = scipy.optimize.fmin_l_bfgs_b(

        full_loss, theta_init, maxiter=10, fprime=full_grad, disp=True)