首页 > 解决方案 > 对于 pbtxt 格式的 tensorflow SavedModel,操作/节点的设备放置在哪里定义?

问题描述

我有一个SavedModelwithsaved_model.pbtxt并且variables\在这个 repo 的单个 GPU 上进行了预训练:https ://github.com/shalles/deeplab_v3 。我正在尝试SavedModel使用 tensorflow-serving 来服务它,它只能在多 GPU 机器中使用 GPU:0。我从https://github.com/tensorflow/serving/issues/311了解到,tensorflow-serving 使用 tensorflow 加载图形,并且该模型是在单个 GPU 上训练的。我试图用clear_devices=True标志保存模型但没有帮助,仍然在 GPU:0 上运行。

然后我尝试从https://www.tensorflow.org/guide/extend/model_files#deviceGraphDef读取in ,我知道分配给一个节点/操作的设备是在.saved_model.pbtxtNodeDef

我的问题是,在 this 中,只为assaved_model.pbtxt中的某些操作/节点分配了 CPU ,而没有专门分配 GPU。所有这些在 GPU 上执行的操作在它们的.NodeDefdevice: "/device:CPU:0"deviceNodeDef

我想知道 GPU 操作的设备放置信息保存在哪里,SavedModel我可以在图表中更改设备信息吗?谢谢你的帮助。

例如,在此saved_model.pbtxtCPU 操作被定义为:

node {
     name: "save/RestoreV2/tensor_names"
     op: "Const"
     device: "/device:CPU:0"
     ...
     }

计算操作是:

node {
  name: "resnet_v2_50/block1/unit_1/bottleneck_v2/conv2/kernel/Regularizer/l2_regularizer"
  op: "Mul"
  input: "resnet_v2_50/block1/unit_1/bottleneck_v2/conv2/kernel/Regularizer/l2_regularizer/scale"
  input: "resnet_v2_50/block1/unit_1/bottleneck_v2/conv2/kernel/Regularizer/l2_regularizer/L2Loss"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "_class"
    value {
      list {
        s: "loc:@resnet_v2_50/block1/unit_1/bottleneck_v2/conv2/weights"
      }
    }
  }
  attr {
    key: "_output_shapes"
    value {
      list {
        shape {
        }
      }
    }
  }
}

标签: tensorflowtensorflow-serving

解决方案


推荐阅读