Expect input tensor of rank 4 with dimensions (batch_size, height, width, depth).
| 64 | // Expect input tensor of rank 4 with dimensions (batch_size, height, width, |
| 65 | // depth). |
| 66 | void Compute(OpKernelContext* context) override { |
| 67 | const Tensor& input = context->input(0); |
| 68 | const TensorShape& input_shape = input.shape(); |
| 69 | const int32 num_dims = input_shape.dims(); |
| 70 | OP_REQUIRES( |
| 71 | context, num_dims == 4, |
| 72 | errors::InvalidArgument( |
| 73 | "input must be 4-dimensional (batch_size, height, width, depth)", |
| 74 | input_shape.DebugString())); |
| 75 | |
| 76 | const int64 batch_size = input_shape.dim_size(0); |
| 77 | |
| 78 | const Tensor& window_size = context->input(1); |
| 79 | OP_REQUIRES(context, |
| 80 | (window_size.shape().dims() == 1) && |
| 81 | window_size.shape().dim_size(0) == 2, |
| 82 | errors::InvalidArgument( |
| 83 | "input must be a vector of size 2 (height, width)", |
| 84 | window_size.shape().DebugString())); |
| 85 | |
| 86 | const int64 output_height = window_size.tensor<int, 1>()(0); |
| 87 | const int64 output_width = window_size.tensor<int, 1>()(1); |
| 88 | TensorShape output_shape = input_shape; |
| 89 | OP_REQUIRES_OK(context, output_shape.SetDimWithStatus(1, output_height)); |
| 90 | OP_REQUIRES_OK(context, output_shape.SetDimWithStatus(2, output_width)); |
| 91 | |
| 92 | const Tensor& offsets = context->input(2); |
| 93 | OP_REQUIRES(context, offsets.shape().dims() == 2, |
| 94 | errors::InvalidArgument("input must be a matrix", |
| 95 | offsets.shape().DebugString())); |
| 96 | OP_REQUIRES(context, offsets.shape().dim_size(0) == batch_size, |
| 97 | errors::InvalidArgument("first dimension should be batch", |
| 98 | offsets.shape().DebugString())); |
| 99 | OP_REQUIRES( |
| 100 | context, offsets.shape().dim_size(1) == 2, |
| 101 | errors::InvalidArgument("second dimension should be of size 2 (y,x)", |
| 102 | offsets.shape().DebugString())); |
| 103 | |
| 104 | Tensor* output = nullptr; |
| 105 | OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output)); |
| 106 | if (output->NumElements() == 0) { |
| 107 | // Nothing else to do. |
| 108 | return; |
| 109 | } |
| 110 | |
| 111 | std::vector<Eigen::IndexPair<float> > offset_vec; |
| 112 | offset_vec.reserve(batch_size); |
| 113 | for (int i = 0; i < batch_size; ++i) { |
| 114 | float offset_y = offsets.tensor<float, 2>()(i, 0); |
| 115 | float offset_x = offsets.tensor<float, 2>()(i, 1); |
| 116 | // Eigen::ExtractGlimpses expects offsets as (x,y), whereas the |
| 117 | // calling TensorFlow operates with (y,x) as indices. |
| 118 | offset_vec.push_back(Eigen::IndexPair<float>(offset_x, offset_y)); |
| 119 | } |
| 120 | |
| 121 | output->tensor<float, 4>().swap_layout().device( |
| 122 | context->eigen_cpu_device()) = |
| 123 | Eigen::ExtractGlimpses(input.tensor<float, 4>().swap_layout(), |
nothing calls this directly
no test coverage detected