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
|
#include "tensorflow/core/util/saved_tensor_slice_util.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/strings/ordered_code.h"
namespace tensorflow {
namespace checkpoint {
const char kSavedTensorSlicesKey[] = "";
string EncodeTensorNameSlice(const string& name, const TensorSlice& slice) {
string buffer;
// All the tensor slice keys will start with a 0
tensorflow::strings::OrderedCode::WriteNumIncreasing(&buffer, 0);
tensorflow::strings::OrderedCode::WriteString(&buffer, name);
tensorflow::strings::OrderedCode::WriteNumIncreasing(&buffer, slice.dims());
for (int d = 0; d < slice.dims(); ++d) {
// A trivial extent (meaning we take EVERYTHING) will default to -1 for both
// start and end. These will be properly parsed.
tensorflow::strings::OrderedCode::WriteSignedNumIncreasing(&buffer,
slice.start(d));
tensorflow::strings::OrderedCode::WriteSignedNumIncreasing(&buffer,
slice.length(d));
}
return buffer;
}
Status DecodeTensorNameSlice(const string& code, string* name,
tensorflow::TensorSlice* slice) {
StringPiece src(code);
uint64 x;
if (!tensorflow::strings::OrderedCode::ReadNumIncreasing(&src, &x)) {
return errors::Internal("Failed to parse the leading number: src = ", src);
}
if (x != 0) {
return errors::Internal(
"The leading number should always be 0 for any valid key: src = ", src);
}
if (!tensorflow::strings::OrderedCode::ReadString(&src, name)) {
return errors::Internal("Failed to parse the tensor name: src = ", src);
}
if (!tensorflow::strings::OrderedCode::ReadNumIncreasing(&src, &x)) {
return errors::Internal("Failed to parse the tensor rank: src = ", src);
}
if (x == 0) {
return errors::Internal("Expecting positive rank of the tensor, got ", x,
", src = ", src);
}
if (x >= kint32max) {
return errors::Internal("Too many elements ", x);
}
slice->SetFullSlice(x);
for (int d = 0; d < static_cast<int32>(x); ++d) {
// We expected 2x integers
int64 start, length;
if (!tensorflow::strings::OrderedCode::ReadSignedNumIncreasing(&src,
&start)) {
return errors::Internal("Failed to parse start: src = ", src);
}
if (!tensorflow::strings::OrderedCode::ReadSignedNumIncreasing(&src,
&length)) {
return errors::Internal("Failed to parse length: src = ", src);
}
if (length >= 0) {
// a non-trivial extent
slice->set_start(d, start);
slice->set_length(d, length);
}
}
return Status::OK();
}
} // namespace checkpoint
} // namespace tensorflow
|
