File size: 2,821 Bytes
872b0a0 | 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 | #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 // GOOGLE_CUDA
|