mirror of
https://github.com/wassname/Flex-Convolution.git
synced 2026-08-21 11:10:02 +08:00
113 lines
3.9 KiB
C++
113 lines
3.9 KiB
C++
/* Copyright 2017 ComputerGraphics Tuebingen. All Rights Reserved.
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
==============================================================================*/
|
|
// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
|
|
|
|
#include "flex_conv_op.h"
|
|
|
|
#include <stdio.h>
|
|
#include <type_traits>
|
|
|
|
#include "tensorflow/core/framework/op.h"
|
|
#include "tensorflow/core/framework/op_kernel.h"
|
|
#include "tensorflow/core/framework/register_types.h"
|
|
|
|
namespace tensorflow {
|
|
|
|
// Forward-Pass (CPU, GPU)
|
|
// --------------------------------------------------
|
|
template <typename Device, typename Dtype>
|
|
class FlexConvOp : public OpKernel {
|
|
public:
|
|
explicit FlexConvOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
|
|
|
|
void Compute(OpKernelContext* ctx) override {
|
|
// printf("--> Compute CPU Version <--\n");
|
|
const Tensor& features_ = ctx->input(0);
|
|
const Tensor& theta_ = ctx->input(1);
|
|
const Tensor& bias_ = ctx->input(2);
|
|
const Tensor& neighborhood_ = ctx->input(3);
|
|
const Tensor& positions_ = ctx->input(4);
|
|
|
|
const int B = neighborhood_.shape().dim_size(0);
|
|
const int N = neighborhood_.shape().dim_size(2);
|
|
const int Dout = theta_.shape().dim_size(3);
|
|
|
|
Tensor* output_ = nullptr;
|
|
OP_REQUIRES_OK(
|
|
ctx, ctx->allocate_output(0, TensorShape({B, Dout, N}), &output_));
|
|
|
|
::tensorflow::functor::FlexConvFunctor<Device, Dtype>()(
|
|
ctx, features_, theta_, bias_, neighborhood_, positions_, output_);
|
|
}
|
|
|
|
private:
|
|
TF_DISALLOW_COPY_AND_ASSIGN(FlexConvOp);
|
|
};
|
|
|
|
// Backward-Pass (CPU, GPU)
|
|
// --------------------------------------------------
|
|
template <typename Device, typename Dtype>
|
|
class FlexConvGradOp : public OpKernel {
|
|
public:
|
|
explicit FlexConvGradOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
|
|
|
|
void Compute(OpKernelContext* ctx) override {
|
|
// printf("--> Compute CPU Version <--\n");
|
|
const Tensor& features_ = ctx->input(0);
|
|
const Tensor& theta_ = ctx->input(1);
|
|
const Tensor& bias_ = ctx->input(2);
|
|
const Tensor& neighborhood_ = ctx->input(3);
|
|
const Tensor& positions_ = ctx->input(4);
|
|
|
|
const Tensor& topdiff_ = ctx->input(5);
|
|
|
|
// specify output shape
|
|
Tensor* grad_features_ = nullptr;
|
|
Tensor* grad_theta_ = nullptr;
|
|
Tensor* grad_bias_ = nullptr;
|
|
|
|
const int Degree = theta_.shape().dim_size(0);
|
|
|
|
OP_REQUIRES_OK(ctx,
|
|
ctx->allocate_output(0, features_.shape(), &grad_features_));
|
|
OP_REQUIRES_OK(ctx, ctx->allocate_output(1, theta_.shape(), &grad_theta_));
|
|
OP_REQUIRES_OK(ctx, ctx->allocate_output(2, bias_.shape(), &grad_bias_));
|
|
|
|
::tensorflow::functor::FlexConvGrad<Device, Dtype>()(
|
|
ctx, features_, theta_, bias_, neighborhood_, positions_, topdiff_,
|
|
grad_features_, grad_theta_, grad_bias_);
|
|
}
|
|
};
|
|
|
|
#define REGISTER_CUSTOM_OP(NAME, DEVICE, T) \
|
|
REGISTER_KERNEL_BUILDER( \
|
|
Name(#NAME).Device(DEVICE_##DEVICE).TypeConstraint<T>("T"), \
|
|
NAME##Op<DEVICE##Device, T>)
|
|
|
|
REGISTER_CUSTOM_OP(FlexConv, CPU, float);
|
|
REGISTER_CUSTOM_OP(FlexConvGrad, CPU, float);
|
|
REGISTER_CUSTOM_OP(FlexConv, CPU, double);
|
|
REGISTER_CUSTOM_OP(FlexConvGrad, CPU, double);
|
|
|
|
#ifdef GOOGLE_CUDA
|
|
REGISTER_CUSTOM_OP(FlexConv, GPU, float);
|
|
REGISTER_CUSTOM_OP(FlexConvGrad, GPU, float);
|
|
REGISTER_CUSTOM_OP(FlexConv, GPU, double);
|
|
REGISTER_CUSTOM_OP(FlexConvGrad, GPU, double);
|
|
#endif // GOOGLE_CUDA
|
|
#undef REGISTER_CUSTOM_OP
|
|
|
|
} // namespace tensorflow
|