Tensorflow custom op - how can I read and write from tensors?
I am writing a custom Tensorflow op using a tutorial and I am having trouble understanding how to read and write to / from tensors.
let's say I have a tensor in my OpKernel which I am getting from
const Tensor& values_tensor = context->input(0);
(where context = OpKernelConstruction*
)
if this Tensor has the shape of, say, [2, 10, 20], how can I index it (eg, auto x = values_tensor[1, 4, 12]
etc.)?
equivalent if i
Tensor *output_tensor = NULL;
OP_REQUIRES_OK(context, context->allocate_output(
0,
{batch_size, value_len - window_size, window_size},
&output_tensor
));
how can I assign output_tensor
, for example output_tensor[1, 2, 3] = 11
, etc.?
sorry for the dumb question, but the docs really got me off here, and the examples in Tensorflow core code for embedded operating systems somehow obfuscate this to the point that I'm very confused :)
Thank you!
source to share
The easiest way to read and write objects tensorflow::Tensor
is to convert them to Native Tensor using the tensorflow::Tensor::tensor<T, NDIMS>()
. Note that you must specify the type (C ++) of the elements in the tensor as a template parameter T
.
For example, to read a specific value from a tensor DT_FLOAT32
:
const Tensor& values_tensor = context->input(0);
auto x = value_tensor.tensor<float, 3>()(1, 4, 12);
To write a specific value to a tensor DT_FLOAT32
:
Tensor* output_tensor = ...;
output_tensor->tensor<float, 3>()(1, 2, 3) = 11.0;
There are also convenience methods for accessing scalar , vector , or matrix .
source to share