Skip to main content
 首页 » 编程设计

python之TensorFlow 分布式运行时模型并行 CIFAR-10

2025年05月04日37thcjp

我已尝试修改 CIFAR-10 示例以在新的 TensorFlow 分布式运行时上运行。但是,在尝试运行该程序时出现以下错误:

InvalidArgumentError: Cannot assign a device to node 'softmax_linear/biases/ExponentialMovingAverage':  
Could not satisfy explicit device specification '/job:local/task:0/device:CPU:0' 

我使用以下命令启动集群。在我运行的第一个节点上:

bazel-bin/tensorflow/core/distributed_runtime/rpc/grpc_tensorflow_server --cluster_spec='local|10.31.101.101:7777;10.31.101.224:7778' --job_name=local --task_id=0 

...在我运行的第二个节点上:

bazel-bin/tensorflow/core/distributed_runtime/rpc/grpc_tensorflow_server --cluster_spec='local|10.31.101.101:7777;10.31.101.224:7778' --job_name=local --task_id=1 

对于CIFAR-10 multi-GPU code ,我进行了简单的修改,替换了 train() 函数中的两行。下面一行:

with tf.Graph().as_default(), tf.device('/cpu:0'): 

...替换为:

with tf.Graph().as_default(), tf.device('/job:local/task:0/cpu:0'): 

和下一行:

with tf.device('/gpu:%d' % i): 

...替换为:

with tf.device('/job:local/task:0/gpu:%d' % i): 

以我的理解,第二个替换应该负责模型替换。运行一个更简单的示例,例如下面的代码,效果很好:

with tf.device('/job:local/task:0/cpu:0'): 
    c = tf.constant("Hello, distributed TensorFlow!") 
    sess.run(c) 
    print(c) 

请您参考如下方法:

我无法从你的程序中看出,但我的猜测是你还必须修改 line that creates the session指定您的一项工作任务的地址。例如,给定上面的配置,您可以这样写:

sess = tf.Session( 
    "grpc://10.31.101.101:7777", 
    config=tf.ConfigProto( 
        allow_soft_placement=True, 
        log_device_placement=FLAGS.log_device_placement)) 

碰巧的是,我们一直在努力改进该错误消息以减少混淆。如果您在 GitHub 中更新到最新版本并运行相同的代码,您应该会看到一条错误消息,解释为什么无法满足设备规范。