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
|
// See docs in ../ops/parse_ops.cc.
#include <algorithm>
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/public/tensor.h"
#include "tensorflow/core/public/tensor_shape.h"
namespace tensorflow {
template <typename T>
class DecodeRawOp : public OpKernel {
public:
explicit DecodeRawOp(OpKernelConstruction* context) : OpKernel(context) {
OP_REQUIRES_OK(context, context->GetAttr("little_endian", &little_endian_));
OP_REQUIRES_OK(context, context->GetAttr("out_type", &out_type_));
}
void Compute(OpKernelContext* context) override {
const auto& input = context->input(0);
int str_size = -1;
auto flat_in = input.flat<string>();
for (int i = 0; i < flat_in.size(); ++i) {
const string& in_str = flat_in(i);
if (str_size == -1) {
str_size = in_str.size();
} else {
OP_REQUIRES(context, str_size == in_str.size(),
errors::InvalidArgument(
"DecodeRaw requires input strings to all be the same "
"size, but element ",
i, " has size ", str_size, " != ", in_str.size()));
}
}
TensorShape out_shape = input.shape();
if (str_size == -1) { // Empty input
out_shape.AddDim(1);
Tensor* output_tensor = nullptr;
OP_REQUIRES_OK(context, context->allocate_output("output", out_shape,
&output_tensor));
return;
}
OP_REQUIRES(
context, str_size % sizeof(T) == 0,
errors::InvalidArgument("Input to DecodeRaw has length ", str_size,
" that is not a multiple of ", sizeof(T),
", the size of ", DataTypeString(out_type_)));
const int added_dim = str_size / sizeof(T);
out_shape.AddDim(added_dim);
Tensor* output_tensor = nullptr;
OP_REQUIRES_OK(
context, context->allocate_output("output", out_shape, &output_tensor));
auto out = output_tensor->flat_inner_dims<T>();
DCHECK_EQ(flat_in.size(), out.dimensions()[0]);
OP_REQUIRES(
context,
little_endian_ == ::tensorflow::port::kLittleEndian || sizeof(T) == 1,
errors::Unimplemented("Unimplemented support for little_endian=",
little_endian_ ? "true" : "false"));
// Endianness matches, so just copy each string byte-for-byte.
T* out_data = out.data();
for (int i = 0; i < flat_in.size(); ++i) {
const T* in_data = reinterpret_cast<const T*>(flat_in(i).data());
memcpy(out_data, in_data, str_size);
out_data += added_dim;
}
}
private:
bool little_endian_;
DataType out_type_;
};
#define REGISTER(type) \
REGISTER_KERNEL_BUILDER( \
Name("DecodeRaw").Device(DEVICE_CPU).TypeConstraint<type>("out_type"), \
DecodeRawOp<type>)
REGISTER(float);
REGISTER(double);
REGISTER(int32);
REGISTER(uint8);
REGISTER(int16);
REGISTER(int8);
REGISTER(int64);
#undef REGISTER
} // namespace tensorflow
|
