| #define EIGEN_USE_THREADS |
|
|
| #include <memory> |
| #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" |
| #include "tensorflow/core/framework/op_kernel.h" |
| #include "tensorflow/core/framework/register_types.h" |
| #include "tensorflow/core/framework/tensor.h" |
| #include "tensorflow/core/framework/tensor_shape.h" |
| #include "tensorflow/core/framework/types.h" |
| #include "tensorflow/core/lib/core/status.h" |
| #include "tensorflow/core/platform/logging.h" |
| #include "tensorflow/core/framework/op.h" |
| #include "tensorflow/core/framework/shape_inference.h" |
| #include "tensorflow/core/framework/common_shape_fns.h" |
|
|
| using namespace tensorflow; |
|
|
| using CPUDevice = Eigen::ThreadPoolDevice; |
| using GPUDevice = Eigen::GpuDevice; |
|
|
| void Downsample(const GPUDevice& d, |
| typename TTypes<float, 4>::ConstTensor images, |
| typename TTypes<float, 4>::Tensor output); |
|
|
| class DownsampleOp : public OpKernel { |
| public: |
| explicit DownsampleOp(OpKernelConstruction* c) : OpKernel(c) { |
| OP_REQUIRES_OK(c, c->GetAttr("scale", &scale)); |
| } |
|
|
| void Compute(OpKernelContext* context) override { |
| const Tensor& input = context->input(0); |
|
|
| typename TTypes<float, 4>::ConstTensor input_data = input.tensor<float, 4>(); |
|
|
|
|
| OP_REQUIRES(context, |
| input_data.dimension(1) % scale == 0 && |
| input_data.dimension(2) % scale == 0, |
| errors::InvalidArgument("Input height and width must be divisible by scale")); |
|
|
| const int batch = input_data.dimension(0); |
| const int height = input_data.dimension(1) / scale; |
| const int width = input_data.dimension(2) / scale; |
| const int channels = input_data.dimension(3); |
|
|
| auto output_shape = TensorShape({batch, height, width, channels}); |
|
|
| Tensor* output = NULL; |
| OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output)); |
|
|
| typename TTypes<float, 4>::Tensor output_data = output->tensor<float, 4>(); |
|
|
| Downsample(context->eigen_device<GPUDevice>(), input_data, output_data); |
| } |
| private: |
| int scale; |
| }; |
|
|
| using shape_inference::DimensionHandle; |
| using shape_inference::ShapeHandle; |
|
|
| REGISTER_OP("Downsample") |
| .Input("images: float") |
| .Attr("scale: int = 2") |
| .Output("out_images: float") |
| .SetShapeFn([](shape_inference::InferenceContext* c) { |
| ShapeHandle in = c->input(0); |
| int scale; |
| DimensionHandle batch = c->Dim(in, 0); |
| DimensionHandle channels = c->Dim(in, 3); |
| DimensionHandle height; |
| DimensionHandle width; |
|
|
| c->GetAttr("scale", &scale); |
| c->Divide(c->Dim(in, 1), scale, true, &height); |
| c->Divide(c->Dim(in, 2), scale, true, &width); |
|
|
| c->set_output(0, c->MakeShape({batch, height, width, channels})); |
| return Status::OK(); |
| }); |
|
|
| #if GOOGLE_CUDA |
|
|
| REGISTER_KERNEL_BUILDER(Name("Downsample").Device(DEVICE_GPU), DownsampleOp); |
|
|
| #endif |
|
|
