apache-spark - 将 Tensorflow 模型的预测输出保存到 hdfs 文件中

Question

我有一个经过训练的 tf 模型，我想将它应用到 hdfs 中大约有十亿个样本的大数据集。要点是我需要将 tf 模型的预测写入 hdfs 文件。但是我在tensorflow中找不到关于如何将数据保存到hdfs文件的相关API，只能找到关于读取hdfs文件的api
到目前为止，我的做法是将训练好的 tf 模型保存到本地的 pb 文件中，然后在 spark 或 Mapreduce 代码中使用 Java api 加载 pb 文件。 spark 或 mapreduce 的问题是运行速度很慢并且失败并出现 exceeds memory 错误。这是我的演示:

public class TF_model implements Serializable{

public Session session;

public TF_model(String model_path){
    try{
        Graph graph = new Graph();
        InputStream stream = this.getClass().getClassLoader().getResourceAsStream(model_path);
        byte[] graphBytes = IOUtils.toByteArray(stream);
        graph.importGraphDef(graphBytes);
        this.session = new Session(graph);
    }
    catch (Exception e){
        System.out.println("failed to load tensorflow model");
    }
}
// this is the function to predict a sample in hdfs
public int[][] predict(int[] token_id_array){
    Tensor z = session.runner()
            .feed("words_ids_placeholder", Tensor.create(new int[][]{token_id_array}))
            .fetch("softmax_prediction").run().get(0);
    double[][][] softmax_prediction = new double[1][token_id_array.length][2];
    z.copyTo(softmax_prediction);
    return softmax_prediction[0];
}}

下面是我的 Spark 代码:

val rdd = spark.sparkContext.textFile(file_path)
val predct_result= rdd.mapPartitions(pa=>{
    val tf_model = new TF_model("model.pb")
    pa.map(line=>{
      val transformed = transform(line) // omitted the transform code
      val rs = tf_model .predict(transformed)
      rs
    })
  })

我也尝试过在 hadoop 中部署 tensorflow，但找不到将大数据集写入 HDFS 的方法。

Answer 1

您可以从 hdfs 读取一次模型文件，然后使用 sc.broadcast 将您的图形的字节数组广播到分区。最后，启动负载图和预测。只是为了避免从 hdfs 多次读取文件。