Opening Tensorflow models in scala

how_to_save.py

sess = tf.InteractiveSession()

# [graph creation here]

with open(os.path.join(MODEL_DIR, "model.pb"), "wb") as f:
  f.write(sess.graph_def.SerializeToString())

TensorflowModel.scala

package epic.tensorflow

import org.bytedeco.javacpp.{tensorflow, BytePointer}
import org.bytedeco.javacpp.tensorflow.GraphDef


/**
 * Created by kitaev on 1/22/16.
 */
class TensorflowModel(filename: String) {
  val graphdef = new GraphDef()
  val graphdef_data = java.nio.file.Files.readAllBytes(java.nio.file.Paths.get(filename))
  val graphdef_data_ptr = new BytePointer(java.nio.ByteBuffer.wrap(graphdef_data))

  val import_ok = tensorflow.ParseProtoUnlimited(graphdef, graphdef_data_ptr)
  if (!import_ok) {
    throw new Exception(s"Graphdef import failed from: ${filename}")
  }

  graphdef.set_version(0) // TODO(nikita): figure out how versioning is supposed to work

  val s = TensorflowSession.sess.Extend(graphdef)
  if (!s.ok()) {
    throw new Exception(s"${s.ToString().getString()}")
  }
}

TensorflowSession.scala

package epic.tensorflow

import org.bytedeco.javacpp.tensorflow
import org.bytedeco.javacpp.tensorflow._

/**
 * Created by kitaev on 1/22/16.
 */

object TensorflowSession {
  val sess = new tensorflow.Session(new SessionOptions())
  private val graphdef = new tensorflow.GraphDef()
  sess.Create(graphdef)

  // [project-specific code redacted]
}