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

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

1 个答案: