mirror of
https://github.com/wassname/Flex-Convolution.git
synced 2026-09-12 12:01:02 +08:00
cherry pick files for public release
This commit is contained in:
Patrick Wieschollek
@@ -0,0 +1,173 @@
|
||||
/* 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 "tensorflow/core/framework/op.h"
|
||||
|
||||
namespace tensorflow {
|
||||
|
||||
namespace functor {
|
||||
|
||||
template <typename Dtype>
|
||||
struct FlexConvFunctor<CPUDevice, Dtype> {
|
||||
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& features_,
|
||||
const Tensor& theta_, const Tensor& bias_,
|
||||
const Tensor& neighborhood_, const Tensor& positions_,
|
||||
Tensor* output_) {
|
||||
const auto features = features_.tensor<Dtype, 3>();
|
||||
const auto theta = theta_.tensor<Dtype, 4>();
|
||||
const auto bias = bias_.tensor<Dtype, 2>();
|
||||
const auto neighborhood = neighborhood_.tensor<int, 3>();
|
||||
const auto positions = positions_.tensor<Dtype, 3>();
|
||||
|
||||
auto output = output_->tensor<Dtype, 3>();
|
||||
|
||||
// get dimensions
|
||||
const int B = neighborhood_.dim_size(0);
|
||||
const int K = neighborhood_.dim_size(1);
|
||||
const int N = neighborhood_.dim_size(2);
|
||||
const int Dp = theta_.dim_size(1);
|
||||
const int Din = theta_.dim_size(2);
|
||||
const int Dout = theta_.dim_size(3);
|
||||
|
||||
output.setZero();
|
||||
|
||||
for (int b = 0; b < B; ++b) {
|
||||
for (int n = 0; n < N; ++n) {
|
||||
for (int k_ = 0; k_ < K; ++k_) {
|
||||
int k = neighborhood(b, k_, n);
|
||||
|
||||
for (int dout = 0; dout < Dout; ++dout) {
|
||||
for (int din = 0; din < Din; ++din) {
|
||||
const Dtype v = features(b, din, k);
|
||||
|
||||
Dtype W = bias(din, dout);
|
||||
for (int dp = 0; dp < Dp; ++dp) {
|
||||
Dtype delta = positions(b, dp, k) -
|
||||
positions(b, dp, neighborhood(b, 0, n));
|
||||
W += theta(0, dp, din, dout) * delta;
|
||||
}
|
||||
output(b, dout, n) = output(b, dout, n) + W * v;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template struct FlexConvFunctor<CPUDevice, float>;
|
||||
|
||||
template <typename Dtype>
|
||||
struct FlexConvGrad<CPUDevice, Dtype> {
|
||||
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& features_,
|
||||
const Tensor& theta_, const Tensor& bias_,
|
||||
const Tensor& neighborhood_, const Tensor& positions_,
|
||||
const Tensor& topdiff_, Tensor* grad_features_,
|
||||
Tensor* grad_theta_, Tensor* grad_bias_) {
|
||||
const auto features = features_.tensor<Dtype, 3>();
|
||||
const auto theta = theta_.tensor<Dtype, 4>();
|
||||
const auto bias = bias_.tensor<Dtype, 2>();
|
||||
const auto neighborhood = neighborhood_.tensor<int, 3>();
|
||||
const auto positions = positions_.tensor<Dtype, 3>();
|
||||
const auto topdiff = topdiff_.tensor<Dtype, 3>();
|
||||
|
||||
auto grad_features = grad_features_->tensor<Dtype, 3>();
|
||||
auto grad_theta = grad_theta_->tensor<Dtype, 4>();
|
||||
auto grad_bias = grad_bias_->tensor<Dtype, 2>();
|
||||
|
||||
// get dimensions
|
||||
const int B = neighborhood_.dim_size(0);
|
||||
const int K = neighborhood_.dim_size(1);
|
||||
const int N = neighborhood_.dim_size(2);
|
||||
const int Ddegree = theta_.dim_size(0);
|
||||
const int Dp = theta_.dim_size(1);
|
||||
const int Din = theta_.dim_size(2);
|
||||
const int Dout = theta_.dim_size(3);
|
||||
|
||||
grad_features.setZero();
|
||||
grad_theta.setZero();
|
||||
grad_bias.setZero();
|
||||
|
||||
// ========================= bias ==============================
|
||||
for (int b = 0; b < B; ++b) {
|
||||
for (int n = 0; n < N; ++n) {
|
||||
for (int k_ = 0; k_ < K; ++k_) {
|
||||
int k = neighborhood(b, k_, n);
|
||||
|
||||
for (int j = 0; j < Din; ++j) {
|
||||
for (int l = 0; l < Dout; ++l) {
|
||||
grad_bias(j, l) += features(b, j, k) * topdiff(b, l, n);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ========================= theta ==============================
|
||||
for (int b = 0; b < B; ++b) {
|
||||
for (int n = 0; n < N; ++n) {
|
||||
for (int k_ = 0; k_ < K; ++k_) {
|
||||
int k = neighborhood(b, k_, n);
|
||||
|
||||
for (int j = 0; j < Din; ++j) {
|
||||
for (int l = 0; l < Dout; ++l) {
|
||||
for (int i = 0; i < Dp; ++i) {
|
||||
const Dtype delta =
|
||||
positions(b, i, k) - positions(b, i, neighborhood(b, 0, n));
|
||||
// printf("delta %f\n", delta);
|
||||
|
||||
for (int dd = 0; dd < Ddegree; ++dd) {
|
||||
grad_theta(dd, i, j, l) +=
|
||||
features(b, j, k) * pow(delta, dd + 1) * topdiff(b, l, n);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ========================= features ==============================
|
||||
for (int b = 0; b < B; ++b) {
|
||||
for (int n = 0; n < N; ++n) {
|
||||
for (int k_ = 0; k_ < K; ++k_) {
|
||||
int k = neighborhood(b, k_, n);
|
||||
|
||||
for (int j = 0; j < Din; ++j) {
|
||||
for (int l = 0; l < Dout; ++l) {
|
||||
Dtype W = bias(j, l);
|
||||
for (int i = 0; i < Dp; ++i) {
|
||||
const Dtype delta =
|
||||
positions(b, i, k) - positions(b, i, neighborhood(b, 0, n));
|
||||
for (int dd = 0; dd < Ddegree; ++dd)
|
||||
W += theta(dd, i, j, l) * pow(delta, dd + 1);
|
||||
}
|
||||
grad_features(b, j, k) += W * topdiff(b, l, n);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// template struct FlexConvGrad<CPUDevice, int>;
|
||||
template struct FlexConvGrad<CPUDevice, float>;
|
||||
// template struct FlexConvGrad<CPUDevice, double>;
|
||||
|
||||
} // namespace functor
|
||||
} // namespace tensorflow
|
||||
@@ -0,0 +1,532 @@
|
||||
/* 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
|
||||
|
||||
#if GOOGLE_CUDA
|
||||
|
||||
#define EIGEN_USE_GPU
|
||||
|
||||
#include <cub/cub.cuh>
|
||||
|
||||
#include "flex_conv_op.h"
|
||||
#include "tensorflow/core/util/cuda_kernel_helper.h"
|
||||
|
||||
namespace FlexConvCuda {
|
||||
|
||||
using CudaLaunchConfig = ::tensorflow::CudaLaunchConfig;
|
||||
|
||||
constexpr __host__ __device__ int pmin(int x, int y) { return x <= y ? x : y; }
|
||||
|
||||
template <typename Dtype, typename NBtype, int Dp = 3, int C_N = 256,
|
||||
int C_Dout = 32, int C_Din = 64>
|
||||
struct ForwardKernel;
|
||||
|
||||
template <typename Dtype, typename NBtype, int Dp, int C_N, int C_Dout,
|
||||
int C_Din>
|
||||
__global__ void runForwardKernel(
|
||||
const ForwardKernel<Dtype, NBtype, Dp, C_N, C_Dout, C_Din> kernel) {
|
||||
kernel();
|
||||
}
|
||||
|
||||
template <typename Dtype, typename NBtype, int Dp, int C_N, int C_Dout,
|
||||
int C_Din>
|
||||
struct ForwardKernel {
|
||||
enum {
|
||||
PMIN = 3 // only for unrolling
|
||||
};
|
||||
|
||||
void launch(int B) {
|
||||
dim3 block(C_N);
|
||||
dim3 grid((N - 1) / C_N + 1, (Dout - 1) / C_Dout + 1, B);
|
||||
|
||||
size_t shm_size = (Dp + 1) * C_Din * C_Dout * sizeof(Dtype);
|
||||
|
||||
runForwardKernel<<<grid, block, shm_size>>>((*this));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void operator()() const {
|
||||
extern __shared__ Dtype s_shm[];
|
||||
|
||||
Dtype* s_theta = (float*)&s_shm[0];
|
||||
Dtype* s_bias = (float*)&s_shm[Dp * C_Din * C_Dout];
|
||||
|
||||
// glob ids
|
||||
int b = blockIdx.z;
|
||||
int n = blockIdx.x * C_N + threadIdx.x;
|
||||
|
||||
Dtype result[C_Dout];
|
||||
for (int dout = 0; dout < C_Dout; ++dout) {
|
||||
result[dout] = 0.0;
|
||||
}
|
||||
|
||||
Dtype p0[Dp];
|
||||
#pragma unroll pmin(Dp, PMIN)
|
||||
for (int dp = 0; dp < Dp && n < N; ++dp) {
|
||||
p0[dp] = d_positions[b * Dp * N + dp * N + n];
|
||||
}
|
||||
|
||||
for (int o_din = 0; o_din < Din; o_din += C_Din) {
|
||||
// load shm
|
||||
__syncthreads();
|
||||
for (int tid = threadIdx.x; tid < Dp * C_Din * C_Dout; tid += C_N) {
|
||||
int dp = tid / (C_Din * C_Dout);
|
||||
int din = (tid % (C_Din * C_Dout)) / C_Dout;
|
||||
int dout = tid % C_Dout;
|
||||
|
||||
int g_dout = (dout + blockIdx.y * C_Dout);
|
||||
int g_din = o_din + din;
|
||||
|
||||
if (g_dout < Dout && g_din < Din) {
|
||||
s_theta[dp * C_Din * C_Dout + din * C_Dout + dout] =
|
||||
d_theta[dp * Din * Dout + g_din * Dout + g_dout];
|
||||
|
||||
if (!dp) s_bias[din * C_Dout + dout] = d_bias[g_din * Dout + g_dout];
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (n < N) {
|
||||
// Loop over K
|
||||
for (int k = 0; k < K && n < N; ++k) {
|
||||
NBtype nk = d_neighborhood[b * K * N + k * N + n];
|
||||
|
||||
Dtype q[Dp];
|
||||
#pragma unroll pmin(Dp, PMIN)
|
||||
for (int dp = 0; dp < Dp; ++dp) {
|
||||
q[dp] = d_positions[b * Dp * N + dp * N + nk] - p0[dp];
|
||||
}
|
||||
|
||||
// Loop over Din
|
||||
for (int din = 0; din < C_Din && (o_din + din) < Din; ++din) {
|
||||
Dtype fk = d_features[b * Din * N + (o_din + din) * N + nk];
|
||||
|
||||
// Loop over partial Dout
|
||||
for (int dout = 0;
|
||||
dout < C_Dout && (dout + blockIdx.y * C_Dout) < Dout; ++dout) {
|
||||
Dtype w = 0.0;
|
||||
|
||||
for (int dp = 0; dp < Dp; ++dp)
|
||||
w += q[dp] * s_theta[dp * C_Din * C_Dout + din * C_Dout + dout];
|
||||
w += s_bias[din * C_Dout + dout];
|
||||
result[dout] += w * fk;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (int dout = 0;
|
||||
dout < C_Dout && (dout + blockIdx.y * C_Dout) < Dout && n < N;
|
||||
++dout) {
|
||||
d_output[b * Dout * N + (dout + blockIdx.y * C_Dout) * N + n] =
|
||||
result[dout];
|
||||
}
|
||||
}
|
||||
|
||||
// features: incoming features [B, Din, N].
|
||||
// position: each datapoint in nd space [B, Dp, N].
|
||||
// neighborhood: all K nearest neighbors [B, K, N].
|
||||
const Dtype* d_features;
|
||||
const Dtype* d_positions;
|
||||
const NBtype* d_neighborhood;
|
||||
|
||||
// theta: parameters for kernel function [Dp,
|
||||
// Din, Dout]. bias: parameters for kernel function [Din, Dout].
|
||||
const Dtype* d_theta;
|
||||
const Dtype* d_bias;
|
||||
|
||||
// output: each feature description for each point [B, Dout, N].
|
||||
Dtype* d_output;
|
||||
|
||||
int N;
|
||||
int K;
|
||||
int Din;
|
||||
int Dout;
|
||||
};
|
||||
|
||||
template <typename Dtype>
|
||||
struct BackwardThetaKernel;
|
||||
|
||||
template <typename T>
|
||||
__global__ void runBackwardKernel(const BackwardThetaKernel<T> kernel) {
|
||||
kernel();
|
||||
}
|
||||
|
||||
template <typename Dtype>
|
||||
struct BackwardThetaKernel {
|
||||
enum { C_N = 256, DP_MAX = 3, DEGREE_MAX = 2 };
|
||||
|
||||
void launch() {
|
||||
dim3 block(C_N);
|
||||
dim3 grid(Dout, Din);
|
||||
|
||||
runBackwardKernel<<<grid, block>>>((*this));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void operator()() const {
|
||||
typedef cub::BlockReduce<Dtype, C_N> BlockReduce;
|
||||
__shared__ typename BlockReduce::TempStorage temp_storage;
|
||||
|
||||
Dtype theta_diff[DP_MAX];
|
||||
for (int dp = 0; dp < Dp; ++dp) theta_diff[dp] = 0;
|
||||
|
||||
Dtype bias_diff = 0;
|
||||
|
||||
int dout = blockIdx.x;
|
||||
int din = blockIdx.y;
|
||||
|
||||
for (int b = 0; b < B; ++b) {
|
||||
for (int n = threadIdx.x; n < N; n += C_N) {
|
||||
Dtype topdiff = d_topdiff[b * Dout * N + dout * N + n];
|
||||
|
||||
for (int k = 0; k < K; ++k) {
|
||||
int nk0 = d_neigh[b * N * K + 0 * N + n];
|
||||
int nk = d_neigh[b * N * K + k * N + n];
|
||||
|
||||
Dtype feature = d_features[b * Din * N + din * N + nk];
|
||||
for (int dp = 0; dp < Dp; ++dp) {
|
||||
Dtype diffpos = d_pos[b * Dp * N + dp * N + nk] -
|
||||
d_pos[b * Dp * N + dp * N + nk0];
|
||||
theta_diff[dp] += feature * diffpos * topdiff;
|
||||
}
|
||||
bias_diff += feature * topdiff;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (int dp = 0; dp < Dp; ++dp) {
|
||||
// for (int dd = 0; dd < Ddegree; ++dd) {
|
||||
Dtype thread_data = theta_diff[dp];
|
||||
Dtype aggregate = BlockReduce(temp_storage).Sum(thread_data, N);
|
||||
|
||||
if (!threadIdx.x) {
|
||||
d_theta_out[dp * Din * Dout + din * Dout + dout] = aggregate;
|
||||
}
|
||||
// }
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
Dtype thread_data = bias_diff;
|
||||
|
||||
Dtype aggregate = BlockReduce(temp_storage).Sum(thread_data, N);
|
||||
|
||||
if (!threadIdx.x) d_bias_out[din * Dout + dout] = aggregate;
|
||||
}
|
||||
|
||||
const Dtype* d_topdiff;
|
||||
|
||||
const Dtype* d_pos;
|
||||
const Dtype* d_features;
|
||||
const int* d_neigh;
|
||||
|
||||
const Dtype* d_theta;
|
||||
const Dtype* d_bias;
|
||||
|
||||
Dtype* d_theta_out;
|
||||
Dtype* d_bias_out;
|
||||
|
||||
int B;
|
||||
int N;
|
||||
int K;
|
||||
int Ddegree;
|
||||
int Dp;
|
||||
int Din;
|
||||
int Dout;
|
||||
};
|
||||
|
||||
template <typename Dtype>
|
||||
struct BackwardFeatureKernel;
|
||||
|
||||
template <typename T>
|
||||
__global__ void runBackwardKernel(const BackwardFeatureKernel<T> kernel) {
|
||||
kernel();
|
||||
}
|
||||
|
||||
template <typename Dtype>
|
||||
struct BackwardFeatureKernel {
|
||||
enum {
|
||||
C_N = 32,
|
||||
C_Dout = 32, // multiple of Warpsize is better
|
||||
|
||||
C_Din = 8 // reduce first
|
||||
};
|
||||
|
||||
void launch(int B) {
|
||||
dim3 fblock(C_N, C_Din);
|
||||
dim3 fgrid((N - 1) / C_N + 1, (Din - 1) / C_Din + 1, B);
|
||||
|
||||
const int theta_size = Dp * C_Din * C_Dout;
|
||||
const int bias_size = C_Din * C_Dout;
|
||||
const int topdiff_size = C_N * C_Dout;
|
||||
const int pos_size = C_N * K * Dp;
|
||||
const int nk_size = C_N * K;
|
||||
|
||||
int shm =
|
||||
(theta_size + bias_size + topdiff_size + pos_size) * sizeof(Dtype) +
|
||||
(nk_size) * sizeof(int);
|
||||
|
||||
runBackwardKernel<<<fgrid, fblock, shm>>>((*this));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void operator()() const {
|
||||
extern __shared__ float s_shm[];
|
||||
|
||||
int i_n = threadIdx.x;
|
||||
int i_din = threadIdx.y;
|
||||
|
||||
int b = blockIdx.z;
|
||||
int n = blockIdx.x * C_N + i_n;
|
||||
int din = blockIdx.y * C_Din + i_din;
|
||||
|
||||
Dtype* s_theta = (Dtype*)&s_shm[0];
|
||||
Dtype* s_bias = (Dtype*)&s_theta[Dp * C_Din * C_Dout];
|
||||
Dtype* s_topdiff = (Dtype*)&s_bias[C_Din * C_Dout];
|
||||
Dtype* s_pos = (Dtype*)&s_topdiff[C_N * C_Dout];
|
||||
int* s_nk = (int*)&s_pos[C_N * K * Dp];
|
||||
|
||||
for (int k = threadIdx.y; k < K && n < N; k += blockDim.y) {
|
||||
int nk = d_neigh[b * K * N + k * N + n];
|
||||
s_nk[k * C_N + i_n] = nk;
|
||||
|
||||
for (int i_dp = 0; i_dp < Dp; ++i_dp) {
|
||||
s_pos[k * C_N * Dp + i_dp * C_N + i_n] =
|
||||
d_pos[b * Dp * N + i_dp * N + nk];
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
for (int i_dp = 0; i_dp < Dp; ++i_dp) {
|
||||
Dtype val0 = s_pos[0 * C_N * Dp + i_dp * C_N + i_n];
|
||||
__syncthreads();
|
||||
for (int k = threadIdx.y; k < K && n < N; k += blockDim.y) {
|
||||
s_pos[k * C_N * Dp + i_dp * C_N + i_n] -= val0;
|
||||
}
|
||||
}
|
||||
|
||||
for (int dout_outer = 0; dout_outer < (Dout - 1) / C_Dout + 1;
|
||||
++dout_outer) {
|
||||
__syncthreads();
|
||||
|
||||
// fill s_theta
|
||||
int dout = dout_outer * C_Dout + i_n;
|
||||
if (din < Din && dout < Dout) {
|
||||
for (int i_dp = 0; i_dp < Dp; ++i_dp)
|
||||
s_theta[i_dp * C_Din * C_Dout + i_din * C_Dout + i_n] =
|
||||
d_theta[i_dp * Din * Dout + din * Dout + dout];
|
||||
|
||||
s_bias[i_din * C_Dout + i_n] = d_bias[din * Dout + dout];
|
||||
}
|
||||
|
||||
if (n < N) {
|
||||
for (int i_dout = threadIdx.y;
|
||||
i_dout < C_Dout && (dout_outer * C_Dout + i_dout) < Dout;
|
||||
i_dout += blockDim.y)
|
||||
s_topdiff[i_dout * C_N + i_n] =
|
||||
d_topdiff[b * Dout * N + (dout_outer * C_Dout + i_dout) * N + n];
|
||||
}
|
||||
|
||||
for (int dout_inner = 0;
|
||||
dout_inner < C_Dout && (dout_outer * C_Dout + dout_inner) < Dout;
|
||||
++dout_inner) {
|
||||
for (int k = 0; k < K; k++) {
|
||||
__syncthreads();
|
||||
|
||||
if (n < N && din < Din) {
|
||||
Dtype W = 0;
|
||||
for (int dp = 0; dp < Dp; ++dp) {
|
||||
const Dtype diffpos = s_pos[k * C_N * Dp + dp * C_N + i_n];
|
||||
W += s_theta[dp * C_Din * C_Dout + i_din * C_Dout + dout_inner] *
|
||||
diffpos;
|
||||
}
|
||||
W += s_bias[i_din * C_Dout + dout_inner];
|
||||
Dtype value = W * s_topdiff[dout_inner * C_N + i_n];
|
||||
|
||||
atomicAdd(
|
||||
&d_features_out[b * Din * N + din * N + s_nk[k * C_N + i_n]],
|
||||
value);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const Dtype* d_topdiff;
|
||||
const Dtype* d_pos;
|
||||
const Dtype* d_features;
|
||||
const int* d_neigh;
|
||||
const Dtype* d_theta;
|
||||
const Dtype* d_bias;
|
||||
|
||||
Dtype* d_features_out;
|
||||
|
||||
int N;
|
||||
int K;
|
||||
int Dp;
|
||||
int Din;
|
||||
int Dout;
|
||||
};
|
||||
|
||||
} // namespace FlexConvCuda
|
||||
|
||||
namespace tensorflow {
|
||||
namespace functor {
|
||||
|
||||
template <typename Dtype>
|
||||
struct FlexConvFunctor<GPUDevice, Dtype> {
|
||||
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& features,
|
||||
const Tensor& theta, const Tensor& bias,
|
||||
const Tensor& neighborhood, const Tensor& positions,
|
||||
Tensor* output) {
|
||||
typedef int NBtype;
|
||||
|
||||
const int B = neighborhood.dim_size(0);
|
||||
const int K = neighborhood.dim_size(1);
|
||||
const int N = neighborhood.dim_size(2);
|
||||
const int Dp = theta.dim_size(1);
|
||||
const int Din = theta.dim_size(2);
|
||||
const int Dout = theta.dim_size(3);
|
||||
|
||||
FlexConvCuda::ForwardKernel<Dtype, NBtype, 3, 128, 32, 64> fwk;
|
||||
fwk.N = N;
|
||||
fwk.K = K;
|
||||
fwk.Din = Din;
|
||||
fwk.Dout = Dout;
|
||||
|
||||
fwk.d_features = features.flat<Dtype>().data();
|
||||
fwk.d_positions = positions.flat<Dtype>().data();
|
||||
fwk.d_neighborhood = neighborhood.flat<NBtype>().data();
|
||||
fwk.d_theta = theta.flat<Dtype>().data();
|
||||
fwk.d_bias = bias.flat<Dtype>().data();
|
||||
fwk.d_output = output->flat<Dtype>().data();
|
||||
|
||||
fwk.launch(B);
|
||||
|
||||
if (!ctx->eigen_gpu_device().ok()) {
|
||||
ctx->SetStatus(tensorflow::errors::Internal(
|
||||
"FlexConvInvFunctor::forward::ForwardKernel execution failed"));
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template struct FlexConvFunctor<GPUDevice, float>;
|
||||
|
||||
template <typename Dtype>
|
||||
struct FlexConvGrad<GPUDevice, Dtype> {
|
||||
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& features_,
|
||||
const Tensor& theta_, const Tensor& bias_,
|
||||
const Tensor& neighborhood_, const Tensor& positions_,
|
||||
const Tensor& topdiff_, Tensor* grad_features_,
|
||||
Tensor* grad_theta_, Tensor* grad_bias_) {
|
||||
const auto features = features_.tensor<Dtype, 3>();
|
||||
const auto theta = theta_.tensor<Dtype, 4>();
|
||||
const auto bias = bias_.tensor<Dtype, 2>();
|
||||
const auto neighborhood = neighborhood_.tensor<int, 3>();
|
||||
const auto positions = positions_.tensor<Dtype, 3>();
|
||||
const auto topdiff = topdiff_.tensor<Dtype, 3>();
|
||||
|
||||
auto grad_features = grad_features_->tensor<Dtype, 3>();
|
||||
auto grad_theta = grad_theta_->tensor<Dtype, 4>();
|
||||
auto grad_bias = grad_bias_->tensor<Dtype, 2>();
|
||||
|
||||
// get dimensions
|
||||
const int B = neighborhood_.dim_size(0);
|
||||
const int K = neighborhood_.dim_size(1);
|
||||
const int N = neighborhood_.dim_size(2);
|
||||
const int Dp = theta_.dim_size(1);
|
||||
const int Din = theta_.dim_size(2);
|
||||
const int Dout = theta_.dim_size(3);
|
||||
|
||||
const int* neighborhood_ptr =
|
||||
reinterpret_cast<const int*>(neighborhood.data());
|
||||
const Dtype* positions_ptr =
|
||||
reinterpret_cast<const Dtype*>(positions.data());
|
||||
const Dtype* features_ptr = reinterpret_cast<const Dtype*>(features.data());
|
||||
const Dtype* theta_ptr = reinterpret_cast<const Dtype*>(theta.data());
|
||||
const Dtype* bias_ptr = reinterpret_cast<const Dtype*>(bias.data());
|
||||
|
||||
const Dtype* topdiff_ptr = reinterpret_cast<const Dtype*>(topdiff.data());
|
||||
|
||||
Dtype* grad_features_ptr = reinterpret_cast<float*>(grad_features.data());
|
||||
Dtype* grad_theta_ptr = reinterpret_cast<Dtype*>(grad_theta.data());
|
||||
Dtype* grad_bias_ptr = reinterpret_cast<Dtype*>(grad_bias.data());
|
||||
|
||||
cudaMemset(grad_features_ptr, 0, B * Din * N * sizeof(Dtype));
|
||||
|
||||
::tensorflow::CudaLaunchConfig cfg =
|
||||
::tensorflow::GetCudaLaunchConfig(N, ctx->eigen_device<GPUDevice>());
|
||||
|
||||
typedef FlexConvCuda::BackwardFeatureKernel<Dtype> BFK;
|
||||
|
||||
BFK bfk;
|
||||
bfk.N = N;
|
||||
bfk.K = K;
|
||||
bfk.Dp = Dp;
|
||||
bfk.Din = Din;
|
||||
bfk.Dout = Dout;
|
||||
|
||||
bfk.d_pos = positions_ptr;
|
||||
bfk.d_neigh = neighborhood_ptr;
|
||||
bfk.d_features = features_ptr;
|
||||
bfk.d_theta = theta_ptr;
|
||||
bfk.d_bias = bias_ptr;
|
||||
|
||||
bfk.d_topdiff = topdiff_ptr;
|
||||
|
||||
bfk.d_features_out = grad_features_ptr;
|
||||
|
||||
bfk.launch(B);
|
||||
|
||||
if (!ctx->eigen_gpu_device().ok()) {
|
||||
ctx->SetStatus(
|
||||
tensorflow::errors::Internal("CUDA: BackwardFeatureKernel Error!\n"));
|
||||
}
|
||||
|
||||
typedef FlexConvCuda::BackwardThetaKernel<Dtype> BTK;
|
||||
|
||||
BTK btk;
|
||||
btk.B = B;
|
||||
btk.N = N;
|
||||
btk.K = K;
|
||||
btk.Dp = Dp;
|
||||
btk.Din = Din;
|
||||
btk.Dout = Dout;
|
||||
|
||||
btk.d_pos = positions_ptr;
|
||||
btk.d_neigh = neighborhood_ptr;
|
||||
btk.d_features = features_ptr;
|
||||
btk.d_theta = theta_ptr;
|
||||
btk.d_bias = bias_ptr;
|
||||
|
||||
btk.d_topdiff = topdiff_ptr;
|
||||
|
||||
btk.d_theta_out = grad_theta_ptr;
|
||||
btk.d_bias_out = grad_bias_ptr;
|
||||
|
||||
btk.launch();
|
||||
|
||||
if (!ctx->eigen_gpu_device().ok()) {
|
||||
ctx->SetStatus(
|
||||
tensorflow::errors::Internal("CUDA: BackwardThetaKernel Error!\n"));
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template struct FlexConvGrad<GPUDevice, float>;
|
||||
|
||||
} // namespace functor
|
||||
} // namespace tensorflow
|
||||
|
||||
#endif // GOOGLE_CUDA
|
||||
@@ -0,0 +1,122 @@
|
||||
/* 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_);
|
||||
}
|
||||
};
|
||||
|
||||
// Register the CPU kernels.
|
||||
#define REGISTER_FLEXCONV_OP_CPU(T) \
|
||||
REGISTER_KERNEL_BUILDER( \
|
||||
Name("FlexConv").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
|
||||
FlexConvOp<CPUDevice, T>) \
|
||||
REGISTER_KERNEL_BUILDER( \
|
||||
Name("FlexConvGrad").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
|
||||
FlexConvGradOp<CPUDevice, T>)
|
||||
|
||||
TF_CALL_float(REGISTER_FLEXCONV_OP_CPU);
|
||||
#undef REGISTER_FLEXCONV_OP_CPU
|
||||
|
||||
// Register the GPU kernels.
|
||||
// #ifdef GOOGLE_CUDA
|
||||
|
||||
#define REGISTER_FLEXCONV_OP_GPU(T) \
|
||||
REGISTER_KERNEL_BUILDER( \
|
||||
Name("FlexConv").Device(DEVICE_GPU).TypeConstraint<T>("T"), \
|
||||
FlexConvOp<GPUDevice, T>) \
|
||||
REGISTER_KERNEL_BUILDER( \
|
||||
Name("FlexConvGrad").Device(DEVICE_GPU).TypeConstraint<T>("T"), \
|
||||
FlexConvGradOp<GPUDevice, T>)
|
||||
|
||||
TF_CALL_float(REGISTER_FLEXCONV_OP_GPU);
|
||||
#undef REGISTER_FLEXCONV_OP_GPU
|
||||
|
||||
// #endif // GOOGLE_CUDA
|
||||
|
||||
} // namespace tensorflow
|
||||
@@ -0,0 +1,53 @@
|
||||
/* 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
|
||||
|
||||
#ifndef USER_OPS_KERNELS_FLEX_CONV_OP_H_
|
||||
#define USER_OPS_KERNELS_FLEX_CONV_OP_H_
|
||||
|
||||
#include "tensorflow/core/framework/op_kernel.h"
|
||||
|
||||
namespace tensorflow {
|
||||
class OpKernelContext;
|
||||
class Tensor;
|
||||
|
||||
using CPUDevice = Eigen::ThreadPoolDevice;
|
||||
using GPUDevice = Eigen::GpuDevice;
|
||||
} // namespace tensorflow
|
||||
|
||||
namespace tensorflow {
|
||||
namespace functor {
|
||||
|
||||
template <typename Device, typename Dtype>
|
||||
struct FlexConvFunctor {
|
||||
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& features_,
|
||||
const Tensor& theta_, const Tensor& bias_,
|
||||
const Tensor& neighborhood_, const Tensor& positions_,
|
||||
Tensor* output_);
|
||||
};
|
||||
|
||||
template <typename Device, typename Dtype>
|
||||
struct FlexConvGrad {
|
||||
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& features_,
|
||||
const Tensor& theta_, const Tensor& bias_,
|
||||
const Tensor& neighborhood_, const Tensor& positions_,
|
||||
const Tensor& topdiff_, Tensor* grad_features_,
|
||||
Tensor* grad_theta_, Tensor* grad_bias_);
|
||||
};
|
||||
|
||||
} // namespace functor
|
||||
} // namespace tensorflow
|
||||
|
||||
#endif // USER_OPS_KERNELS_FLEX_CONV_OP_H_
|
||||
@@ -0,0 +1,173 @@
|
||||
/* 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_deconv_op.h"
|
||||
#include "tensorflow/core/framework/op.h"
|
||||
|
||||
namespace tensorflow {
|
||||
|
||||
namespace functor {
|
||||
|
||||
template <typename Dtype>
|
||||
struct FlexDeconvFunctor<CPUDevice, Dtype> {
|
||||
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& features_,
|
||||
const Tensor& theta_, const Tensor& bias_,
|
||||
const Tensor& neighborhood_, const Tensor& positions_,
|
||||
Tensor* output_) {
|
||||
const auto features = features_.tensor<Dtype, 3>();
|
||||
const auto theta = theta_.tensor<Dtype, 4>();
|
||||
const auto bias = bias_.tensor<Dtype, 2>();
|
||||
const auto neighborhood = neighborhood_.tensor<int, 3>();
|
||||
const auto positions = positions_.tensor<Dtype, 3>();
|
||||
|
||||
auto output = output_->tensor<Dtype, 3>();
|
||||
|
||||
// get dimensions
|
||||
const int B = neighborhood_.dim_size(0);
|
||||
const int K = neighborhood_.dim_size(1);
|
||||
const int N = neighborhood_.dim_size(2);
|
||||
const int Dp = theta_.dim_size(1);
|
||||
const int Din = theta_.dim_size(2);
|
||||
const int Dout = theta_.dim_size(3);
|
||||
|
||||
output.setZero();
|
||||
|
||||
for (int b = 0; b < B; ++b) {
|
||||
for (int n = 0; n < N; ++n) {
|
||||
const int self_k = neighborhood(b, 0, n);
|
||||
for (int k_ = 0; k_ < K; ++k_) {
|
||||
const int other_k = neighborhood(b, k_, n);
|
||||
|
||||
for (int dout = 0; dout < Dout; ++dout) {
|
||||
for (int din = 0; din < Din; ++din) {
|
||||
const Dtype v = features(b, din, self_k);
|
||||
|
||||
Dtype W = bias(din, dout);
|
||||
for (int dp = 0; dp < Dp; ++dp) {
|
||||
Dtype delta =
|
||||
positions(b, dp, other_k) - positions(b, dp, self_k);
|
||||
W += theta(0, dp, din, dout) * delta;
|
||||
}
|
||||
output(b, dout, other_k) = output(b, dout, other_k) + W * v;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template struct FlexDeconvFunctor<CPUDevice, float>;
|
||||
|
||||
template <typename Dtype>
|
||||
struct FlexDeconvGrad<CPUDevice, Dtype> {
|
||||
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& features_,
|
||||
const Tensor& theta_, const Tensor& bias_,
|
||||
const Tensor& neighborhood_, const Tensor& positions_,
|
||||
const Tensor& topdiff_, Tensor* grad_features_,
|
||||
Tensor* grad_theta_, Tensor* grad_bias_) {
|
||||
const auto features = features_.tensor<Dtype, 3>();
|
||||
const auto theta = theta_.tensor<Dtype, 4>();
|
||||
const auto bias = bias_.tensor<Dtype, 2>();
|
||||
const auto neighborhood = neighborhood_.tensor<int, 3>();
|
||||
const auto positions = positions_.tensor<Dtype, 3>();
|
||||
const auto topdiff = topdiff_.tensor<Dtype, 3>();
|
||||
|
||||
auto grad_features = grad_features_->tensor<Dtype, 3>();
|
||||
auto grad_theta = grad_theta_->tensor<Dtype, 4>();
|
||||
auto grad_bias = grad_bias_->tensor<Dtype, 2>();
|
||||
|
||||
// get dimensions
|
||||
const int B = neighborhood_.dim_size(0);
|
||||
const int K = neighborhood_.dim_size(1);
|
||||
const int N = neighborhood_.dim_size(2);
|
||||
const int Dp = theta_.dim_size(1);
|
||||
const int Din = theta_.dim_size(2);
|
||||
const int Dout = theta_.dim_size(3);
|
||||
|
||||
grad_features.setZero();
|
||||
grad_theta.setZero();
|
||||
grad_bias.setZero();
|
||||
|
||||
// ========================= bias ==============================
|
||||
for (int b = 0; b < B; ++b) {
|
||||
for (int n = 0; n < N; ++n) {
|
||||
const int self_k = neighborhood(b, 0, n);
|
||||
for (int k_ = 0; k_ < K; ++k_) {
|
||||
const int other_k = neighborhood(b, k_, n);
|
||||
|
||||
for (int din = 0; din < Din; ++din) {
|
||||
for (int dout = 0; dout < Dout; ++dout) {
|
||||
grad_bias(din, dout) +=
|
||||
features(b, din, self_k) * topdiff(b, dout, other_k);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ========================= theta ==============================
|
||||
for (int b = 0; b < B; ++b) {
|
||||
for (int n = 0; n < N; ++n) {
|
||||
const int self_k = neighborhood(b, 0, n);
|
||||
for (int k_ = 0; k_ < K; ++k_) {
|
||||
const int other_k = neighborhood(b, k_, n);
|
||||
|
||||
for (int din = 0; din < Din; ++din) {
|
||||
for (int dout = 0; dout < Dout; ++dout) {
|
||||
for (int dp = 0; dp < Dp; ++dp) {
|
||||
const Dtype delta =
|
||||
positions(b, dp, other_k) - positions(b, dp, self_k);
|
||||
grad_theta(0, dp, din, dout) += features(b, din, self_k) *
|
||||
delta *
|
||||
topdiff(b, dout, other_k);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ========================= features ==============================
|
||||
for (int b = 0; b < B; ++b) {
|
||||
for (int n = 0; n < N; ++n) {
|
||||
const int self_k = neighborhood(b, 0, n);
|
||||
for (int k_ = 0; k_ < K; ++k_) {
|
||||
const int other_k = neighborhood(b, k_, n);
|
||||
|
||||
for (int din = 0; din < Din; ++din) {
|
||||
for (int dout = 0; dout < Dout; ++dout) {
|
||||
Dtype W = bias(din, dout);
|
||||
for (int dp = 0; dp < Dp; ++dp) {
|
||||
const Dtype delta =
|
||||
positions(b, dp, other_k) - positions(b, dp, self_k);
|
||||
W += theta(0, dp, din, dout) * delta;
|
||||
}
|
||||
grad_features(b, din, self_k) += W * topdiff(b, dout, other_k);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// template struct FlexDeconvGrad<CPUDevice, int>;
|
||||
template struct FlexDeconvGrad<CPUDevice, float>;
|
||||
// template struct FlexDeconvGrad<CPUDevice, double>;
|
||||
|
||||
} // namespace functor
|
||||
} // namespace tensorflow
|
||||
@@ -0,0 +1,227 @@
|
||||
/* 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
|
||||
|
||||
#if GOOGLE_CUDA
|
||||
|
||||
#define EIGEN_USE_GPU
|
||||
|
||||
#include <cub/cub.cuh>
|
||||
|
||||
#include "flex_deconv_op.h"
|
||||
#include "tensorflow/core/util/cuda_kernel_helper.h"
|
||||
|
||||
namespace {
|
||||
inline int up2(int len, int th) { return (len - 1) / th + 1; }
|
||||
|
||||
template <typename Dtype>
|
||||
__global__ void forward(const int B, const int N, const int K, const int Dp,
|
||||
const int Din, const int Dout, const Dtype* positions,
|
||||
const Dtype* features, const int* neighborhood,
|
||||
const Dtype* theta, const Dtype* bias, Dtype* output) {
|
||||
/*
|
||||
positions B, Dp, N
|
||||
features B, Din, N
|
||||
neighborhood B, K, N
|
||||
theta Dp, Din, Dout
|
||||
bias Din, Dout
|
||||
output B, Dout, N
|
||||
*/
|
||||
|
||||
const int b = blockIdx.z;
|
||||
|
||||
for (int n = blockIdx.y * blockDim.y + threadIdx.y; n < N;
|
||||
n += blockDim.y * gridDim.y) {
|
||||
const int self_k = neighborhood[b * K * N + 0 * N + n];
|
||||
|
||||
for (int k_ = 0; k_ < K; ++k_) {
|
||||
const int other_k = neighborhood[b * K * N + k_ * N + n];
|
||||
|
||||
for (int dout = blockIdx.x * blockDim.x + threadIdx.x; dout < Dout;
|
||||
dout += blockDim.x * gridDim.x) {
|
||||
for (int din = 0; din < Din; ++din) {
|
||||
const Dtype v = features[b * Din * N + din * N + self_k];
|
||||
Dtype W = bias[din * Dout + dout];
|
||||
|
||||
for (int dp = 0; dp < Dp; ++dp) {
|
||||
Dtype delta = positions[b * Dp * N + dp * N + other_k] -
|
||||
positions[b * Dp * N + dp * N + self_k];
|
||||
W += theta[dp * Din * Dout + din * Dout + dout] * delta;
|
||||
}
|
||||
|
||||
Dtype Wv = W * v;
|
||||
tensorflow::CudaAtomicAdd(&output[b * Dout * N + dout * N + other_k],
|
||||
Wv);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Dtype>
|
||||
__global__ void backward(const int B, const int N, const int K, const int Dp,
|
||||
const int Din, const int Dout,
|
||||
|
||||
const Dtype* positions, const Dtype* features,
|
||||
const int* neighborhood,
|
||||
|
||||
const Dtype* theta, const Dtype* bias,
|
||||
|
||||
const Dtype* top_diff,
|
||||
|
||||
Dtype* grad_features, Dtype* grad_theta,
|
||||
Dtype* grad_bias) {
|
||||
/*
|
||||
B, Dp, N positions, grad_positions
|
||||
B, Din, N features, grad_features
|
||||
B, K, N neighborhood
|
||||
Dp, Din, Dout theta, grad_theta
|
||||
Din, Dout bias, grad_bias
|
||||
B, Dout, N output, top_diff
|
||||
*/
|
||||
|
||||
const int b = blockIdx.z;
|
||||
|
||||
// Compute
|
||||
// ---------------------------------------------------------------
|
||||
|
||||
for (int n = blockIdx.y * blockDim.y + threadIdx.y; n < N;
|
||||
n += blockDim.y * gridDim.y) {
|
||||
const int self_k = neighborhood[b * K * N + 0 * N + n];
|
||||
|
||||
for (int k_ = 0; k_ < K; ++k_) {
|
||||
const int other_k = neighborhood[b * K * N + k_ * N + n];
|
||||
|
||||
for (int dout = blockIdx.x * blockDim.x + threadIdx.x; dout < Dout;
|
||||
dout += blockDim.x * gridDim.x) {
|
||||
for (int din = 0; din < Din; ++din) {
|
||||
const Dtype current_top_diff =
|
||||
top_diff[b * Dout * N + dout * N + other_k];
|
||||
const Dtype v = features[b * Din * N + din * N + self_k];
|
||||
|
||||
// update bias
|
||||
Dtype bias_update = v * current_top_diff;
|
||||
tensorflow::CudaAtomicAdd(&grad_bias[din * Dout + dout], bias_update);
|
||||
|
||||
Dtype W = bias[din * Dout + dout];
|
||||
|
||||
// update theta
|
||||
for (int dp = 0; dp < Dp; ++dp) {
|
||||
Dtype delta = positions[b * Dp * N + dp * N + other_k] -
|
||||
positions[b * Dp * N + dp * N + self_k];
|
||||
Dtype theta_update = v * delta * current_top_diff;
|
||||
tensorflow::CudaAtomicAdd(
|
||||
&grad_theta[dp * Din * Dout + din * Dout + dout], theta_update);
|
||||
|
||||
W += theta[dp * Din * Dout + din * Dout + dout] * delta;
|
||||
}
|
||||
|
||||
// update features
|
||||
Dtype feature_update = W * current_top_diff;
|
||||
tensorflow::CudaAtomicAdd(
|
||||
&grad_features[b * Din * N + din * N + self_k], feature_update);
|
||||
// tensorflow::CudaAtomicAdd(&grad_features[b * Din * N + din * N +
|
||||
// self_k], 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
namespace tensorflow {
|
||||
namespace functor {
|
||||
|
||||
template <typename Dtype>
|
||||
struct FlexDeconvFunctor<GPUDevice, Dtype> {
|
||||
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& features_,
|
||||
const Tensor& theta_, const Tensor& bias_,
|
||||
const Tensor& neighborhood_, const Tensor& positions_,
|
||||
Tensor* output_) {
|
||||
// printf("GPU::FlexDeconvFunctor:operator()\n");
|
||||
|
||||
// get dimensions
|
||||
const int B = neighborhood_.dim_size(0);
|
||||
const int K = neighborhood_.dim_size(1);
|
||||
const int N = neighborhood_.dim_size(2);
|
||||
const int Dp = theta_.dim_size(1);
|
||||
const int Din = theta_.dim_size(2);
|
||||
const int Dout = theta_.dim_size(3);
|
||||
|
||||
const int threads = 32;
|
||||
dim3 block(threads, threads, 1);
|
||||
dim3 grid(up2(Dout, threads), up2(N, threads), B);
|
||||
|
||||
cudaMemset(output_->flat<Dtype>().data(), 0,
|
||||
output_->NumElements() * sizeof(Dtype));
|
||||
|
||||
forward<Dtype><<<grid, block>>>(
|
||||
B, N, K, Dp, Din, Dout, positions_.flat<Dtype>().data(),
|
||||
features_.flat<Dtype>().data(), neighborhood_.flat<int>().data(),
|
||||
theta_.flat<Dtype>().data(), bias_.flat<Dtype>().data(),
|
||||
output_->flat<Dtype>().data());
|
||||
}
|
||||
};
|
||||
|
||||
template struct FlexDeconvFunctor<GPUDevice, float>;
|
||||
|
||||
template <typename Dtype>
|
||||
struct FlexDeconvGrad<GPUDevice, Dtype> {
|
||||
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& features_,
|
||||
const Tensor& theta_, const Tensor& bias_,
|
||||
const Tensor& neighborhood_, const Tensor& positions_,
|
||||
const Tensor& topdiff_, Tensor* grad_features_,
|
||||
Tensor* grad_theta_, Tensor* grad_bias_) {
|
||||
// get dimensions
|
||||
const int B = neighborhood_.dim_size(0);
|
||||
const int K = neighborhood_.dim_size(1);
|
||||
const int N = neighborhood_.dim_size(2);
|
||||
const int Dp = theta_.dim_size(1);
|
||||
const int Din = theta_.dim_size(2);
|
||||
const int Dout = theta_.dim_size(3);
|
||||
|
||||
const int threads = 32;
|
||||
dim3 block(threads, threads, 1);
|
||||
dim3 grid(up2(Dout, threads), up2(N, threads), B);
|
||||
|
||||
cudaMemset(grad_features_->flat<Dtype>().data(), 0,
|
||||
grad_features_->NumElements() * sizeof(Dtype));
|
||||
cudaMemset(grad_theta_->flat<Dtype>().data(), 0,
|
||||
grad_theta_->NumElements() * sizeof(Dtype));
|
||||
cudaMemset(grad_bias_->flat<Dtype>().data(), 0,
|
||||
grad_bias_->NumElements() * sizeof(Dtype));
|
||||
|
||||
backward<Dtype><<<grid, block>>>(
|
||||
B, N, K, Dp, Din, Dout,
|
||||
|
||||
positions_.flat<Dtype>().data(), features_.flat<Dtype>().data(),
|
||||
neighborhood_.flat<int>().data(),
|
||||
|
||||
theta_.flat<Dtype>().data(), bias_.flat<Dtype>().data(),
|
||||
|
||||
topdiff_.flat<Dtype>().data(),
|
||||
|
||||
grad_features_->flat<Dtype>().data(), grad_theta_->flat<Dtype>().data(),
|
||||
grad_bias_->flat<Dtype>().data());
|
||||
}
|
||||
};
|
||||
|
||||
template struct FlexDeconvGrad<GPUDevice, float>;
|
||||
|
||||
} // namespace functor
|
||||
} // namespace tensorflow
|
||||
|
||||
#endif // GOOGLE_CUDA
|
||||
@@ -0,0 +1,103 @@
|
||||
/* 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_deconv_op.h"
|
||||
|
||||
#include <stdio.h>
|
||||
#include <type_traits>
|
||||
|
||||
#include "tensorflow/core/framework/op.h"
|
||||
#include "tensorflow/core/framework/op_kernel.h"
|
||||
|
||||
namespace tensorflow {
|
||||
|
||||
// Forward-Pass (CPU, GPU)
|
||||
// --------------------------------------------------
|
||||
template <typename Device, typename Dtype>
|
||||
class FlexDeconvOp : public OpKernel {
|
||||
public:
|
||||
explicit FlexDeconvOp(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::FlexDeconvFunctor<Device, Dtype>()(
|
||||
ctx, features_, theta_, bias_, neighborhood_, positions_, output_);
|
||||
}
|
||||
|
||||
private:
|
||||
TF_DISALLOW_COPY_AND_ASSIGN(FlexDeconvOp);
|
||||
};
|
||||
|
||||
// Backward-Pass (CPU, GPU)
|
||||
// --------------------------------------------------
|
||||
template <typename Device, typename Dtype>
|
||||
class FlexDeconvGradOp : public OpKernel {
|
||||
public:
|
||||
explicit FlexDeconvGradOp(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;
|
||||
|
||||
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::FlexDeconvGrad<Device, Dtype>()(
|
||||
ctx, features_, theta_, bias_, neighborhood_, positions_, topdiff_,
|
||||
grad_features_, grad_theta_, grad_bias_);
|
||||
}
|
||||
};
|
||||
|
||||
#define OPNAME(NAME) NAME##Op
|
||||
#define REGISTER(NAME, Dtype) \
|
||||
REGISTER_KERNEL_BUILDER( \
|
||||
Name(#NAME).Device(DEVICE_CPU).TypeConstraint<Dtype>("T"), \
|
||||
OPNAME(NAME) < CPUDevice, Dtype >); \
|
||||
REGISTER_KERNEL_BUILDER( \
|
||||
Name(#NAME).Device(DEVICE_GPU).TypeConstraint<Dtype>("T"), \
|
||||
OPNAME(NAME) < GPUDevice, Dtype >);
|
||||
|
||||
REGISTER(FlexDeconv, float);
|
||||
REGISTER(FlexDeconvGrad, float);
|
||||
|
||||
} // namespace tensorflow
|
||||
@@ -0,0 +1,53 @@
|
||||
/* 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
|
||||
|
||||
#ifndef USER_OPS_KERNELS_FLEX_DECONV_OP_H_
|
||||
#define USER_OPS_KERNELS_FLEX_DECONV_OP_H_
|
||||
|
||||
#include "tensorflow/core/framework/op_kernel.h"
|
||||
|
||||
namespace tensorflow {
|
||||
class OpKernelContext;
|
||||
class Tensor;
|
||||
|
||||
using CPUDevice = Eigen::ThreadPoolDevice;
|
||||
using GPUDevice = Eigen::GpuDevice;
|
||||
} // namespace tensorflow
|
||||
|
||||
namespace tensorflow {
|
||||
namespace functor {
|
||||
|
||||
template <typename Device, typename Dtype>
|
||||
struct FlexDeconvFunctor {
|
||||
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& features_,
|
||||
const Tensor& theta_, const Tensor& bias_,
|
||||
const Tensor& neighborhood_, const Tensor& positions_,
|
||||
Tensor* output_);
|
||||
};
|
||||
|
||||
template <typename Device, typename Dtype>
|
||||
struct FlexDeconvGrad {
|
||||
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& features_,
|
||||
const Tensor& theta_, const Tensor& bias_,
|
||||
const Tensor& neighborhood_, const Tensor& positions_,
|
||||
const Tensor& topdiff_, Tensor* grad_features_,
|
||||
Tensor* grad_theta_, Tensor* grad_bias_);
|
||||
};
|
||||
|
||||
} // namespace functor
|
||||
} // namespace tensorflow
|
||||
|
||||
#endif // USER_OPS_KERNELS_FLEX_DECONV_OP_H_
|
||||
@@ -0,0 +1,102 @@
|
||||
/* 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_pool_op.h"
|
||||
#include "tensorflow/core/framework/op.h"
|
||||
|
||||
namespace tensorflow {
|
||||
|
||||
namespace functor {
|
||||
|
||||
template <typename Dtype>
|
||||
struct FlexPoolFunctor<CPUDevice, Dtype> {
|
||||
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& features_,
|
||||
const Tensor& neighborhood_, Tensor* output_,
|
||||
Tensor* argmax_) {
|
||||
const auto features = features_.tensor<Dtype, 3>();
|
||||
const auto neighborhood = neighborhood_.tensor<int, 3>();
|
||||
|
||||
auto output = output_->tensor<Dtype, 3>();
|
||||
auto argmax = argmax_->tensor<int, 3>();
|
||||
|
||||
// get dimensions
|
||||
const int B = neighborhood_.dim_size(0);
|
||||
const int K = neighborhood_.dim_size(1);
|
||||
const int N = neighborhood_.dim_size(2);
|
||||
const int D = features_.dim_size(1);
|
||||
|
||||
output.setConstant(Eigen::NumTraits<Dtype>::lowest());
|
||||
argmax.setZero(); // stores global id
|
||||
|
||||
for (int b = 0; b < B; ++b) {
|
||||
for (int d = 0; d < D; ++d) {
|
||||
for (int n = 0; n < N; ++n) {
|
||||
// max in neighborhood
|
||||
for (int k_ = 0; k_ < K; ++k_) {
|
||||
const int other_global_id = neighborhood(b, k_, n);
|
||||
if (output(b, d, n) < features(b, d, other_global_id)) {
|
||||
argmax(b, d, n) = other_global_id;
|
||||
output(b, d, n) = features(b, d, other_global_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template struct FlexPoolFunctor<CPUDevice, float>;
|
||||
|
||||
template <typename Dtype>
|
||||
struct FlexPoolGrad<CPUDevice, Dtype> {
|
||||
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& features_,
|
||||
const Tensor& neighborhood_, const Tensor& topdiff_,
|
||||
const Tensor& argmax_, Tensor* grad_features_) {
|
||||
// as only argmax contributes to the output
|
||||
// only argmax receives the topdiff
|
||||
|
||||
const auto features = features_.tensor<Dtype, 3>();
|
||||
const auto neighborhood = neighborhood_.tensor<int, 3>();
|
||||
const auto topdiff = topdiff_.tensor<Dtype, 3>();
|
||||
const auto argmax = argmax_.tensor<int, 3>();
|
||||
|
||||
auto grad_features = grad_features_->tensor<Dtype, 3>();
|
||||
|
||||
// get dimensions
|
||||
const int B = neighborhood_.dim_size(0);
|
||||
const int K = neighborhood_.dim_size(1);
|
||||
const int N = neighborhood_.dim_size(2);
|
||||
const int D = features_.dim_size(1);
|
||||
// printf("B %i K %i N %i D %i\n", B, K ,N, D);
|
||||
|
||||
grad_features.setZero();
|
||||
|
||||
for (int b = 0; b < B; ++b) {
|
||||
for (int d = 0; d < D; ++d) {
|
||||
for (int n = 0; n < N; ++n) {
|
||||
grad_features(b, d, argmax(b, d, n)) += topdiff(b, d, n);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// template struct FlexPoolGrad<CPUDevice, int>;
|
||||
template struct FlexPoolGrad<CPUDevice, float>;
|
||||
// template struct FlexPoolGrad<CPUDevice, double>;
|
||||
|
||||
} // namespace functor
|
||||
} // namespace tensorflow
|
||||
@@ -0,0 +1,171 @@
|
||||
/* 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
|
||||
|
||||
#if GOOGLE_CUDA
|
||||
|
||||
#define EIGEN_USE_GPU
|
||||
|
||||
#include <cub/cub.cuh>
|
||||
#include <limits>
|
||||
|
||||
#include "flex_pool_op.h"
|
||||
#include "tensorflow/core/util/cuda_kernel_helper.h"
|
||||
|
||||
namespace {
|
||||
inline int up2(int len, int th) { return (len - 1) / th + 1; }
|
||||
|
||||
template <typename Dtype>
|
||||
__global__ void forward(const int B, const int N, const int K, const int D,
|
||||
const Dtype* features, const int* neighborhood,
|
||||
Dtype* output, int* argmax, float float_min_value) {
|
||||
// features: each feature description for each point [B, D, N].
|
||||
// neighborhood: all K nearest neighbors [B, K, N].
|
||||
// output: each feature description for each point [B, D, N].
|
||||
// argmax: global id in neighborhood who was winning the pooling [B, D, N].
|
||||
const int b = blockIdx.z;
|
||||
|
||||
for (int d = blockIdx.y * blockDim.y + threadIdx.y; d < D;
|
||||
d += blockDim.y * gridDim.y) {
|
||||
for (int n = blockIdx.x * blockDim.x + threadIdx.x; n < N;
|
||||
n += blockDim.x * gridDim.x) {
|
||||
float best_value = float_min_value;
|
||||
int best_id = 0;
|
||||
|
||||
const int current_flat = b * D * N + d * N + n;
|
||||
|
||||
for (int k_ = 0; k_ < K; ++k_) {
|
||||
const int other_global_id = neighborhood[b * K * N + k_ * N + n];
|
||||
const float v = features[b * D * N + d * N + other_global_id];
|
||||
|
||||
if (best_value < v) {
|
||||
best_id = other_global_id;
|
||||
best_value = v;
|
||||
}
|
||||
}
|
||||
|
||||
output[current_flat] = best_value;
|
||||
argmax[current_flat] = best_id;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Dtype>
|
||||
__global__ void backward(const int B, const int N, const int K, const int D,
|
||||
|
||||
const Dtype* features, const int* neighborhood,
|
||||
const Dtype* topdiff, const int* argmax,
|
||||
|
||||
Dtype* grad_features) {
|
||||
// features: each feature description for each point [B, D, N].
|
||||
// neighborhood: all K nearest neighbors [B, K, N].
|
||||
// gradients: topdiff[B, D, N].
|
||||
// argmax: argmax[B, D, N].
|
||||
// grad_features: gradient to each feature description for each point [B, D,
|
||||
// N].
|
||||
const int b = blockIdx.z;
|
||||
|
||||
for (int d = blockIdx.y * blockDim.y + threadIdx.y; d < D;
|
||||
d += blockDim.y * gridDim.y) {
|
||||
for (int n = blockIdx.x * blockDim.x + threadIdx.x; n < N;
|
||||
n += blockDim.x * gridDim.x) {
|
||||
const int top_id_flat = b * D * N + d * N + n;
|
||||
const int argmax_id = argmax[top_id_flat];
|
||||
const int bottom_id_flat = b * D * N + d * N + argmax_id;
|
||||
|
||||
// TODO(patwie): scattered write, yeah :-(
|
||||
tensorflow::CudaAtomicAdd(&grad_features[bottom_id_flat],
|
||||
topdiff[top_id_flat]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
namespace tensorflow {
|
||||
namespace functor {
|
||||
|
||||
template <typename Dtype>
|
||||
struct FlexPoolFunctor<GPUDevice, Dtype> {
|
||||
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& features_,
|
||||
const Tensor& neighborhood_, Tensor* output_,
|
||||
Tensor* argmax_) {
|
||||
// get dimensions
|
||||
const int B = neighborhood_.dim_size(0);
|
||||
const int K = neighborhood_.dim_size(1);
|
||||
const int N = neighborhood_.dim_size(2);
|
||||
const int D = features_.dim_size(1);
|
||||
|
||||
const int threads = 32;
|
||||
dim3 block(threads, threads, 1);
|
||||
dim3 grid(up2(N, threads), up2(D, threads), B);
|
||||
|
||||
forward<Dtype><<<grid, block>>>(
|
||||
B, N, K, D,
|
||||
|
||||
features_.flat<Dtype>().data(), neighborhood_.flat<int>().data(),
|
||||
|
||||
output_->flat<Dtype>().data(), argmax_->flat<int>().data(),
|
||||
|
||||
std::numeric_limits<Dtype>::lowest());
|
||||
|
||||
if (!ctx->eigen_gpu_device().ok()) {
|
||||
ctx->SetStatus(
|
||||
tensorflow::errors::Internal("CUDA: FlexPoolFunctor Error!"));
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template struct FlexPoolFunctor<GPUDevice, float>;
|
||||
|
||||
template <typename Dtype>
|
||||
struct FlexPoolGrad<GPUDevice, Dtype> {
|
||||
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& features_,
|
||||
const Tensor& neighborhood_, const Tensor& topdiff_,
|
||||
const Tensor& argmax_, Tensor* grad_features_) {
|
||||
// get dimensions
|
||||
const int B = neighborhood_.dim_size(0);
|
||||
const int K = neighborhood_.dim_size(1);
|
||||
const int N = neighborhood_.dim_size(2);
|
||||
const int D = features_.dim_size(1);
|
||||
|
||||
const int threads = 32;
|
||||
dim3 block(threads, threads, 1);
|
||||
dim3 grid(up2(N, threads), up2(D, threads), B);
|
||||
|
||||
cudaMemset(grad_features_->flat<Dtype>().data(), 0,
|
||||
grad_features_->NumElements() * sizeof(Dtype));
|
||||
|
||||
backward<Dtype><<<grid, block>>>(
|
||||
B, N, K, D,
|
||||
|
||||
features_.flat<Dtype>().data(), neighborhood_.flat<int>().data(),
|
||||
|
||||
topdiff_.flat<Dtype>().data(), argmax_.flat<int>().data(),
|
||||
|
||||
grad_features_->flat<Dtype>().data());
|
||||
|
||||
if (!ctx->eigen_gpu_device().ok()) {
|
||||
ctx->SetStatus(tensorflow::errors::Internal("CUDA: FlexPoolGrad Error!"));
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template struct FlexPoolGrad<GPUDevice, float>;
|
||||
|
||||
} // namespace functor
|
||||
} // namespace tensorflow
|
||||
|
||||
#endif // GOOGLE_CUDA
|
||||
@@ -0,0 +1,113 @@
|
||||
/* 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_pool_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 FlexPoolOp : public OpKernel {
|
||||
public:
|
||||
explicit FlexPoolOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
|
||||
|
||||
void Compute(OpKernelContext* ctx) override {
|
||||
// printf("--> Compute CPU Version <--\n");
|
||||
const Tensor& features_ = ctx->input(0);
|
||||
const Tensor& neighborhood_ = ctx->input(1);
|
||||
|
||||
const int B = features_.dim_size(0);
|
||||
const int D = features_.dim_size(1);
|
||||
const int N = features_.dim_size(2);
|
||||
|
||||
Tensor* output_ = nullptr;
|
||||
OP_REQUIRES_OK(ctx,
|
||||
ctx->allocate_output(0, TensorShape({B, D, N}), &output_));
|
||||
|
||||
Tensor* argmax_ = nullptr;
|
||||
OP_REQUIRES_OK(ctx,
|
||||
ctx->allocate_output(1, TensorShape({B, D, N}), &argmax_));
|
||||
|
||||
::tensorflow::functor::FlexPoolFunctor<Device, Dtype>()(
|
||||
ctx, features_, neighborhood_, output_, argmax_);
|
||||
}
|
||||
|
||||
private:
|
||||
TF_DISALLOW_COPY_AND_ASSIGN(FlexPoolOp);
|
||||
};
|
||||
|
||||
// Backward-Pass (CPU, GPU)
|
||||
// --------------------------------------------------
|
||||
template <typename Device, typename Dtype>
|
||||
class FlexPoolGradOp : public OpKernel {
|
||||
public:
|
||||
explicit FlexPoolGradOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
|
||||
|
||||
void Compute(OpKernelContext* ctx) override {
|
||||
// printf("--> Compute CPU Version <--\n");
|
||||
const Tensor& features_ = ctx->input(0);
|
||||
const Tensor& neighborhood_ = ctx->input(1);
|
||||
const Tensor& topdiff_ = ctx->input(2);
|
||||
const Tensor& argmax_ = ctx->input(3);
|
||||
|
||||
// specify output shape
|
||||
Tensor* grad_features_ = nullptr;
|
||||
|
||||
OP_REQUIRES_OK(ctx,
|
||||
ctx->allocate_output(0, features_.shape(), &grad_features_));
|
||||
|
||||
::tensorflow::functor::FlexPoolGrad<Device, Dtype>()(
|
||||
ctx, features_, neighborhood_, topdiff_, argmax_, grad_features_);
|
||||
}
|
||||
};
|
||||
|
||||
// Register the CPU kernels.
|
||||
#define REGISTER_FLEXPOOL_OP_CPU(T) \
|
||||
REGISTER_KERNEL_BUILDER( \
|
||||
Name("FlexPool").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
|
||||
FlexPoolOp<CPUDevice, T>) \
|
||||
REGISTER_KERNEL_BUILDER( \
|
||||
Name("FlexPoolGrad").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
|
||||
FlexPoolGradOp<CPUDevice, T>)
|
||||
|
||||
TF_CALL_float(REGISTER_FLEXPOOL_OP_CPU);
|
||||
#undef REGISTER_FLEXPOOL_OP_CPU
|
||||
|
||||
// Register the GPU kernels.
|
||||
#ifdef GOOGLE_CUDA
|
||||
|
||||
#define REGISTER_FLEXPOOL_OP_GPU(T) \
|
||||
REGISTER_KERNEL_BUILDER( \
|
||||
Name("FlexPool").Device(DEVICE_GPU).TypeConstraint<T>("T"), \
|
||||
FlexPoolOp<GPUDevice, T>) \
|
||||
REGISTER_KERNEL_BUILDER( \
|
||||
Name("FlexPoolGrad").Device(DEVICE_GPU).TypeConstraint<T>("T"), \
|
||||
FlexPoolGradOp<GPUDevice, T>)
|
||||
|
||||
TF_CALL_float(REGISTER_FLEXPOOL_OP_GPU);
|
||||
#undef REGISTER_FLEXPOOL_OP_GPU
|
||||
|
||||
#endif // GOOGLE_CUDA
|
||||
|
||||
} // namespace tensorflow
|
||||
@@ -0,0 +1,50 @@
|
||||
/* 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
|
||||
|
||||
#ifndef USER_OPS_KERNELS_FLEX_POOL_OP_H_
|
||||
#define USER_OPS_KERNELS_FLEX_POOL_OP_H_
|
||||
|
||||
#include "tensorflow/core/framework/op_kernel.h"
|
||||
|
||||
namespace tensorflow {
|
||||
class OpKernelContext;
|
||||
class Tensor;
|
||||
|
||||
using CPUDevice = Eigen::ThreadPoolDevice;
|
||||
using GPUDevice = Eigen::GpuDevice;
|
||||
} // namespace tensorflow
|
||||
|
||||
namespace tensorflow {
|
||||
namespace functor {
|
||||
|
||||
template <typename Device, typename Dtype>
|
||||
struct FlexPoolFunctor {
|
||||
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& features_,
|
||||
const Tensor& neighborhood_, Tensor* output_,
|
||||
Tensor* argmax_);
|
||||
};
|
||||
|
||||
template <typename Device, typename Dtype>
|
||||
struct FlexPoolGrad {
|
||||
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& features_,
|
||||
const Tensor& neighborhood_, const Tensor& topdiff_,
|
||||
const Tensor& argmax_, Tensor* grad_features_);
|
||||
};
|
||||
|
||||
} // namespace functor
|
||||
} // namespace tensorflow
|
||||
|
||||
#endif // USER_OPS_KERNELS_FLEX_POOL_OP_H_
|
||||
Reference in New Issue
Block a user