如何在 TensorFlow(python API)中将多个设备手动分配给一个操作?
例如,为了给一个图分配一个 cpu,我会这样做:
with tf.device("/cpu:0"):
# Define operation here
如果我想将 cpu 0、1、2、3 都分配给同一个操作怎么办?
似乎我不能使用设备名称列表作为 tf.device 的 arg。任何帮助或指导将不胜感激。谢谢!
请您参考如下方法:
你可以拥有multiple with contexts像这样:
with tf.device("/cpu:0") as cpu0, tf.device("/cpu:1") as cpu1:
pass
