将 TensorFlow 模型的预测输出保存到 hdfs 文件中

Question

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

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];
}}

下面是我的火花代码：

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 多次读取文件。