org.tensorflow.framework.MetaGraphDef Java Examples

/**
 * Note: Please use Models from zoltar-models module.
 *
 * <p>Returns a TensorFlow model with metadata given {@link SavedModelBundle} export directory URI
 * and {@link Options}.
 */
public static TensorFlowModel create(
    final Model.Id id,
    final URI modelResource,
    final Options options,
    final String signatureDefinition)
    throws IOException {
  // GCS requires that directory URIs have a trailing slash, so add the slash if it's missing
  // and the URI starts with 'gs'.
  final URI normalizedUri =
      !CloudStorageFileSystem.URI_SCHEME.equalsIgnoreCase(modelResource.getScheme())
              || modelResource.toString().endsWith("/")
          ? modelResource
          : URI.create(modelResource.toString() + "/");
  final URI localDir = FileSystemExtras.downloadIfNonLocal(normalizedUri);
  final SavedModelBundle model =
      SavedModelBundle.load(localDir.toString(), options.tags().toArray(new String[0]));
  final MetaGraphDef metaGraphDef;
  try {
    metaGraphDef = extractMetaGraphDefinition(model);
  } catch (TensorflowMetaGraphDefParsingException e) {
    throw new IOException(e);
  }
  final SignatureDef signatureDef = metaGraphDef.getSignatureDefOrThrow(signatureDefinition);

  return new AutoValue_TensorFlowModel(
      id,
      model,
      options,
      metaGraphDef,
      signatureDef,
      toNameMap(signatureDef.getInputsMap()),
      toNameMap(signatureDef.getOutputsMap()));
}
 

private static MetaGraphDef extractMetaGraphDefinition(final SavedModelBundle bundle)
    throws TensorflowMetaGraphDefParsingException {
  final MetaGraphDef metaGraphDef;
  try {
    metaGraphDef = MetaGraphDef.parseFrom(bundle.metaGraphDef());
  } catch (InvalidProtocolBufferException e) {
    throw new TensorflowMetaGraphDefParsingException(
        "Failed parsing tensorflow metagraph " + "definition", e);
  }

  return metaGraphDef;
}
 

public CollectionDef getCollectionDef(String key){
	MetaGraphDef metaGraphDef = getMetaGraphDef();

	Map<String, CollectionDef> collectionMap = metaGraphDef.getCollectionDefMap();

	CollectionDef collectionDef = collectionMap.get(key);
	if(collectionDef == null){
		throw new IllegalArgumentException(key);
	}

	return collectionDef;
}
 

static IntermediateGraph importGraph(String modelName, SavedModelBundle bundle) throws IOException {
    MetaGraphDef tfGraph = MetaGraphDef.parseFrom(bundle.metaGraphDef());

    IntermediateGraph intermediateGraph = new IntermediateGraph(modelName);
    importSignatures(tfGraph, intermediateGraph);
    importOperations(tfGraph, intermediateGraph, bundle);
    verifyOutputTypes(tfGraph, intermediateGraph);

    return intermediateGraph;
}
 

private static void importOperations(MetaGraphDef tfGraph,
                                     IntermediateGraph intermediateGraph,
                                     SavedModelBundle bundle) {
    for (String signatureName : intermediateGraph.signatures()) {
        for (String outputName : intermediateGraph.outputs(signatureName).values()) {
            importOperation(outputName, tfGraph.getGraphDef(), intermediateGraph, bundle);
        }
    }
}
 

private static void verifyOutputTypes(MetaGraphDef tfGraph, IntermediateGraph intermediateGraph) {
    for (String signatureName : intermediateGraph.signatures()) {
        for (String outputName : intermediateGraph.outputs(signatureName).values()) {
            IntermediateOperation operation = intermediateGraph.get(outputName);
            NodeDef node = getTensorFlowNodeFromGraph(IntermediateOperation.namePartOf(operation.name()), tfGraph.getGraphDef());
            OrderedTensorType type = operation.type().orElseThrow(
                    () -> new IllegalArgumentException("Output of '" + outputName + "' has no type."));
            TypeConverter.verifyType(node, type);
        }
    }

}
 

public MetaGraphDef getMetaGraphDef(){
	return this.metaGraphDef;
}
 

private void setMetaGraphDef(MetaGraphDef metaGraphDef){
	this.metaGraphDef = metaGraphDef;
}
 

public abstract MetaGraphDef metaGraphDefinition();