wassname/Flex-Convolution
1
0
Fork 0
mirror of https://github.com/wassname/Flex-Convolution.git synced 2026-09-10 11:40:26 +08:00
Code Issues Packages Projects Releases Wiki Activity

cherry pick files for public release

Browse Source
This commit is contained in:
Patrick Wieschollek
2018-09-23 12:44:08 +02:00
committed by PatWie
commit 5a9cd25aa9
32 changed files with 3947 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
user_ops/CMakeLists.txt
+40
View File
@@ -0,0 +1,40 @@
# University Tuebingen, 2018
# Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
cmake_minimum_required( VERSION 2.8 )
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
project( FlexConv )
list(APPEND CMAKE_MODULE_PATH ${PROJECT_SOURCE_DIR})
find_package(CUDA 9 REQUIRED)
find_package(TensorFlow REQUIRED)
if (DEFINED ENV{CUB_INC})
message(STATUS "Use Cuda-CUB from " $ENV{CUB_INC})
set(CUB_INC $ENV{CUB_INC})
else()
message(FATAL_ERROR "requires 'export CUB_INC=/path/to/cub")
endif()
# set necessary flags
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${SSE_FLAGS} -march=native -fopenmp -O3 -D_GLIBCXX_USE_CXX11_ABI=${TensorFlow_ABI}")
set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} -fPIC --shared -D_GLIBCXX_USE_CXX11_ABI=${TensorFlow_ABI}" )
set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -std=c++11 -O3 -Xptxas=-v --expt-relaxed-constexpr -D GOOGLE_CUDA=1 --gpu-architecture=sm_52 -D_GLIBCXX_USE_CXX11_ABI=${TensorFlow_ABI}" )
# quick fix for drone-ci
include_directories(SYSTEM "/usr/local/")
# fix cgtuebingen
include_directories(SYSTEM "/graphics/opt/opt_Ubuntu18.04/cuda/toolkit_9.2")
include_directories(SYSTEM ${CUB_INC})
include_directories(SYSTEM ${TensorFlow_INCLUDE_DIR})
include_directories(SYSTEM kernels)
add_tensorflow_gpu_operation("flex_conv")
add_tensorflow_gpu_operation("flex_deconv")
add_tensorflow_gpu_operation("flex_pool")
user_ops/FindTensorFlow.cmake
+273
View File
@@ -0,0 +1,273 @@
# Patrick Wieschollek, <mail@patwie.com>
# FindTENSORFLOW.cmake
# https://github.com/PatWie/tensorflow-cmake/blob/master/cmake/modules/FindTensorFlow.cmake
# -------------
#
# Find TensorFlow library and includes
#
# Automatically set variables have prefix "TensorFlow",
# while variables you need to specify have prefix "TENSORFLOW"
# This module will set the following variables in your project:
#
# ``TensorFlow_VERSION``
# exact TensorFlow version obtained from runtime
# ``TensorFlow_ABI``
# ABI specification of TensorFlow library obtained from runtime
# ``TensorFlow_INCLUDE_DIR``
# where to find tensorflow header files obtained from runtime
# ``TensorFlow_LIBRARY``
# the libraries to link against to use TENSORFLOW obtained from runtime
# ``TensorFlow_FOUND TRUE``
# If false, do not try to use TENSORFLOW.
#
# for some examples, you will need to specify on of the following paths
# ``TensorFlow_SOURCE_DIR``
# Path to source of TensorFlow, when env-var 'TENSORFLOW_SOURCE_DIR' is set and path exists
# ``TensorFlow_C_LIBRARY``
# Path to libtensorflow_cc.so (require env-var 'TENSORFLOW_BUILD_DIR')
#
#
# USAGE
# ------
# add "list(APPEND CMAKE_MODULE_PATH ${PROJECT_SOURCE_DIR}../../path/to/this/file)" to your project
#
# "add_tensorflow_gpu_operation" is a macro to compile a custom operation
#
# add_tensorflow_gpu_operation("<op-name>") expects the following files to exists:
# - kernels/<op-name>_kernel.cc
# - kernels/<op-name>_kernel_gpu.cu.cc (kernels/<op-name>_kernel.cu is supported as well)
# - kernels/<op-name>_op.cc
# - kernels/<op-name>_op.h
# - ops/<op-name>.cc
if(APPLE)
message(WARNING "This FindTensorflow.cmake is not tested on APPLE\n"
"Please report if this works\n"
"https://github.com/PatWie/tensorflow-cmake")
endif(APPLE)
if(WIN32)
message(WARNING "This FindTensorflow.cmake is not tested on WIN32\n"
"Please report if this works\n"
"https://github.com/PatWie/tensorflow-cmake")
endif(WIN32)
set(PYTHON_EXECUTABLE "python3" CACHE STRING "specify the python version TensorFlow is installed on.")
if(TensorFlow_FOUND)
# reuse cached variables
message(STATUS "Reuse cached information from TensorFlow ${TensorFlow_VERSION} ")
else()
message(STATUS "Detecting TensorFlow using ${PYTHON_EXECUTABLE}"
" (use -DPYTHON_EXECUTABLE=... otherwise)")
execute_process(
COMMAND ${PYTHON_EXECUTABLE} -c "import tensorflow as tf; print(tf.__version__); print(tf.__cxx11_abi_flag__); print(tf.sysconfig.get_include()); print(tf.sysconfig.get_lib() + '/libtensorflow_framework.so')"
OUTPUT_VARIABLE TF_INFORMATION_STRING
OUTPUT_STRIP_TRAILING_WHITESPACE
RESULT_VARIABLE retcode)
if(NOT "${retcode}" STREQUAL "0")
message(FATAL_ERROR "Detecting TensorFlow info - failed \n Did you installed TensorFlow?")
else()
message(STATUS "Detecting TensorFlow info - done")
endif()
string(REPLACE "\n" ";" TF_INFORMATION_LIST ${TF_INFORMATION_STRING})
list(GET TF_INFORMATION_LIST 0 TF_DETECTED_VERSION)
list(GET TF_INFORMATION_LIST 1 TF_DETECTED_ABI)
list(GET TF_INFORMATION_LIST 2 TF_DETECTED_INCLUDE_DIR)
list(GET TF_INFORMATION_LIST 3 TF_DETECTED_LIBRARY)
# set(TF_DETECTED_VERSION 1.8)
set(_packageName "TF")
if (DEFINED TF_DETECTED_VERSION)
string (REGEX MATCHALL "[0-9]+" _versionComponents "${TF_DETECTED_VERSION}")
list (LENGTH _versionComponents _len)
if (${_len} GREATER 0)
list(GET _versionComponents 0 TF_DETECTED_VERSION_MAJOR)
endif()
if (${_len} GREATER 1)
list(GET _versionComponents 1 TF_DETECTED_VERSION_MINOR)
endif()
if (${_len} GREATER 2)
list(GET _versionComponents 2 TF_DETECTED_VERSION_PATCH)
endif()
if (${_len} GREATER 3)
list(GET _versionComponents 3 TF_DETECTED_VERSION_TWEAK)
endif()
set (TF_DETECTED_VERSION_COUNT ${_len})
else()
set (TF_DETECTED_VERSION_COUNT 0)
endif()
# -- prevent pre 1.9 versions
# Note: TensorFlow 1.7 supported custom ops and all header files.
# TensorFlow 1.8 broke that promise and 1.9, 1.10 are fine again.
# This cmake-file is only tested against 1.9+.
if("${TF_DETECTED_VERSION}" VERSION_LESS "1.9")
message(FATAL_ERROR "Your installed TensorFlow version ${TF_DETECTED_VERSION} is too old.")
endif()
if(TF_FIND_VERSION_EXACT)
# User requested exact match of TensorFlow.
# TensorFlow release cycles are currently just depending on (major, minor)
# But we test against both.
set(_TensorFlow_TEST_VERSIONS
"${TF_FIND_VERSION_MAJOR}.${TF_FIND_VERSION_MINOR}.${TF_FIND_VERSION_PATCH}"
"${TF_FIND_VERSION_MAJOR}.${TF_FIND_VERSION_MINOR}")
else(TF_FIND_VERSION_EXACT)
# User requested not an exact TensorFlow version.
# However, only TensorFlow versions 1.9, 1.10 support all header files
# for custom ops.
set(_TensorFlow_KNOWN_VERSIONS ${TensorFlow_ADDITIONAL_VERSIONS}
"1.9" "1.9.0" "1.10" "1.10.0")
set(_TensorFlow_TEST_VERSIONS)
if(TF_FIND_VERSION)
set(_TF_FIND_VERSION_SHORT "${TF_FIND_VERSION_MAJOR}.${TF_FIND_VERSION_MINOR}")
# Select acceptable versions.
foreach(version ${_TensorFlow_KNOWN_VERSIONS})
if(NOT "${version}" VERSION_LESS "${TF_FIND_VERSION}")
# This version is high enough.
list(APPEND _TensorFlow_TEST_VERSIONS "${version}")
endif()
endforeach(version)
else(TF_FIND_VERSION)
# Any version is acceptable.
set(_TensorFlow_TEST_VERSIONS "${_TensorFlow_KNOWN_VERSIONS}")
endif(TF_FIND_VERSION)
endif()
# test all given versions
set(TensorFlow_FOUND FALSE)
FOREACH(_TensorFlow_VER ${_TensorFlow_TEST_VERSIONS})
if("${TF_DETECTED_VERSION_MAJOR}.${TF_DETECTED_VERSION_MINOR}" STREQUAL "${_TensorFlow_VER}")
# found appropriate version
set(TensorFlow_VERSION ${TF_DETECTED_VERSION})
set(TensorFlow_ABI ${TF_DETECTED_ABI})
set(TensorFlow_INCLUDE_DIR ${TF_DETECTED_INCLUDE_DIR})
set(TensorFlow_LIBRARY ${TF_DETECTED_LIBRARY})
set(TensorFlow_FOUND TRUE)
message(STATUS "Found TensorFlow: (found appropriate version \"${TensorFlow_VERSION}\")")
message(STATUS "TensorFlow-ABI is ${TensorFlow_ABI}")
message(STATUS "TensorFlow-INCLUDE_DIR is ${TensorFlow_INCLUDE_DIR}")
message(STATUS "TensorFlow-LIBRARY is ${TensorFlow_LIBRARY}")
add_definitions("-DTENSORFLOW_ABI=${TensorFlow_ABI}")
add_definitions("-DTENSORFLOW_VERSION=${TensorFlow_VERSION}")
break()
endif()
ENDFOREACH(_TensorFlow_VER)
if(NOT TensorFlow_FOUND)
message(FATAL_ERROR "Your installed TensorFlow version ${TF_DETECTED_VERSION_MAJOR}.${TF_DETECTED_VERSION_MINOR} is not supported\n"
"We tested against ${_TensorFlow_TEST_VERSIONS}")
endif(NOT TensorFlow_FOUND)
endif()
find_library(TensorFlow_C_LIBRARY
NAMES libtensorflow_cc.so
PATHS $ENV{TENSORFLOW_BUILD_DIR}
DOC "TensorFlow CC library." )
if(TensorFlow_C_LIBRARY)
message(STATUS "TensorFlow-CC-LIBRARY is ${TensorFlow_C_LIBRARY}")
else()
message(STATUS "No TensorFlow-CC-LIBRARY detected")
endif()
find_path(TensorFlow_SOURCE_DIR
NAMES
tensorflow/c
tensorflow/cc
tensorflow/core
tensorflow/core/framework
tensorflow/core/platform
tensorflow/python
third_party
PATHS $ENV{TENSORFLOW_SOURCE_DIR})
if(TensorFlow_SOURCE_DIR)
message(STATUS "TensorFlow-SOURCE-DIRECTORY is ${TensorFlow_SOURCE_DIR}")
else()
message(STATUS "No TensorFlow source repository detected")
endif()
macro(TensorFlow_REQUIRE_C_LIBRARY)
if(TensorFlow_C_LIBRARY)
else()
message(FATAL_ERROR "Project requires libtensorflow_cc.so, please specify the path in ENV-VAR 'TENSORFLOW_BUILD_DIR'")
endif()
endmacro()
macro(TensorFlow_REQUIRE_SOURCE)
if(TensorFlow_SOURCE_DIR)
else()
message(FATAL_ERROR "Project requires TensorFlow source directory, please specify the path in ENV-VAR 'TENSORFLOW_SOURCE_DIR'")
endif()
endmacro()
macro(add_tensorflow_cpu_operation op_name)
# Compiles a CPU-only operation without invoking NVCC
message(STATUS "will build custom TensorFlow operation \"${op_name}\" (CPU only)")
add_library(${op_name}_op SHARED kernels/${op_name}_op.cc kernels/${op_name}_kernel.cc ops/${op_name}.cc )
set_target_properties(${op_name}_op PROPERTIES PREFIX "")
target_link_libraries(${op_name}_op LINK_PUBLIC ${TensorFlow_LIBRARY})
endmacro()
macro(add_tensorflow_gpu_operation op_name)
# Compiles a CPU + GPU operation with invoking NVCC
message(STATUS "will build custom TensorFlow operation \"${op_name}\" (CPU+GPU)")
set(kernel_file "")
if(EXISTS "kernels/${op_name}_kernel.cu")
message(WARNING "you should rename your file ${op_name}_kernel.cu to ${op_name}_kernel_gpu.cu.cc")
set(kernel_file kernels/${op_name}_kernel.cu)
else()
set_source_files_properties(kernels/${op_name}_kernel_gpu.cu.cc PROPERTIES CUDA_SOURCE_PROPERTY_FORMAT OBJ)
set(kernel_file kernels/${op_name}_kernel_gpu.cu.cc)
endif()
cuda_add_library(${op_name}_op_cu SHARED ${kernel_file})
set_target_properties(${op_name}_op_cu PROPERTIES PREFIX "")
add_library(${op_name}_op SHARED kernels/${op_name}_op.cc kernels/${op_name}_kernel.cc ops/${op_name}.cc )
set_target_properties(${op_name}_op PROPERTIES PREFIX "")
set_target_properties(${op_name}_op PROPERTIES COMPILE_FLAGS "-DGOOGLE_CUDA")
target_link_libraries(${op_name}_op LINK_PUBLIC ${op_name}_op_cu ${TensorFlow_LIBRARY})
endmacro()
# simplify TensorFlow dependencies
add_library(TensorFlow_DEP INTERFACE)
TARGET_INCLUDE_DIRECTORIES(TensorFlow_DEP INTERFACE ${TensorFlow_SOURCE_DIR})
TARGET_INCLUDE_DIRECTORIES(TensorFlow_DEP INTERFACE ${TensorFlow_INCLUDE_DIR})
TARGET_LINK_LIBRARIES(TensorFlow_DEP INTERFACE -Wl,--allow-multiple-definition -Wl,--whole-archive ${TensorFlow_C_LIBRARY} -Wl,--no-whole-archive)
TARGET_LINK_LIBRARIES(TensorFlow_DEP INTERFACE -Wl,--allow-multiple-definition -Wl,--whole-archive ${TensorFlow_LIBRARY} -Wl,--no-whole-archive)
include(FindPackageHandleStandardArgs)
find_package_handle_standard_args(
TENSORFLOW
FOUND_VAR TENSORFLOW_FOUND
REQUIRED_VARS
TensorFlow_LIBRARY
TensorFlow_INCLUDE_DIR
VERSION_VAR
TensorFlow_VERSION
)
mark_as_advanced(TF_INFORMATION_STRING TF_DETECTED_VERSION TF_DETECTED_VERSION_MAJOR TF_DETECTED_VERSION_MINOR TF_DETECTED_VERSION TF_DETECTED_ABI
TF_DETECTED_INCLUDE_DIR TF_DETECTED_LIBRARY
TensorFlow_C_LIBRARY TensorFlow_LIBRARY TensorFlow_SOURCE_DIR TensorFlow_INCLUDE_DIR TensorFlow_ABI)
SET(TensorFlow_INCLUDE_DIR ${TensorFlow_INCLUDE_DIR} CACHE PATH "path to tensorflow header files")
SET(TensorFlow_VERSION ${TensorFlow_VERSION} CACHE INTERNAL "The Python executable Version")
SET(TensorFlow_ABI ${TensorFlow_ABI} CACHE STRING "The Python executable Version")
SET(TensorFlow_LIBRARY ${TensorFlow_LIBRARY} CACHE PATH "The Python executable Version")
SET(TensorFlow_FOUND ${TensorFlow_FOUND} CACHE BOOL "The Python executable Version")
user_ops/PointTestCase.py
+130
View File
@@ -0,0 +1,130 @@
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# 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
import numpy as np
import tensorflow as tf
np.random.seed(42)
tf.set_random_seed(42)
class FakePointCloud(object):
"""docstring for FakePointCloud"""
def __init__(self, B, N, K, Din, Dout, Dp, N2, scaling=1):
super(FakePointCloud, self).__init__()
assert K < N
self.B = B
self.N = N
self.K = K
self.Din = Din
self.Dout = Dout
self.Dp = Dp
self.N2 = N2
def expected_feature_shape(self):
return [self.B, self.Din, self.N]
def expected_output_shape(self):
return [self.B, self.Dout, self.N]
def random_values(shape, human_readable=False):
"""Return random values within range [-10, 10] and precision 2
"""
length = np.prod(shape)
return np.arange(length).astype(np.float32).reshape(shape) / float(length)
def summary(numeric_grad, graph_grad, name, eps=0.001, max_outputs=20):
a, b = numeric_grad.flatten(), graph_grad.flatten()
print("summary: %s" % name)
print("\ttheirs\t\tours\t\tabs-diff")
for i in range(np.prod(numeric_grad.shape)):
if np.abs(a[i] - b[i]) > eps and max_outputs > 0:
print('%i\t%f\t%f\t%f' % (i, a[i], b[i], np.abs(a[i] - b[i])))
max_outputs -= 1
if max_outputs == 20:
for i in range(max_outputs):
print('%i\t%f\t%f\t%f' % (i, a[i], b[i], np.abs(a[i] - b[i])))
# print( np.stack([numeric_grad, graph_grad], axis=-1)
print("%s - abs-diff (sum): " % name, np.abs(
graph_grad - numeric_grad).sum())
print("%s - abs-diff (max): " % name, np.abs(
graph_grad - numeric_grad).max())
print("%s - abs-diff (mean): " % name, np.abs(
graph_grad - numeric_grad).mean())
# TestPointCloud(B, N, K, Din, Dout, Dp, N2)
# TPC = FakePointCloud(2, 32, 16, 5, 6, 3, 16)
# TPC = FakePointCloud(2, 16, 8, 5, 6, 3, 8)
TPC = FakePointCloud(2, 32, 4, 2, 6, 3, 16)
# TPC = FakePointCloud(2, 64, 8, 1, 6, 3, 64)
class PointTestCase(tf.test.TestCase):
def __init__(self, methodName="runTest", data=None):
if data is None:
data = TPC
self.position = random_values([data.B, data.Dp, data.N])
self.features = random_values([data.B, data.Din, data.N])
# make sure, each neighbor hood has no duplicates and first entry is point n
# THIS IS IMPORTANT!!
self.neighborhood = np.zeros((data.B, data.K, data.N), dtype=np.int32)
for b in range(data.B):
for n in range(data.N):
x = np.arange(data.N)
# does not support axis, hence the loop
np.random.shuffle(x)
offset = np.argwhere(x == n)[0][0]
# roll array such that n is first entry
x = np.roll(x, -offset)
self.neighborhood[b, :, n] = x[:data.K].astype(np.int32)
self.neighborhood_ds = np.zeros((data.B, data.K, data.N2), dtype=np.int32)
for b in range(data.B):
for n in range(data.N2):
x = np.arange(data.N2)
# does not support axis, hence the loop
np.random.shuffle(x)
offset = np.argwhere(x == n)[0][0]
# roll array such that n is first entry
x = np.roll(x, -offset)
self.neighborhood[b, :, n] = x[:data.K].astype(np.int32)
self.theta = random_values([1, data.Dp, data.Din, data.Dout])
self.bias = random_values([data.Din, data.Dout])
super(PointTestCase, self).__init__(methodName)
def init_ops(self):
# needs to be called in each method, otherwise graph is empty
# probably tf.reset_graph between calls
self.features_op = tf.convert_to_tensor(self.features)
self.position_op = tf.convert_to_tensor(self.position)
self.neighborhood_op = tf.convert_to_tensor(self.neighborhood)
self.neighborhood_ds_op = tf.convert_to_tensor(self.neighborhood_ds)
self.theta_op = tf.convert_to_tensor(self.theta)
self.bias_op = tf.convert_to_tensor(self.bias)
user_ops/__init__.py
+176
View File
@@ -0,0 +1,176 @@
# Copyright 2018 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.
# ============================================================================
"""Tensorflow op performing flex convolution operation."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import tensorflow as tf
from tensorflow.python.framework import ops
from tensorflow.contrib.util import loader
from tensorflow.python.platform import resource_loader
_flex_convolution_op_so = loader.load_op_library(
resource_loader.get_path_to_datafile("flex_conv_op.so"))
_flex_pooling_op_so = loader.load_op_library(
resource_loader.get_path_to_datafile("flex_pool_op.so"))
_flex_deconvolution_op_so = loader.load_op_library(
resource_loader.get_path_to_datafile("flex_deconv_op.so"))
# undocumented version
flex_conv = _flex_convolution_op_so.flex_conv
flex_conv_grad = _flex_convolution_op_so.flex_conv_grad
flex_pool = _flex_pooling_op_so.flex_pool
flex_pool_grad = _flex_pooling_op_so.flex_pool_grad
flex_deconv = _flex_deconvolution_op_so.flex_deconv
flex_deconv_grad = _flex_deconvolution_op_so.flex_deconv_grad
# pylint: disable=redefined-builtin
def flex_convolution(features,
position,
neighborhood,
theta,
bias,
name=None):
"""Flex-Convolution computation.
Computes a convolution over arbitrary neighborhoods with elements of
arbitrary positions:
output(c', l) = sum_{c} sum_{l'} w(c, l, l') * f(c, l')
Args:
features: A `Tensor` of the format [B, Din, N].
position: A `Tensor` of the format [B, Dp, N].
neighborhood: A `Tensor` of the format [B, K, N] (tf.int32).
theta: A `Tensor` of the format [1, Dp, Din, Dout].
bias: A `Tensor` of the format [Din, Dout].
name: A name for the operation (optional).
Returns:
A `Tensor` of the format [B, Dout, N].
"""
with ops.name_scope(name, "flex_convolution"):
return flex_conv(features, theta, bias, neighborhood, position)
@ops.RegisterGradient("FlexConv")
def _FlexConvGrad(op, *grads): # noqa
features = ops.convert_to_tensor(op.inputs[0])
theta = ops.convert_to_tensor(op.inputs[1])
bias = ops.convert_to_tensor(op.inputs[2])
neighborhood = ops.convert_to_tensor(op.inputs[3], dtype=tf.int32)
positions = ops.convert_to_tensor(op.inputs[4])
topdiff = ops.convert_to_tensor(grads[0])
df, dt, db = flex_conv_grad(
features, theta, bias, neighborhood, positions, topdiff)
df = ops.convert_to_tensor(df, name='gradient_features')
dt = ops.convert_to_tensor(dt, name='gradient_theta')
db = ops.convert_to_tensor(db, name='gradient_bias')
return [df, dt, db, None, None]
# pylint: disable=redefined-builtin
def flex_pooling(features,
neighborhood,
name=None):
"""Flex-Pooling computation.
Computes a pooling over arbitrary neighborhoods:
output(n) = max_l' f(l')
Args:
features: A `Tensor` of the format [B, D, N].
neighborhood: A `Tensor` of the format [B, K, N] (tf.int32).
name: A name for the operation (optional).
Returns:
A `Tensor` of the format [B, D, N] containing the max values.
A `Tensor` of the format [B, D, N] containing the max indicies.
"""
with ops.name_scope(name, "flex_pooling"):
return flex_pool(features, neighborhood)
@ops.RegisterGradient("FlexPool")
def _FlexPoolGrad(op, *grads): # noqa
features = ops.convert_to_tensor(op.inputs[0])
neighborhood = ops.convert_to_tensor(op.inputs[1])
argmax = ops.convert_to_tensor(op.outputs[1])
topdiff = ops.convert_to_tensor(grads[0])
df = flex_pool_grad(features, neighborhood, topdiff, argmax)
df = ops.convert_to_tensor(df, name='gradient_features')
return [df, None]
# pylint: disable=redefined-builtin
def flex_convolution_transpose(features,
position,
neighborhood,
theta,
bias,
name=None):
"""Flex-Convolution computation.
Computes a tranposed convolution over arbitrary neighborhoods with elements of
arbitrary positions.
Args:
features: A `Tensor` of the format [B, Din, N].
position: A `Tensor` of the format [B, Dp, N].
neighborhood: A `Tensor` of the format [B, K, N] (tf.int32).
theta: A `Tensor` of the format [1, Dp, Din, Dout].
bias: A `Tensor` of the format [Din, Dout].
name: A name for the operation (optional).
Returns:
A `Tensor` of the format [B, Dout, N].
"""
with ops.name_scope(name, "flex_convolution_transpose"):
return flex_deconv(features, theta, bias, neighborhood, position)
@ops.RegisterGradient("FlexDeconv")
def _FlexDeconvGrad(op, *grads): # noqa
features = ops.convert_to_tensor(op.inputs[0])
theta = ops.convert_to_tensor(op.inputs[1])
bias = ops.convert_to_tensor(op.inputs[2])
neighborhood = ops.convert_to_tensor(op.inputs[3], dtype=tf.int32)
positions = ops.convert_to_tensor(op.inputs[4])
topdiff = ops.convert_to_tensor(grads[0])
df, dt, db = flex_deconv_grad(
features, theta, bias, neighborhood, positions, topdiff)
df = ops.convert_to_tensor(df, name='gradient_features')
dt = ops.convert_to_tensor(dt, name='gradient_theta')
db = ops.convert_to_tensor(db, name='gradient_bias')
return [df, dt, db, None, None]
user_ops/kernels/flex_conv_kernel.cc
+173
View File
@@ -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
user_ops/kernels/flex_conv_kernel_gpu.cu.cc
+532
View File
@@ -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
user_ops/kernels/flex_conv_op.cc
+122
View File
@@ -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
user_ops/kernels/flex_conv_op.h
+53
View File
@@ -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_
user_ops/kernels/flex_deconv_kernel.cc
+173
View File
@@ -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
user_ops/kernels/flex_deconv_kernel_gpu.cu.cc
+227
View File
@@ -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
user_ops/kernels/flex_deconv_op.cc
+103
View File
@@ -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
user_ops/kernels/flex_deconv_op.h
+53
View File
@@ -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_
user_ops/kernels/flex_pool_kernel.cc
+102
View File
@@ -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
user_ops/kernels/flex_pool_kernel_gpu.cu.cc
+171
View File
@@ -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
user_ops/kernels/flex_pool_op.cc
+113
View File
@@ -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
user_ops/kernels/flex_pool_op.h
+50
View File
@@ -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_
user_ops/ops/flex_conv.cc
+133
View File
@@ -0,0 +1,133 @@
/* 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 "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"
namespace tensorflow {
using ::tensorflow::shape_inference::InferenceContext;
using ::tensorflow::shape_inference::ShapeHandle;
REGISTER_OP("FlexConv")
.Input("features: T")
.Input("theta: T")
.Input("bias: T")
.Input("neighborhood: int32")
.Input("position: T")
.Output("output: T")
.Attr("T: realnumbertype")
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
const auto features = c->input(0);
const auto theta = c->input(1);
const auto bias = c->input(2);
const auto neighborhood = c->input(3);
const auto position = c->input(4);
// we require the input to have 4 axes
::tensorflow::shape_inference::ShapeHandle shape_hnd;
TF_RETURN_IF_ERROR(c->WithRank(features, 3, &shape_hnd)); // B x Din x Ng
TF_RETURN_IF_ERROR(
c->WithRank(theta, 4, &shape_hnd)); // 1 x Dp x Din x Dout
TF_RETURN_IF_ERROR(c->WithRank(bias, 2, &shape_hnd)); // Din x Dout
TF_RETURN_IF_ERROR(
c->WithRank(neighborhood, 3, &shape_hnd)); // B x K x N
TF_RETURN_IF_ERROR(c->WithRank(position, 3, &shape_hnd)); // B x Dp x Ng
shape_inference::DimensionHandle merged;
// assert B equal
TF_RETURN_IF_ERROR(
c->Merge(c->Dim(features, 0), c->Dim(neighborhood, 0), &merged));
TF_RETURN_IF_ERROR(
c->Merge(c->Dim(features, 0), c->Dim(position, 0), &merged));
// assert Ng equal
TF_RETURN_IF_ERROR(
c->Merge(c->Dim(features, 2), c->Dim(neighborhood, 2),
&merged)); // TODO(fabi?) fix global access and remove
TF_RETURN_IF_ERROR(
c->Merge(c->Dim(features, 2), c->Dim(position, 2), &merged));
// assert Dp equal
TF_RETURN_IF_ERROR(
c->Merge(c->Dim(theta, 1), c->Dim(position, 1), &merged));
// assert Dout equal
TF_RETURN_IF_ERROR(c->Merge(c->Dim(theta, 3), c->Dim(bias, 1), &merged));
// assert Din equal
TF_RETURN_IF_ERROR(
c->Merge(c->Dim(features, 1), c->Dim(theta, 2), &merged));
TF_RETURN_IF_ERROR(
c->Merge(c->Dim(features, 1), c->Dim(bias, 0), &merged));
// specify output-shape
auto B = c->Dim(features, 0);
auto Dout = c->Dim(bias, 1);
auto N = c->Dim(neighborhood, 2);
c->set_output(0, c->MakeShape({B, Dout, N}));
return Status::OK();
})
.Doc(R"doc(
Apply Sparse Convolution to inputs.
This applies a convolution to a neighborhood of inputs. The formula for computing the output is as follows:
`output_j = \sum_i w(x_i, x0) * x_i`
`w(x_i, x0) = ??`
features: each feature description for each point [B, Din, N].
theta: parameters for kernel function [1, Dp, Din, Dout].
bias: bias for kernel function [Din, Dout].
neighborhood: all K nearest neighbors [B, K, N].
position: each datapoint in 3d space [B, Dp, N].
output: each feature description for each point [B, Dout, N].
)doc");
REGISTER_OP("FlexConvGrad")
.Input("features: T")
.Input("theta: T")
.Input("bias: T")
.Input("neighborhood: int32")
.Input("position: T")
.Input("gradients: T")
.Output("grad_features: T")
.Output("grad_theta: T")
.Output("grad_bias: T")
.Attr("T: realnumbertype")
.SetShapeFn([](InferenceContext* c) {
c->set_output(0, c->input(0)); // features
c->set_output(1, c->input(1)); // theta
c->set_output(2, c->input(2)); // bias
return ::tensorflow::Status::OK();
})
.Doc(R"doc(
Returns gradients of Sparse Convolution to inputs.
gradients: topdiff[B, N, Dout].
neighborhood: all K nearest neighbors [B, K, N].
position: each datapoint in 3d space [B, Dp, N].
features: each feature description for each point [B, Din, N].
theta: parameters for kernel function [1, Dp, Din, Dout].
bias: bias for kernel function [Din, Dout].
grad_features: gradient to each feature description for each point [B, N, Din].
grad_theta: gradient to parameters for kernel function [1, Dp, Din, Dout].
grad_bias: gradient to bias for kernel function [Din, Dout].
)doc");
} // namespace tensorflow
user_ops/ops/flex_deconv.cc
+129
View File
@@ -0,0 +1,129 @@
/* 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 "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"
namespace tensorflow {
using ::tensorflow::shape_inference::InferenceContext;
using ::tensorflow::shape_inference::ShapeHandle;
REGISTER_OP("FlexDeconv")
.Input("features: T")
.Input("theta: T")
.Input("bias: T")
.Input("neighborhood: int32")
.Input("position: T")
.Output("output: T")
.Attr("T: realnumbertype")
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
const auto features = c->input(0);
const auto theta = c->input(1);
const auto bias = c->input(2);
const auto neighborhood = c->input(3);
const auto position = c->input(4);
// we require the input to have 4 axes
::tensorflow::shape_inference::ShapeHandle shape_hnd;
TF_RETURN_IF_ERROR(c->WithRank(features, 3, &shape_hnd)); // B x Din x N
TF_RETURN_IF_ERROR(
c->WithRank(theta, 4, &shape_hnd)); // 1 x Dp x Din x Dout
TF_RETURN_IF_ERROR(c->WithRank(bias, 2, &shape_hnd)); // Din x Dout
TF_RETURN_IF_ERROR(
c->WithRank(neighborhood, 3, &shape_hnd)); // B x K x N
TF_RETURN_IF_ERROR(c->WithRank(position, 3, &shape_hnd)); // B x Dp x N
shape_inference::DimensionHandle merged;
// assert B equal
TF_RETURN_IF_ERROR(
c->Merge(c->Dim(features, 0), c->Dim(neighborhood, 0), &merged));
TF_RETURN_IF_ERROR(
c->Merge(c->Dim(features, 0), c->Dim(position, 0), &merged));
// assert N equal
TF_RETURN_IF_ERROR(
c->Merge(c->Dim(features, 2), c->Dim(neighborhood, 2), &merged));
TF_RETURN_IF_ERROR(
c->Merge(c->Dim(features, 2), c->Dim(position, 2), &merged));
// assert Dp equal
TF_RETURN_IF_ERROR(
c->Merge(c->Dim(theta, 1), c->Dim(position, 1), &merged));
// assert Dout equal
TF_RETURN_IF_ERROR(c->Merge(c->Dim(theta, 3), c->Dim(bias, 1), &merged));
// assert Din equal
TF_RETURN_IF_ERROR(
c->Merge(c->Dim(features, 1), c->Dim(theta, 2), &merged));
TF_RETURN_IF_ERROR(
c->Merge(c->Dim(features, 1), c->Dim(bias, 0), &merged));
// specify output-shape
auto B = c->Dim(features, 0);
auto Dout = c->Dim(bias, 1);
auto N = c->Dim(neighborhood, 2);
c->set_output(0, c->MakeShape({B, Dout, N}));
return Status::OK();
})
.Doc(R"doc(
Apply Sparse Convolution to inputs.
This applies a deconvolution to a neighborhood of inputs. F_i spread its value to all neighbors.
features: each feature description for each point [B, Din, N].
theta: parameters for kernel function [1, Dp, Din, Dout].
bias: bias for kernel function [Din, Dout].
neighborhood: all K nearest neighbors [B, K, N].
position: each datapoint in 3d space [B, Dp, N].
output: each feature description for each point [B, Dout, N].
)doc");
REGISTER_OP("FlexDeconvGrad")
.Input("features: T")
.Input("theta: T")
.Input("bias: T")
.Input("neighborhood: int32")
.Input("position: T")
.Input("gradients: T")
.Output("grad_features: T")
.Output("grad_theta: T")
.Output("grad_bias: T")
.Attr("T: realnumbertype")
.SetShapeFn([](InferenceContext* c) {
c->set_output(0, c->input(0)); // features
c->set_output(1, c->input(1)); // theta
c->set_output(2, c->input(2)); // bias
return ::tensorflow::Status::OK();
})
.Doc(R"doc(
Returns gradients of Sparse Deconvolution to inputs.
gradients: topdiff[B, N, Dout].
neighborhood: all K nearest neighbors [B, K, N].
position: each datapoint in 3d space [B, Dp, N].
features: each feature description for each point [B, Din, N].
theta: parameters for kernel function [1, Dp, Din, Dout].
bias: bias for kernel function [Din, Dout].
grad_features: gradient to each feature description for each point [B, N, Din].
grad_theta: gradient to parameters for kernel function [1, Dp, Din, Dout].
grad_bias: gradient to bias for kernel function [Din, Dout].
)doc");
} // namespace tensorflow
user_ops/ops/flex_pool.cc
+93
View File
@@ -0,0 +1,93 @@
/* 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 "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"
namespace tensorflow {
using ::tensorflow::shape_inference::InferenceContext;
using ::tensorflow::shape_inference::ShapeHandle;
REGISTER_OP("FlexPool")
.Input("features: T")
.Input("neighborhood: int32")
.Output("output: T")
.Output("argmax: int32")
.Attr("T: realnumbertype")
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
const auto features = c->input(0);
const auto neighborhood = c->input(1);
// we require the input to have 3 axes
::tensorflow::shape_inference::ShapeHandle shape_hnd;
TF_RETURN_IF_ERROR(c->WithRank(features, 3, &shape_hnd)); // B x D x N
TF_RETURN_IF_ERROR(
c->WithRank(neighborhood, 3, &shape_hnd)); // B x K x N
shape_inference::DimensionHandle merged;
// assert B equal
TF_RETURN_IF_ERROR(
c->Merge(c->Dim(features, 0), c->Dim(neighborhood, 0), &merged));
// assert N equal
TF_RETURN_IF_ERROR(
c->Merge(c->Dim(features, 2), c->Dim(neighborhood, 2), &merged));
// specify output-shape
auto B = c->Dim(features, 0);
auto D = c->Dim(features, 1);
auto N = c->Dim(features, 2);
c->set_output(0, c->MakeShape({B, D, N}));
c->set_output(1, c->MakeShape({B, D, N}));
return Status::OK();
})
.Doc(R"doc(
Apply Sparse Pooling to inputs.
This applies a max-pooling to a neighborhood of inputs. T
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]. This is needed for gradients.
)doc");
REGISTER_OP("FlexPoolGrad")
.Input("features: T")
.Input("neighborhood: int32")
.Input("gradients: T")
.Input("argmax: int32")
.Output("grad_features: T")
.Attr("T: realnumbertype")
.SetShapeFn([](InferenceContext* c) {
c->set_output(0, c->input(0)); // features
return ::tensorflow::Status::OK();
})
.Doc(R"doc(
Returns gradients of MaxPool to inputs.
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].
)doc");
} // namespace tensorflow
user_ops/profile_flexconv.py
+116
View File
@@ -0,0 +1,116 @@
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# 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
import numpy as np
import tensorflow as tf
from PointTestCase import FakePointCloud, random_values
from __init__ import flex_conv
"""
export LD_LIBRARY_PATH=/graphics/opt/opt_Ubuntu16.04/cuda/toolkit_9.0/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH
"""
np.random.seed(42)
tf.set_random_seed(42)
N = 8
TPC = FakePointCloud(8, 4096, N, 64, 64, 2, 1, N)
class PointTestCase(object):
def __init__(self, data):
self.position = random_values([data.B, data.Dp, data.N])
self.features = random_values([data.B, data.Din, data.N])
# make sure, each neighbor hood has no duplicates and first entry is point n
# THIS IS IMPORTANT!!
self.neighborhood = np.zeros((data.B, data.K, data.N), dtype=np.int32)
for b in range(data.B):
for n in range(data.N):
x = np.arange(data.N)
# does not support axis, hence the loop
np.random.shuffle(x)
offset = np.argwhere(x == n)[0][0]
# roll array such that n is first entry
x = np.roll(x, -offset)
self.neighborhood[b, :, n] = x[:data.K].astype(np.int32)
self.neighborhood_ds = np.zeros((data.B, data.K, data.N2), dtype=np.int32)
for b in range(data.B):
for n in range(data.N2):
x = np.arange(data.N2)
# does not support axis, hence the loop
np.random.shuffle(x)
offset = np.argwhere(x == n)[0][0]
# roll array such that n is first entry
x = np.roll(x, -offset)
self.neighborhood[b, :, n] = x[:data.K].astype(np.int32)
self.theta = random_values([data.Degree, data.Dp, data.Din, data.Dout])
self.bias = random_values([data.Din, data.Dout])
def init_ops(self):
self.features_op = tf.Variable(self.features, name='f')
self.position_op = tf.Variable(self.position, name='p')
self.neighborhood_op = tf.Variable(self.neighborhood, name='n')
self.neighborhood_ds_op = tf.Variable(self.neighborhood_ds, name='m')
self.theta_op = tf.Variable(self.theta, name='t')
self.bias_op = tf.Variable(self.bias, name='b')
TestCase = PointTestCase(data=TPC)
TestCase.init_ops()
forward_op = flex_conv(TestCase.features_op,
TestCase.theta_op, TestCase.bias_op, TestCase.neighborhood_op,
TestCase.position_op, degree=1)
builder = tf.profiler.ProfileOptionBuilder
opts = builder(builder.time_and_memory()).order_by('micros').build()
with tf.contrib.tfprof.ProfileContext('./.profiling_outputs/flex_conv') as pctx:
with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as sess:
sess.run(tf.global_variables_initializer())
back_prop = tf.gradients(forward_op, [TestCase.theta_op, TestCase.bias_op, TestCase.features_op])
back_prop2 = tf.gradients(forward_op, [TestCase.theta_op, TestCase.bias_op,
TestCase.features_op, TestCase.position_op])
# warmup
for i in range(2):
actual = sess.run([forward_op, back_prop])
# benchmark
for i in range(10):
pctx.trace_next_step()
pctx.dump_next_step()
_ = sess.run([forward_op])
pctx.profiler.profile_operations(options=opts)
for i in range(10):
pctx.trace_next_step()
pctx.dump_next_step()
_ = sess.run([back_prop])
pctx.profiler.profile_operations(options=opts)
user_ops/test_flex_convolution.py
+109
View File
@@ -0,0 +1,109 @@
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# 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
from PointTestCase import TPC, PointTestCase, summary
import tensorflow as tf
from __init__ import flex_convolution
class FlexConvTest(PointTestCase):
def __init__(self, methodName="runTest"):
super(FlexConvTest, self).__init__(methodName)
def _forward(self, use_gpu=False, force_gpu=False):
self.init_ops()
with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
actual_op = flex_convolution(self.features_op,
self.position_op, self.neighborhood_op,
self.theta_op, self.bias_op)
actual = sess.run(actual_op)
return actual
def test_forward(self):
cpu = self._forward(use_gpu=False)
gpu = self._forward(use_gpu=True)
self.assertAllClose(cpu, gpu, 1e-5, 1e-5)
def _backward_features(self, use_gpu=False):
self.init_ops()
with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu):
actual_op = flex_convolution(self.features_op,
self.position_op, self.neighborhood_op,
self.theta_op, self.bias_op)
graph_features_grad, num_features_grad = tf.test.compute_gradient(
[self.features_op], [self.features.shape], actual_op,
TPC.expected_output_shape())[0]
summary(num_features_grad, graph_features_grad, 'self.features')
err = tf.test.compute_gradient_error([self.features_op],
[self.features.shape],
actual_op, TPC.expected_output_shape())
self.assertLess(err, 1e-2)
def _backward_bias(self, use_gpu=False):
self.init_ops()
with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu):
actual_op = flex_convolution(self.features_op,
self.position_op, self.neighborhood_op,
self.theta_op, self.bias_op)
graph_bias_grad, num_bias_grad = tf.test.compute_gradient(
[self.bias_op], [self.bias.shape], actual_op,
TPC.expected_output_shape())[0]
summary(num_bias_grad, graph_bias_grad, 'self.bias')
err = tf.test.compute_gradient_error([self.bias_op],
[self.bias.shape], actual_op,
TPC.expected_output_shape())
self.assertLess(err, 1e-2)
def _backward_theta(self, use_gpu=False):
self.init_ops()
with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu):
actual_op = flex_convolution(self.features_op,
self.position_op, self.neighborhood_op,
self.theta_op, self.bias_op)
graph_theta_grad, num_theta_grad = tf.test.compute_gradient(
[self.theta_op], [self.theta.shape], actual_op,
TPC.expected_output_shape())[0]
summary(num_theta_grad, graph_theta_grad, 'self.theta')
err = tf.test.compute_gradient_error([self.theta_op],
[self.theta.shape], actual_op,
TPC.expected_output_shape())
self.assertLess(err, 1e-2)
def test_backward_features(self):
self._backward_features(use_gpu=False)
self._backward_features(use_gpu=True)
def test_backward_bias(self):
self._backward_bias(use_gpu=False)
self._backward_bias(use_gpu=True)
def test_backward_theta(self):
self._backward_theta(use_gpu=False)
self._backward_theta(use_gpu=True)
if __name__ == '__main__':
tf.test.main()
user_ops/test_flex_convolution_transpose.py
+116
View File
@@ -0,0 +1,116 @@
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# 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
from PointTestCase import TPC, PointTestCase, summary
import tensorflow as tf
from __init__ import flex_convolution_transpose
class FlexDeconvTest(PointTestCase):
def __init__(self, methodName="runTest"):
super(FlexDeconvTest, self).__init__(methodName)
def _forward(self, use_gpu=False, force_gpu=False):
self.init_ops()
with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
actual = flex_convolution_transpose(self.features_op,
self.position_op, self.neighborhood_op,
self.theta_op, self.bias_op)
actual = sess.run(actual)
return actual
def test_forward(self):
cpu = self._forward(use_gpu=False)
gpu = self._forward(use_gpu=True)
summary(cpu, gpu, 'self.features')
self.assertAllClose(cpu, gpu, rtol=1e-05)
def _backward_features(self, use_gpu=False):
self.init_ops()
with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu):
actual_op = flex_convolution_transpose(self.features_op,
self.position_op,
self.neighborhood_op,
self.theta_op, self.bias_op)
graph_features_grad, num_features_grad = tf.test.compute_gradient(
[self.features_op], [self.features.shape], actual_op,
TPC.expected_output_shape())[0]
summary(num_features_grad, graph_features_grad,
'self.features', max_outputs=1200000)
err = tf.test.compute_gradient_error([self.features_op],
[self.features.shape],
actual_op, TPC.expected_output_shape())
self.assertLess(err, 1e-2)
def _backward_bias(self, use_gpu=False):
self.init_ops()
with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu):
actual_op = flex_convolution_transpose(self.features_op,
self.position_op,
self.neighborhood_op,
self.theta_op, self.bias_op)
graph_bias_grad, num_bias_grad = tf.test.compute_gradient(
[self.bias_op], [self.bias.shape], actual_op,
TPC.expected_output_shape())[0]
summary(num_bias_grad, graph_bias_grad, 'self.bias')
err = tf.test.compute_gradient_error([self.bias_op],
[self.bias.shape], actual_op,
TPC.expected_output_shape())
self.assertLess(err, 1e-2)
def _backward_theta(self, use_gpu=False):
self.init_ops()
with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu):
actual_op = flex_convolution_transpose(self.features_op,
self.position_op,
self.neighborhood_op,
self.theta_op, self.bias_op)
graph_theta_grad, num_theta_grad = tf.test.compute_gradient(
[self.theta_op], [self.theta.shape], actual_op,
TPC.expected_output_shape())[0]
summary(num_theta_grad, graph_theta_grad, 'self.theta')
err = tf.test.compute_gradient_error([self.theta_op],
[self.theta.shape], actual_op,
TPC.expected_output_shape())
self.assertLess(err, 1e-2)
def test_backward_features(self):
self._backward_features(use_gpu=False)
self._backward_features(use_gpu=True)
def test_backward_bias(self):
self._backward_bias(use_gpu=False)
self._backward_bias(use_gpu=True)
def test_backward_theta(self):
self._backward_theta(use_gpu=False)
self._backward_theta(use_gpu=True)
if __name__ == '__main__':
tf.test.main()
user_ops/test_flex_pooling.py
+87
View File
@@ -0,0 +1,87 @@
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# 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
from PointTestCase import PointTestCase
import tensorflow as tf
import numpy as np
from __init__ import flex_pooling
class FlexPoolTest(PointTestCase):
def __init__(self, methodName="runTest"):
super(FlexPoolTest, self).__init__(methodName)
def _forward(self, use_gpu=False):
self.init_ops()
with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
actual_op, winner_op = flex_pooling(self.features_op, self.neighborhood_op)
actual = sess.run(actual_op)
return actual
def test_forward(self):
cpu = self._forward(use_gpu=False)
gpu = self._forward(use_gpu=True)
self.assertAllClose(cpu, gpu)
def test_backward(self):
cpu, winner_cpu = self._backward(use_gpu=False)
gpu, winner_gpu = self._backward(use_gpu=True)
self.assertAllClose(cpu, gpu)
self.assertAllClose(cpu[winner_cpu == 0].sum(), 0)
self.assertAllClose(gpu[winner_gpu == 0].sum(), 0)
def _backward(self, use_gpu=False):
self.init_ops()
with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
actual_op, winner_op = flex_pooling(
self.features_op,
self.neighborhood_op)
graph_features_grad = tf.gradients(actual_op, [self.features_op])[0]
dx, winner = sess.run([graph_features_grad, winner_op])
return dx, winner
def _simple_backward(self, use_gpu=False):
# BN
x = np.array([[[1], [2], [5], [3]]]).transpose(0, 2, 1)
n = np.array([[[0, 1, 2, 3], [1, 2, 3, 0], [2, 3, 0, 1], [3, 0, 1, 2, ]]]).transpose(0, 2, 1)
x = tf.convert_to_tensor(x.astype(np.float32))
n = tf.convert_to_tensor(n.astype(np.int32))
with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
actual_op, winner_op = flex_pooling(x, n)
graph_features_grad = tf.gradients(actual_op, [x])[0]
return sess.run(graph_features_grad)
def test_backward_simple(self):
cpu = self._simple_backward(use_gpu=False)
cpu[0, 0, 2] -= 4
self.assertEqual(cpu.sum(), 0)
gpu = self._simple_backward(use_gpu=True)
gpu[0, 0, 2] -= 4
self.assertEqual(gpu.sum(), 0)
if __name__ == '__main__':
tf.test.main()