blob: f3d87d89c272ea129e7ccb06a9257e11cf91af07 (
plain) (
blame)
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
|
// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package ai.vespa.rankingexpression.importer.onnx;
import com.google.protobuf.ByteString;
import ai.vespa.rankingexpression.importer.OrderedTensorType;
import com.yahoo.tensor.IndexedTensor;
import com.yahoo.tensor.Tensor;
import onnx.Onnx;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.FloatBuffer;
/**
* Converts Onnx tensors into Vespa tensors.
*
* @author lesters
*/
class TensorConverter {
static Tensor toVespaTensor(Onnx.TensorProto tensorProto, OrderedTensorType type) {
Values values = readValuesOf(tensorProto);
IndexedTensor.BoundBuilder builder = (IndexedTensor.BoundBuilder) Tensor.Builder.of(type.type());
for (int i = 0; i < values.size(); i++) {
builder.cellByDirectIndex(type.toDirectIndex(i), values.get(i));
}
return builder.build();
}
private static Values readValuesOf(Onnx.TensorProto tensorProto) {
if (tensorProto.hasRawData()) {
switch (tensorProto.getDataType()) {
case FLOAT: return new RawFloatValues(tensorProto);
}
} else {
switch (tensorProto.getDataType()) {
case FLOAT: return new FloatValues(tensorProto);
}
}
throw new IllegalArgumentException("Cannot convert a tensor with elements of type " +
tensorProto.getDataType() + " to a Vespa tensor");
}
/** Allows reading values from buffers of various numeric types as bytes */
private static abstract class Values {
abstract double get(int i);
abstract int size();
}
private static abstract class RawValues extends Values {
ByteBuffer bytes(Onnx.TensorProto tensorProto) {
ByteString byteString = tensorProto.getRawData();
return byteString.asReadOnlyByteBuffer().order(ByteOrder.LITTLE_ENDIAN);
}
}
private static class RawFloatValues extends RawValues {
private final FloatBuffer values;
private final int size;
RawFloatValues(Onnx.TensorProto tensorProto) {
values = bytes(tensorProto).asFloatBuffer();
size = values.remaining();
}
@Override double get(int i) { return values.get(i); }
@Override int size() { return size; }
}
private static class FloatValues extends Values {
private final Onnx.TensorProto tensorProto;
FloatValues(Onnx.TensorProto tensorProto) {
this.tensorProto = tensorProto;
}
@Override double get(int i) { return tensorProto.getFloatData(i); }
@Override int size() { return tensorProto.getFloatDataCount(); }
}
}
|
