1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
|
// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package com.yahoo.searchlib.rankingexpression.integration.tensorflow;
import com.yahoo.searchlib.rankingexpression.RankingExpression;
import com.yahoo.searchlib.rankingexpression.parser.ParseException;
import com.yahoo.tensor.TensorType;
import com.yahoo.tensor.functions.ScalarFunctions;
import com.yahoo.yolean.Exceptions;
import org.tensorflow.SavedModelBundle;
import org.tensorflow.framework.GraphDef;
import org.tensorflow.framework.MetaGraphDef;
import org.tensorflow.framework.NodeDef;
import org.tensorflow.framework.SignatureDef;
import org.tensorflow.framework.TensorInfo;
import org.tensorflow.framework.TensorShapeProto;
import java.io.File;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
/**
* Converts a saved TensorFlow model into a ranking expression and set of constants.
*
* @author bratseth
*/
public class TensorFlowImporter {
private final OperationMapper operationMapper = new OperationMapper();
/**
* Imports a saved TensorFlow model from a directory.
* The model should be saved as a .pbtxt or .pb file.
* The name of the model is taken as the db/pbtxt file name (not including the file ending).
*
* @param modelDir the directory containing the TensorFlow model files to import
*/
public TensorFlowModel importModel(String modelDir) {
try (SavedModelBundle model = SavedModelBundle.load(modelDir, "serve")) {
return importModel(model);
}
catch (IllegalArgumentException e) {
throw new IllegalArgumentException("Could not import TensorFlow model from directory '" + modelDir + "'", e);
}
}
public TensorFlowModel importModel(File modelDir) {
return importModel(modelDir.toString());
}
/** Imports a TensorFlow model */
public TensorFlowModel importModel(SavedModelBundle model) {
try {
return importGraph(MetaGraphDef.parseFrom(model.metaGraphDef()), model);
}
catch (IOException e) {
throw new IllegalArgumentException("Could not import TensorFlow model '" + model + "'", e);
}
}
private TensorFlowModel importGraph(MetaGraphDef graph, SavedModelBundle model) {
TensorFlowModel result = new TensorFlowModel();
for (Map.Entry<String, SignatureDef> signatureEntry : graph.getSignatureDefMap().entrySet()) {
TensorFlowModel.Signature signature = result.signature(signatureEntry.getKey()); // Prefer key over "methodName"
importInputs(signatureEntry.getValue().getInputsMap(), signature);
for (Map.Entry<String, TensorInfo> output : signatureEntry.getValue().getOutputsMap().entrySet()) {
String outputName = output.getKey();
try {
NodeDef node = getNode(nameOf(output.getValue().getName()), graph.getGraphDef());
importNode(node, graph.getGraphDef(), model, result);
signature.output(outputName, nameOf(output.getValue().getName()));
}
catch (IllegalArgumentException e) {
signature.skippedOutput(outputName, Exceptions.toMessageString(e));
}
}
}
return result;
}
private void importInputs(Map<String, TensorInfo> inputInfoMap, TensorFlowModel.Signature signature) {
inputInfoMap.forEach((key, value) -> {
String argumentName = nameOf(value.getName());
TensorType argumentType = importTensorType(value.getTensorShape());
// Arguments are (Placeholder) nodes, so not local to the signature:
signature.owner().argument(argumentName, argumentType);
signature.input(key, argumentName);
});
}
private TensorType importTensorType(TensorShapeProto tensorShape) {
TensorType.Builder b = new TensorType.Builder();
for (TensorShapeProto.Dim dimension : tensorShape.getDimList()) {
int dimensionSize = (int)dimension.getSize();
if (dimensionSize >= 0)
b.indexed("d" + b.rank(), dimensionSize);
else
b.indexed("d" + b.rank()); // unbound size
}
return b.build();
}
/** Recursively convert a graph of TensorFlow nodes into a Vespa tensor function expression tree */
private TypedTensorFunction importNode(NodeDef tfNode, GraphDef graph, SavedModelBundle model, TensorFlowModel result) {
TypedTensorFunction function = tensorFunctionOf(tfNode, graph, model, result);
try {
// We add all intermediate nodes imported as separate expressions. Only those referenced in a signature output
// will be used. We parse the TensorFunction here to convert it to a RankingExpression tree
result.expression(tfNode.getName(), new RankingExpression(tfNode.getName(), function.function().toString()));
return function;
}
catch (ParseException e) {
throw new RuntimeException("Tensorflow function " + function.function() +
" cannot be parsed as a ranking expression", e);
}
}
private TypedTensorFunction tensorFunctionOf(NodeDef tfNode, GraphDef graph, SavedModelBundle model, TensorFlowModel result) {
// Import arguments lazily below, as some nodes have arguments unused arguments leading to unsupported ops
// TODO: Implement mapping of more functions from https://www.tensorflow.org/api_docs/python/
switch (tfNode.getOp().toLowerCase()) {
// array ops
case "const" : return operationMapper.constant(tfNode, model, result);
case "expanddims" : return operationMapper.expandDims(tfNode, model, importArguments(tfNode, graph, model, result));
case "identity" : return operationMapper.identity(tfNode, model, result);
case "placeholder" : return operationMapper.placeholder(tfNode, result);
case "placeholderwithdefault" : return operationMapper.placeholderWithDefault(tfNode, model, result);
case "reshape" : return operationMapper.reshape(tfNode, model, importArguments(tfNode, graph, model, result));
case "squeeze" : return operationMapper.squeeze(tfNode, importArguments(tfNode, graph, model, result));
// math ops
case "add" : case "add_n" : return operationMapper.join(importArguments(tfNode, graph, model, result), ScalarFunctions.add());
case "acos" : return operationMapper.map(importArguments(tfNode, graph, model, result), ScalarFunctions.acos());
case "matmul" : return operationMapper.matmul(importArguments(tfNode, graph, model, result));
case "maximum" : return operationMapper.join(importArguments(tfNode, graph, model, result), ScalarFunctions.max());
case "mean" : case "reducemean": return operationMapper.mean(tfNode, model, importArguments(tfNode, graph, model, result));
case "multiply": case "mul" : return operationMapper.join(importArguments(tfNode, graph, model, result), ScalarFunctions.multiply());
case "rsqrt": return operationMapper.map(importArguments(tfNode, graph, model, result), ScalarFunctions.rsqrt());
case "where3": case "select" : return operationMapper.select(tfNode, model, result, importArguments(tfNode, graph, model, result));
case "sigmoid": return operationMapper.map(importArguments(tfNode, graph, model, result), ScalarFunctions.sigmoid());
case "squareddifference" : return operationMapper.join(importArguments(tfNode, graph, model, result), ScalarFunctions.squareddifference());
case "subtract" : case "sub" : return operationMapper.join(importArguments(tfNode, graph, model, result), ScalarFunctions.subtract());
// nn ops
case "biasadd" : return operationMapper.join(importArguments(tfNode, graph, model, result), ScalarFunctions.add());
case "elu": return operationMapper.map(importArguments(tfNode, graph, model, result), ScalarFunctions.elu());
case "relu": return operationMapper.map(importArguments(tfNode, graph, model, result), ScalarFunctions.relu());
case "selu": return operationMapper.map(importArguments(tfNode, graph, model, result), ScalarFunctions.selu());
case "softmax" : return operationMapper.softmax(importArguments(tfNode, graph, model, result));
// evaluation no-ops
case "stopgradient" :
case "noop":
return operationMapper.noOp(importArguments(tfNode, graph, model, result));
// not supported
default :
throw new IllegalArgumentException("Conversion of TensorFlow operation '" + tfNode.getOp() + "' is not supported (" + tfNode.getName() + ")");
}
}
private List<TypedTensorFunction> importArguments(NodeDef tfNode, GraphDef graph, SavedModelBundle model,
TensorFlowModel result) {
return tfNode.getInputList().stream()
.map(argNode -> importNode(getNode(nameOf(argNode), graph), graph, model, result))
.collect(Collectors.toList());
}
private NodeDef getNode(String name, GraphDef graph) {
return graph.getNodeList().stream()
.filter(node -> node.getName().equals(name))
.findFirst()
.orElseThrow(() -> new IllegalArgumentException("Could not find node '" + name + "'"));
}
/**
* A method signature input and output has the form name:index.
* This returns the name part without the index.
*/
private String nameOf(String name) {
return name.split(":")[0];
}
}
|
