mirror of
https://github.com/wassname/Flex-Convolution.git
synced 2026-08-30 11:14:38 +08:00
Add more operation and some refactoring of the code to ease usage.
This commit is contained in:
Fabian Groh
+26
-14
@@ -25,39 +25,51 @@ Demonstration of using FlexConvolution, FlexPooling Layer.
|
||||
import numpy as np
|
||||
import tensorflow as tf
|
||||
from tabulate import tabulate
|
||||
from layers import flex_convolution, flex_convolution_transpose, flex_pooling
|
||||
from layers import (flex_convolution,
|
||||
flex_convolution_transpose,
|
||||
flex_pooling,
|
||||
knn_bruteforce)
|
||||
|
||||
B, Din, Dout, Dout2, Dp, N, N2, K, K2 = 1, 2, 4, 8, 3, 10, 5, 5, 3
|
||||
|
||||
B, Din, Dout, Dout2, Dp, N, K = 1, 2, 4, 8, 3, 10, 5
|
||||
|
||||
features = np.random.randn(B, Din, N).astype(np.float32)
|
||||
positions = np.random.randn(B, Dp, N).astype(np.float32)
|
||||
neighbors = np.random.randint(0, N, [B, K, N]).astype(np.int32)
|
||||
neighbors2 = np.random.randint(0, N, [B, K2, N2]).astype(np.int32)
|
||||
|
||||
features = tf.convert_to_tensor(features, name='features')
|
||||
positions = tf.convert_to_tensor(positions, name='positions')
|
||||
neighbors = tf.convert_to_tensor(neighbors, name='neighbors')
|
||||
neighbors2 = tf.convert_to_tensor(neighbors2, name='neighbors2')
|
||||
|
||||
net = [features]
|
||||
# use our FlexConv similar to a traditional convolution layer
|
||||
net.append(flex_convolution(net[-1], positions, neighbors, Dout,
|
||||
|
||||
neighbors = knn_bruteforce(positions, K=5)
|
||||
net.append(flex_convolution(net[-1],
|
||||
positions,
|
||||
neighbors,
|
||||
Dout,
|
||||
activation=tf.nn.relu))
|
||||
# pool and sub-sampling are different operations
|
||||
net.append(flex_pooling(net[-1], neighbors))
|
||||
|
||||
# when ordering the points beforehand sub-sampling is simply
|
||||
features = net[-1][:, :, :N2]
|
||||
positions = positions[:, :, :N2]
|
||||
|
||||
features = net[-1][:, :, :N // 2]
|
||||
positions = positions[:, :, :N // 2]
|
||||
net.append(features)
|
||||
|
||||
neighbors = knn_bruteforce(positions, K=3)
|
||||
# we didn't notice any improvements using the transposed version vs. pooling
|
||||
net.append(flex_convolution_transpose(net[-1], positions, neighbors2, Dout2,
|
||||
net.append(flex_convolution_transpose(net[-1],
|
||||
positions,
|
||||
neighbors,
|
||||
Dout,
|
||||
activation=tf.nn.relu))
|
||||
# of course any commonly used arguments work here as well
|
||||
net.append(flex_convolution(net[-1], positions,
|
||||
neighbors2, Dout2,
|
||||
trainable=False, activation=tf.nn.relu))
|
||||
net.append(flex_convolution(net[-1],
|
||||
positions,
|
||||
neighbors,
|
||||
Dout,
|
||||
activation=tf.nn.relu,
|
||||
trainable=False))
|
||||
|
||||
gradient_wrt_feature = tf.gradients(net[-1], net[0])
|
||||
|
||||
|
||||
@@ -20,25 +20,86 @@
|
||||
import tensorflow as tf
|
||||
from user_ops import flex_convolution as _flex_convolution
|
||||
from user_ops import flex_pooling as _flex_pooling
|
||||
from user_ops import knn_bruteforce as _knn_bruteforce
|
||||
from user_ops import flex_convolution_transpose as _flex_convolution_transpose
|
||||
|
||||
from tensorflow.python.keras import activations
|
||||
from tensorflow.python.keras import initializers
|
||||
from tensorflow.python.layers.base import Layer
|
||||
from tensorflow.python.util.tf_export import tf_export
|
||||
from tensorflow.python.framework import tensor_shape
|
||||
from tensorflow.python.framework import ops
|
||||
|
||||
|
||||
all = ['FlexPooling', 'FlexConvolution', 'FlexConvolutionTranspose',
|
||||
'flex_pooling', 'flex_convolution', 'flex_convolution_transpose']
|
||||
'flex_pooling', 'flex_convolution', 'flex_convolution_transpose',
|
||||
'knn_bruteforce', 'KnnBruteforce']
|
||||
|
||||
|
||||
def _remove_dim(x, axis=2):
|
||||
return tf.squeeze(x, axis=axis)
|
||||
|
||||
|
||||
@tf_export('keras.layers.FlexPooling')
|
||||
class KnnBruteforce(Layer):
|
||||
"""knn bruteforce layer.
|
||||
This layer performs a nearest neighbor lookup for a batch of given positions.
|
||||
|
||||
Arguments:
|
||||
K: size of each neighborhood
|
||||
data_format: A string, one of `simple` (default) or `expanded`.
|
||||
If `simple` the shapes are [B, Din, N], when `expanded` the shapes
|
||||
are assumed to be [B, Din, 1, N] to match `channels_first` in trad
|
||||
convolutions.
|
||||
name: A string, the name of the layer.
|
||||
|
||||
Inputs:
|
||||
positions: A `Tensor` of the format [B, Dp, (1), N].
|
||||
|
||||
Outputs:
|
||||
neighborhoods: A `Tensor` of the format [B, K, (1), N] (tf.int32).
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
K,
|
||||
data_format='simple',
|
||||
name=None):
|
||||
|
||||
super(KnnBruteforce, self).__init__(name=name)
|
||||
assert K > 0
|
||||
assert data_format in ['simple', 'expanded']
|
||||
self.K = K
|
||||
self.data_format = data_format
|
||||
|
||||
def compute_output_shape(self, input_shapes):
|
||||
output_shape = input_shapes[0]
|
||||
output_shape[1] = self.K
|
||||
return output_shape
|
||||
|
||||
def call(self, inputs):
|
||||
|
||||
positions = ops.convert_to_tensor(inputs, dtype=self.dtype)
|
||||
if self.data_format == 'expanded':
|
||||
positions = _remove_dim(positions, 2)
|
||||
|
||||
NN, _, _ = _knn_bruteforce(positions, K=4)
|
||||
NN = tf.transpose(NN, [0, 2, 1])
|
||||
|
||||
if self.data_format == 'expanded':
|
||||
NN = tf.expand_dims(NN, axis=2)
|
||||
|
||||
return NN
|
||||
|
||||
|
||||
def knn_bruteforce(positions,
|
||||
K,
|
||||
data_format='simple',
|
||||
name=None):
|
||||
|
||||
layer = KnnBruteforce(K, data_format=data_format, name=name)
|
||||
|
||||
return layer.apply(positions)
|
||||
|
||||
|
||||
class FlexPooling(Layer):
|
||||
"""flex pooling layer.
|
||||
|
||||
@@ -53,8 +114,19 @@ class FlexPooling(Layer):
|
||||
Arguments:
|
||||
features: A `Tensor` of the format [B, Din, (1), N].
|
||||
neighborhoods: A `Tensor` of the format [B, K, (1), N] (tf.int32).
|
||||
data_format: A string, one of `simple` (default) or `expanded`.
|
||||
If `simple` the shapes are [B, Din, N], when `expanded` the shapes
|
||||
are assumed to be [B, Din, 1, N] to match `channels_first` in trad
|
||||
convolutions.
|
||||
name: A string, the name of the layer.
|
||||
|
||||
Inputs:
|
||||
features: A `Tensor` of the format[B, Din, (1), N].
|
||||
neighborhoods: A `Tensor` of the format [B, K, (1), N] (tf.int32).
|
||||
|
||||
Outputs:
|
||||
features: A `Tensor` of the format[B, Din, (1), N].
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
@@ -62,6 +134,7 @@ class FlexPooling(Layer):
|
||||
name=None):
|
||||
|
||||
super(FlexPooling, self).__init__(name=name)
|
||||
assert data_format in ['simple', 'expanded']
|
||||
self.data_format = data_format
|
||||
|
||||
def compute_output_shape(self, input_shape):
|
||||
@@ -93,11 +166,9 @@ def flex_pooling(features,
|
||||
name=None):
|
||||
|
||||
layer = FlexPooling(data_format=data_format, name=name)
|
||||
|
||||
return layer.apply([features, neighborhoods])
|
||||
|
||||
|
||||
@tf_export('keras.layers.FlexConvolution')
|
||||
class FlexConvolution(Layer):
|
||||
"""flex convolution layer.
|
||||
|
||||
@@ -126,7 +197,7 @@ class FlexConvolution(Layer):
|
||||
features_bias_initializer: An initializer for the bias vector after
|
||||
the convolution. If None, the default initializer will be used.
|
||||
use_feature_bias: Boolean, whether the layer uses a bias.
|
||||
data_format: A string, one of `simple` (default) or `expaned`.
|
||||
data_format: A string, one of `simple` (default) or `expanded`.
|
||||
If `simple` the shapes are [B, Din, N], when `expanded` the shapes
|
||||
are assumed to be [B, Din, 1, N] to match `channels_first` in trad
|
||||
convolutions.
|
||||
@@ -136,9 +207,11 @@ class FlexConvolution(Layer):
|
||||
|
||||
Inputs:
|
||||
features: A `Tensor` of the format[B, Din, (1), N].
|
||||
positions: A `Tensor` of the format[B, Dp, (1), N].
|
||||
neigh
|
||||
borhoods: A `Tensor` of the format [B, K, (1), N] (tf.int32).
|
||||
positions: A `Tensor` of the format[B, Dp, (1), N].
|
||||
neighborhoods: A `Tensor` of the format [B, K, (1), N] (tf.int32).
|
||||
|
||||
Outputs:
|
||||
features: A `Tensor` of the format[B, Dout, (1), N].
|
||||
|
||||
"""
|
||||
|
||||
@@ -155,6 +228,7 @@ borhoods: A `Tensor` of the format [B, K, (1), N] (tf.int32).
|
||||
|
||||
super(FlexConvolution, self).__init__(trainable=trainable,
|
||||
name=name)
|
||||
assert data_format in ['simple', 'expanded']
|
||||
self.filters = int(filters)
|
||||
self.activation = activations.get(activation)
|
||||
self.use_feature_bias = use_feature_bias
|
||||
@@ -226,7 +300,7 @@ borhoods: A `Tensor` of the format [B, K, (1), N] (tf.int32).
|
||||
|
||||
Returns:
|
||||
tf.tensor: A `Tensor` of the format [B, Dout, (1), N] describing
|
||||
the outgoing features.
|
||||
the outgoing features.
|
||||
|
||||
"""
|
||||
if not isinstance(inputs, list):
|
||||
@@ -283,7 +357,6 @@ def flex_convolution(features,
|
||||
return layer.apply([features, positions, neighborhoods])
|
||||
|
||||
|
||||
@tf_export('keras.layers.FlexConvolutionTranspose')
|
||||
class FlexConvolutionTranspose(FlexConvolution):
|
||||
"""flex convolution-transpose layer.
|
||||
|
||||
@@ -312,7 +385,7 @@ class FlexConvolutionTranspose(FlexConvolution):
|
||||
features_bias_initializer: An initializer for the bias vector after
|
||||
the convolution. If None, the default initializer will be used.
|
||||
use_feature_bias: Boolean, whether the layer uses a bias.
|
||||
data_format: A string, one of `simple` (default) or `expaned`.
|
||||
data_format: A string, one of `simple` (default) or `expanded`.
|
||||
If `simple` the shapes are [B, Din, N], when `expanded` the shapes
|
||||
are assumed to be [B, Din, 1, N] to match `channels_first` in trad
|
||||
convolutions.
|
||||
@@ -325,6 +398,9 @@ class FlexConvolutionTranspose(FlexConvolution):
|
||||
positions: A `Tensor` of the format [B, Dp, (1), N].
|
||||
neighborhoods: A `Tensor` of the format [B, K, (1), N] (tf.int32).
|
||||
|
||||
Outputs:
|
||||
features: A `Tensor` of the format[B, Dout, (1), N].
|
||||
|
||||
"""
|
||||
|
||||
def internal_call(self,
|
||||
|
||||
@@ -13,7 +13,12 @@ list(APPEND CMAKE_MODULE_PATH ${PROJECT_SOURCE_DIR})
|
||||
find_package(CUDA REQUIRED)
|
||||
find_package(TensorFlow REQUIRED)
|
||||
|
||||
# include CUDA Toolkit header files
|
||||
set(CUDA_SAMPLE_INC "$ENV{CUDA_INSTALL_PATH}/samples/common/inc")
|
||||
message(STATUS "CUDA_INCLUDE_DIRS: ${CUDA_INCLUDE_DIRS}")
|
||||
message(STATUS "CUDA_SAMPLE_INC: ${CUDA_SAMPLE_INC}")
|
||||
|
||||
# include CUBS
|
||||
if (DEFINED ENV{CUB_INC})
|
||||
message(STATUS "Use Cuda-CUB from " $ENV{CUB_INC})
|
||||
set(CUB_INC $ENV{CUB_INC})
|
||||
@@ -31,6 +36,8 @@ include_directories(SYSTEM "/usr/local/")
|
||||
# fix cgtuebingen
|
||||
include_directories(SYSTEM "/graphics/opt/opt_Ubuntu18.04/cuda/toolkit_9.2")
|
||||
|
||||
include_directories(SYSTEM "${CUDA_SAMPLE_INC}")
|
||||
include_directories(SYSTEM "${CUDA_INCLUDE_DIRS}/../../")
|
||||
include_directories(SYSTEM ${CUB_INC})
|
||||
include_directories(SYSTEM ${TensorFlow_INCLUDE_DIR})
|
||||
include_directories(SYSTEM kernels)
|
||||
@@ -38,3 +45,4 @@ include_directories(SYSTEM kernels)
|
||||
add_tensorflow_gpu_operation("flex_conv")
|
||||
add_tensorflow_gpu_operation("flex_deconv")
|
||||
add_tensorflow_gpu_operation("flex_pool")
|
||||
add_tensorflow_gpu_operation("knn_bruteforce")
|
||||
|
||||
@@ -54,6 +54,7 @@ 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} ")
|
||||
@@ -122,7 +123,7 @@ else()
|
||||
# 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", "1.11", "1.11.0")
|
||||
"1.9" "1.9.0" "1.10" "1.10.0" "1.11" "1.11.0")
|
||||
set(_TensorFlow_TEST_VERSIONS)
|
||||
|
||||
if(TF_FIND_VERSION)
|
||||
@@ -166,6 +167,17 @@ else()
|
||||
"We tested against ${_TensorFlow_TEST_VERSIONS}")
|
||||
endif(NOT TensorFlow_FOUND)
|
||||
|
||||
# test 1.11 version
|
||||
if("${TF_DETECTED_VERSION}" VERSION_EQUAL "1.11")
|
||||
set(TF_DISABLE_ASSERTS "TRUE")
|
||||
endif()
|
||||
|
||||
endif()
|
||||
|
||||
|
||||
if(${TF_DISABLE_ASSERTS})
|
||||
message(STATUS "[WARNING] The TensorFlow version ${TF_DETECTED_VERSION} has a bug (see \#22766). We disable asserts using -DNDEBUG=True ")
|
||||
add_definitions("-DNDEBUG=True")
|
||||
endif()
|
||||
|
||||
find_library(TensorFlow_C_LIBRARY
|
||||
@@ -246,8 +258,8 @@ 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_INCLUDE_DIRECTORIES(TensorFlow_DEP SYSTEM INTERFACE ${TensorFlow_SOURCE_DIR})
|
||||
TARGET_INCLUDE_DIRECTORIES(TensorFlow_DEP SYSTEM 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)
|
||||
|
||||
@@ -263,7 +275,7 @@ find_package_handle_standard_args(
|
||||
)
|
||||
|
||||
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
|
||||
TF_DETECTED_INCLUDE_DIR TF_DETECTED_LIBRARY TF_DISABLE_ASSERTS
|
||||
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")
|
||||
@@ -271,3 +283,4 @@ SET(TensorFlow_VERSION ${TensorFlow_VERSION} CACHE INTERNAL "The Python executab
|
||||
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")
|
||||
SET(TF_DISABLE_ASSERTS ${TF_DISABLE_ASSERTS} CACHE BOOL "Workarounds")
|
||||
@@ -1,130 +0,0 @@
|
||||
#!/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)
|
||||
+42
-22
@@ -19,27 +19,38 @@ from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import tensorflow as tf
|
||||
import os
|
||||
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"))
|
||||
__all__ = []
|
||||
|
||||
|
||||
# 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
|
||||
def load_op(name, has_grad=False, public=False):
|
||||
global __all__
|
||||
path = os.path.join(os.path.dirname(__file__), '%s_op.so' % name)
|
||||
if os.path.isfile(path):
|
||||
_module = loader.load_op_library(path)
|
||||
if has_grad:
|
||||
if public:
|
||||
__all__.append('%s' % name)
|
||||
__all__.append('%s_grad' % name)
|
||||
return getattr(_module, '%s' % name), getattr(_module, '%s_grad' % name)
|
||||
else:
|
||||
if public:
|
||||
__all__.append('%s' % name)
|
||||
return getattr(_module, '%s' % name)
|
||||
else:
|
||||
print('[WARNING]: %s does not exists' % name)
|
||||
|
||||
knn_bruteforce = load_op('knn_bruteforce', has_grad=False, public=True)
|
||||
|
||||
_flex_conv, _flex_conv_grad = load_op(
|
||||
'flex_conv', has_grad=True, public=False)
|
||||
_flex_pool, _flex_pool_grad = load_op(
|
||||
'flex_pool', has_grad=True, public=False)
|
||||
_flex_deconv, _flex_deconv_grad = load_op(
|
||||
'flex_deconv', has_grad=True, public=False)
|
||||
|
||||
|
||||
# pylint: disable=redefined-builtin
|
||||
@@ -69,7 +80,10 @@ def flex_convolution(features,
|
||||
"""
|
||||
|
||||
with ops.name_scope(name, "flex_convolution"):
|
||||
return flex_conv(features, theta, bias, neighborhood, position)
|
||||
return _flex_conv(features, theta, bias, neighborhood, position)
|
||||
|
||||
|
||||
__all__.append('flex_convolution')
|
||||
|
||||
|
||||
@ops.RegisterGradient("FlexConv")
|
||||
@@ -81,7 +95,7 @@ def _FlexConvGrad(op, *grads): # noqa
|
||||
positions = ops.convert_to_tensor(op.inputs[4])
|
||||
topdiff = ops.convert_to_tensor(grads[0])
|
||||
|
||||
df, dt, db = flex_conv_grad(
|
||||
df, dt, db = _flex_conv_grad(
|
||||
features, theta, bias, neighborhood, positions, topdiff)
|
||||
|
||||
df = ops.convert_to_tensor(df, name='gradient_features')
|
||||
@@ -112,7 +126,10 @@ def flex_pooling(features,
|
||||
"""
|
||||
|
||||
with ops.name_scope(name, "flex_pooling"):
|
||||
return flex_pool(features, neighborhood)
|
||||
return _flex_pool(features, neighborhood)
|
||||
|
||||
|
||||
__all__.append('flex_pooling')
|
||||
|
||||
|
||||
@ops.RegisterGradient("FlexPool")
|
||||
@@ -122,7 +139,7 @@ def _FlexPoolGrad(op, *grads): # noqa
|
||||
argmax = ops.convert_to_tensor(op.outputs[1])
|
||||
topdiff = ops.convert_to_tensor(grads[0])
|
||||
|
||||
df = flex_pool_grad(features, neighborhood, topdiff, argmax)
|
||||
df = _flex_pool_grad(features, neighborhood, topdiff, argmax)
|
||||
df = ops.convert_to_tensor(df, name='gradient_features')
|
||||
|
||||
return [df, None]
|
||||
@@ -153,7 +170,10 @@ def flex_convolution_transpose(features,
|
||||
"""
|
||||
|
||||
with ops.name_scope(name, "flex_convolution_transpose"):
|
||||
return flex_deconv(features, theta, bias, neighborhood, position)
|
||||
return _flex_deconv(features, theta, bias, neighborhood, position)
|
||||
|
||||
|
||||
__all__.append('_flex_deconv')
|
||||
|
||||
|
||||
@ops.RegisterGradient("FlexDeconv")
|
||||
@@ -165,7 +185,7 @@ def _FlexDeconvGrad(op, *grads): # noqa
|
||||
positions = ops.convert_to_tensor(op.inputs[4])
|
||||
topdiff = ops.convert_to_tensor(grads[0])
|
||||
|
||||
df, dt, db = flex_deconv_grad(
|
||||
df, dt, db = _flex_deconv_grad(
|
||||
features, theta, bias, neighborhood, positions, topdiff)
|
||||
|
||||
df = ops.convert_to_tensor(df, name='gradient_features')
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
#ifndef LIB_CUDA_UTILS_H_
|
||||
#define LIB_CUDA_UTILS_H_
|
||||
|
||||
template <typename T>
|
||||
__device__ T* DynamicSharedMemory() {
|
||||
extern __shared__ __align__(sizeof(T)) unsigned char s_shm[];
|
||||
return reinterpret_cast<T*>(s_shm);
|
||||
}
|
||||
|
||||
#endif // LIB_CUDA_UTILS_H_
|
||||
@@ -70,6 +70,7 @@ struct FlexConvFunctor<CPUDevice, Dtype> {
|
||||
};
|
||||
|
||||
template struct FlexConvFunctor<CPUDevice, float>;
|
||||
template struct FlexConvFunctor<CPUDevice, double>;
|
||||
|
||||
template <typename Dtype>
|
||||
struct FlexConvGrad<CPUDevice, Dtype> {
|
||||
@@ -167,7 +168,7 @@ struct FlexConvGrad<CPUDevice, Dtype> {
|
||||
|
||||
// template struct FlexConvGrad<CPUDevice, int>;
|
||||
template struct FlexConvGrad<CPUDevice, float>;
|
||||
// template struct FlexConvGrad<CPUDevice, double>;
|
||||
template struct FlexConvGrad<CPUDevice, double>;
|
||||
|
||||
} // namespace functor
|
||||
} // namespace tensorflow
|
||||
|
||||
@@ -12,7 +12,7 @@ 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
|
||||
// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
|
||||
|
||||
#if GOOGLE_CUDA
|
||||
|
||||
@@ -20,6 +20,7 @@ limitations under the License.
|
||||
|
||||
#include <cub/cub.cuh>
|
||||
|
||||
#include "cuda_utils.h"
|
||||
#include "flex_conv_op.h"
|
||||
#include "tensorflow/core/util/cuda_kernel_helper.h"
|
||||
|
||||
@@ -57,10 +58,10 @@ struct ForwardKernel {
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void operator()() const {
|
||||
extern __shared__ Dtype s_shm[];
|
||||
Dtype* s_shm = DynamicSharedMemory<Dtype>();
|
||||
|
||||
Dtype* s_theta = (float*)&s_shm[0];
|
||||
Dtype* s_bias = (float*)&s_shm[Dp * C_Din * C_Dout];
|
||||
Dtype* s_theta = (Dtype*)&s_shm[0];
|
||||
Dtype* s_bias = (Dtype*)&s_shm[Dp * C_Din * C_Dout];
|
||||
|
||||
// glob ids
|
||||
int b = blockIdx.z;
|
||||
@@ -281,7 +282,8 @@ struct BackwardFeatureKernel {
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void operator()() const {
|
||||
extern __shared__ float s_shm[];
|
||||
// extern __shared__ float s_shm[];
|
||||
Dtype* s_shm = DynamicSharedMemory<Dtype>();
|
||||
|
||||
int i_n = threadIdx.x;
|
||||
int i_din = threadIdx.y;
|
||||
@@ -354,7 +356,10 @@ struct BackwardFeatureKernel {
|
||||
W += s_bias[i_din * C_Dout + dout_inner];
|
||||
Dtype value = W * s_topdiff[dout_inner * C_N + i_n];
|
||||
|
||||
atomicAdd(
|
||||
// atomicAdd(
|
||||
// &d_features_out[b * Din * N + din * N + s_nk[k * C_N + i_n]],
|
||||
// value);
|
||||
tensorflow::CudaAtomicAdd(
|
||||
&d_features_out[b * Din * N + din * N + s_nk[k * C_N + i_n]],
|
||||
value);
|
||||
}
|
||||
@@ -384,6 +389,16 @@ struct BackwardFeatureKernel {
|
||||
namespace tensorflow {
|
||||
namespace functor {
|
||||
|
||||
template <class Dtype, class NBtype>
|
||||
struct ForwardKernelType {
|
||||
typedef FlexConvCuda::ForwardKernel<Dtype, NBtype, 3, 128, 32, 32> type;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct ForwardKernelType<float, int> {
|
||||
typedef FlexConvCuda::ForwardKernel<float, int, 3, 128, 32, 64> type;
|
||||
};
|
||||
|
||||
template <typename Dtype>
|
||||
struct FlexConvFunctor<GPUDevice, Dtype> {
|
||||
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& features,
|
||||
@@ -399,7 +414,11 @@ struct FlexConvFunctor<GPUDevice, Dtype> {
|
||||
const int Din = theta.dim_size(2);
|
||||
const int Dout = theta.dim_size(3);
|
||||
|
||||
FlexConvCuda::ForwardKernel<Dtype, NBtype, 3, 128, 32, 64> fwk;
|
||||
// printf("<f> test: %s\n", __PRETTY_FUNCTION__);
|
||||
|
||||
typedef typename ForwardKernelType<Dtype, NBtype>::type FKT;
|
||||
|
||||
FKT fwk;
|
||||
fwk.N = N;
|
||||
fwk.K = K;
|
||||
fwk.Din = Din;
|
||||
@@ -422,6 +441,7 @@ struct FlexConvFunctor<GPUDevice, Dtype> {
|
||||
};
|
||||
|
||||
template struct FlexConvFunctor<GPUDevice, float>;
|
||||
template struct FlexConvFunctor<GPUDevice, double>;
|
||||
|
||||
template <typename Dtype>
|
||||
struct FlexConvGrad<GPUDevice, Dtype> {
|
||||
@@ -459,7 +479,7 @@ struct FlexConvGrad<GPUDevice, Dtype> {
|
||||
|
||||
const Dtype* topdiff_ptr = reinterpret_cast<const Dtype*>(topdiff.data());
|
||||
|
||||
Dtype* grad_features_ptr = reinterpret_cast<float*>(grad_features.data());
|
||||
Dtype* grad_features_ptr = reinterpret_cast<Dtype*>(grad_features.data());
|
||||
Dtype* grad_theta_ptr = reinterpret_cast<Dtype*>(grad_theta.data());
|
||||
Dtype* grad_bias_ptr = reinterpret_cast<Dtype*>(grad_bias.data());
|
||||
|
||||
@@ -525,6 +545,7 @@ struct FlexConvGrad<GPUDevice, Dtype> {
|
||||
};
|
||||
|
||||
template struct FlexConvGrad<GPUDevice, float>;
|
||||
template struct FlexConvGrad<GPUDevice, double>;
|
||||
|
||||
} // namespace functor
|
||||
} // namespace tensorflow
|
||||
|
||||
@@ -12,7 +12,7 @@ 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
|
||||
// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
|
||||
|
||||
#include "flex_conv_op.h"
|
||||
|
||||
@@ -91,32 +91,22 @@ class FlexConvGradOp : public OpKernel {
|
||||
}
|
||||
};
|
||||
|
||||
// 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>)
|
||||
#define REGISTER_CUSTOM_OP(NAME, DEVICE, T) \
|
||||
REGISTER_KERNEL_BUILDER( \
|
||||
Name(#NAME).Device(DEVICE_##DEVICE).TypeConstraint<T>("T"), \
|
||||
NAME##Op<DEVICE##Device, T>)
|
||||
|
||||
TF_CALL_float(REGISTER_FLEXCONV_OP_CPU);
|
||||
#undef REGISTER_FLEXCONV_OP_CPU
|
||||
REGISTER_CUSTOM_OP(FlexConv, CPU, float);
|
||||
REGISTER_CUSTOM_OP(FlexConvGrad, CPU, float);
|
||||
REGISTER_CUSTOM_OP(FlexConv, CPU, double);
|
||||
REGISTER_CUSTOM_OP(FlexConvGrad, CPU, double);
|
||||
|
||||
// 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
|
||||
#ifdef GOOGLE_CUDA
|
||||
REGISTER_CUSTOM_OP(FlexConv, GPU, float);
|
||||
REGISTER_CUSTOM_OP(FlexConvGrad, GPU, float);
|
||||
REGISTER_CUSTOM_OP(FlexConv, GPU, double);
|
||||
REGISTER_CUSTOM_OP(FlexConvGrad, GPU, double);
|
||||
#endif // GOOGLE_CUDA
|
||||
#undef REGISTER_CUSTOM_OP
|
||||
|
||||
} // namespace tensorflow
|
||||
|
||||
@@ -71,6 +71,7 @@ struct FlexDeconvFunctor<CPUDevice, Dtype> {
|
||||
};
|
||||
|
||||
template struct FlexDeconvFunctor<CPUDevice, float>;
|
||||
template struct FlexDeconvFunctor<CPUDevice, double>;
|
||||
|
||||
template <typename Dtype>
|
||||
struct FlexDeconvGrad<CPUDevice, Dtype> {
|
||||
@@ -165,9 +166,8 @@ struct FlexDeconvGrad<CPUDevice, Dtype> {
|
||||
}
|
||||
};
|
||||
|
||||
// template struct FlexDeconvGrad<CPUDevice, int>;
|
||||
template struct FlexDeconvGrad<CPUDevice, float>;
|
||||
// template struct FlexDeconvGrad<CPUDevice, double>;
|
||||
template struct FlexDeconvGrad<CPUDevice, double>;
|
||||
|
||||
} // namespace functor
|
||||
} // namespace tensorflow
|
||||
|
||||
@@ -177,6 +177,7 @@ struct FlexDeconvFunctor<GPUDevice, Dtype> {
|
||||
};
|
||||
|
||||
template struct FlexDeconvFunctor<GPUDevice, float>;
|
||||
template struct FlexDeconvFunctor<GPUDevice, double>;
|
||||
|
||||
template <typename Dtype>
|
||||
struct FlexDeconvGrad<GPUDevice, Dtype> {
|
||||
@@ -220,6 +221,7 @@ struct FlexDeconvGrad<GPUDevice, Dtype> {
|
||||
};
|
||||
|
||||
template struct FlexDeconvGrad<GPUDevice, float>;
|
||||
template struct FlexDeconvGrad<GPUDevice, double>;
|
||||
|
||||
} // namespace functor
|
||||
} // namespace tensorflow
|
||||
|
||||
@@ -12,7 +12,7 @@ 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
|
||||
// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
|
||||
|
||||
#include "flex_deconv_op.h"
|
||||
|
||||
@@ -88,16 +88,22 @@ class FlexDeconvGradOp : public OpKernel {
|
||||
}
|
||||
};
|
||||
|
||||
#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 >);
|
||||
#define REGISTER_CUSTOM_OP(NAME, DEVICE, T) \
|
||||
REGISTER_KERNEL_BUILDER( \
|
||||
Name(#NAME).Device(DEVICE_##DEVICE).TypeConstraint<T>("T"), \
|
||||
NAME##Op<DEVICE##Device, T>)
|
||||
|
||||
REGISTER(FlexDeconv, float);
|
||||
REGISTER(FlexDeconvGrad, float);
|
||||
REGISTER_CUSTOM_OP(FlexDeconv, CPU, float);
|
||||
REGISTER_CUSTOM_OP(FlexDeconvGrad, CPU, float);
|
||||
REGISTER_CUSTOM_OP(FlexDeconv, CPU, double);
|
||||
REGISTER_CUSTOM_OP(FlexDeconvGrad, CPU, double);
|
||||
|
||||
#ifdef GOOGLE_CUDA
|
||||
REGISTER_CUSTOM_OP(FlexDeconv, GPU, float);
|
||||
REGISTER_CUSTOM_OP(FlexDeconvGrad, GPU, float);
|
||||
REGISTER_CUSTOM_OP(FlexDeconv, GPU, double);
|
||||
REGISTER_CUSTOM_OP(FlexDeconvGrad, GPU, double);
|
||||
#endif // GOOGLE_CUDA
|
||||
#undef REGISTER_CUSTOM_OP
|
||||
|
||||
} // namespace tensorflow
|
||||
|
||||
@@ -12,7 +12,7 @@ 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
|
||||
// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
|
||||
|
||||
#include "flex_pool_op.h"
|
||||
#include "tensorflow/core/framework/op.h"
|
||||
@@ -59,6 +59,7 @@ struct FlexPoolFunctor<CPUDevice, Dtype> {
|
||||
};
|
||||
|
||||
template struct FlexPoolFunctor<CPUDevice, float>;
|
||||
template struct FlexPoolFunctor<CPUDevice, double>;
|
||||
|
||||
template <typename Dtype>
|
||||
struct FlexPoolGrad<CPUDevice, Dtype> {
|
||||
@@ -94,9 +95,8 @@ struct FlexPoolGrad<CPUDevice, Dtype> {
|
||||
}
|
||||
};
|
||||
|
||||
// template struct FlexPoolGrad<CPUDevice, int>;
|
||||
template struct FlexPoolGrad<CPUDevice, float>;
|
||||
// template struct FlexPoolGrad<CPUDevice, double>;
|
||||
template struct FlexPoolGrad<CPUDevice, double>;
|
||||
|
||||
} // namespace functor
|
||||
} // namespace tensorflow
|
||||
|
||||
@@ -12,7 +12,7 @@ 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
|
||||
// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
|
||||
|
||||
#if GOOGLE_CUDA
|
||||
|
||||
@@ -30,7 +30,7 @@ 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) {
|
||||
Dtype* output, int* argmax, Dtype 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].
|
||||
@@ -41,14 +41,14 @@ __global__ void forward(const int B, const int N, const int K, const int 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;
|
||||
Dtype 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];
|
||||
const Dtype v = features[b * D * N + d * N + other_global_id];
|
||||
|
||||
if (best_value < v) {
|
||||
best_id = other_global_id;
|
||||
@@ -113,13 +113,9 @@ struct FlexPoolFunctor<GPUDevice, Dtype> {
|
||||
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());
|
||||
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(
|
||||
@@ -129,6 +125,7 @@ struct FlexPoolFunctor<GPUDevice, Dtype> {
|
||||
};
|
||||
|
||||
template struct FlexPoolFunctor<GPUDevice, float>;
|
||||
template struct FlexPoolFunctor<GPUDevice, double>;
|
||||
|
||||
template <typename Dtype>
|
||||
struct FlexPoolGrad<GPUDevice, Dtype> {
|
||||
@@ -149,13 +146,9 @@ struct FlexPoolGrad<GPUDevice, Dtype> {
|
||||
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());
|
||||
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!"));
|
||||
@@ -164,6 +157,7 @@ struct FlexPoolGrad<GPUDevice, Dtype> {
|
||||
};
|
||||
|
||||
template struct FlexPoolGrad<GPUDevice, float>;
|
||||
template struct FlexPoolGrad<GPUDevice, double>;
|
||||
|
||||
} // namespace functor
|
||||
} // namespace tensorflow
|
||||
|
||||
@@ -12,7 +12,7 @@ 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
|
||||
// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
|
||||
|
||||
#include "flex_pool_op.h"
|
||||
|
||||
@@ -33,7 +33,6 @@ class FlexPoolOp : public OpKernel {
|
||||
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);
|
||||
|
||||
@@ -65,7 +64,6 @@ class FlexPoolGradOp : public OpKernel {
|
||||
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);
|
||||
@@ -82,32 +80,22 @@ class FlexPoolGradOp : public OpKernel {
|
||||
}
|
||||
};
|
||||
|
||||
// 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>)
|
||||
#define REGISTER_CUSTOM_OP(NAME, DEVICE, T) \
|
||||
REGISTER_KERNEL_BUILDER( \
|
||||
Name(#NAME).Device(DEVICE_##DEVICE).TypeConstraint<T>("T"), \
|
||||
NAME##Op<DEVICE##Device, T>)
|
||||
|
||||
TF_CALL_float(REGISTER_FLEXPOOL_OP_CPU);
|
||||
#undef REGISTER_FLEXPOOL_OP_CPU
|
||||
REGISTER_CUSTOM_OP(FlexPool, CPU, float);
|
||||
REGISTER_CUSTOM_OP(FlexPoolGrad, CPU, float);
|
||||
REGISTER_CUSTOM_OP(FlexPool, CPU, double);
|
||||
REGISTER_CUSTOM_OP(FlexPoolGrad, CPU, double);
|
||||
|
||||
// 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
|
||||
|
||||
REGISTER_CUSTOM_OP(FlexPool, GPU, float);
|
||||
REGISTER_CUSTOM_OP(FlexPoolGrad, GPU, float);
|
||||
REGISTER_CUSTOM_OP(FlexPool, GPU, double);
|
||||
REGISTER_CUSTOM_OP(FlexPoolGrad, GPU, double);
|
||||
#endif // GOOGLE_CUDA
|
||||
#undef REGISTER_CUSTOM_OP
|
||||
|
||||
} // namespace tensorflow
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
/* 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.
|
||||
==============================================================================*/
|
||||
// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
|
||||
|
||||
#include "knn_bruteforce_op.h"
|
||||
#include "tensorflow/core/framework/op.h"
|
||||
|
||||
namespace tensorflow {
|
||||
|
||||
namespace functor {
|
||||
|
||||
template <typename Dtype, typename NBtype>
|
||||
struct KnnBruteforceFunctor<CPUDevice, Dtype, NBtype> {
|
||||
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& positions_,
|
||||
Tensor* neighborhood_out_, Tensor* distances_,
|
||||
Tensor* timings) {
|
||||
// positions [B, Dp, N]
|
||||
// neighborhood_out_ [B, N, K]
|
||||
// distances_ [B, N, K]
|
||||
const auto positions = positions_.tensor<Dtype, 3>();
|
||||
auto neighborhood_out = neighborhood_out_->tensor<NBtype, 3>();
|
||||
auto distances_out = distances_->tensor<Dtype, 3>();
|
||||
|
||||
const int B = positions_.dim_size(0);
|
||||
const int Dp = positions_.dim_size(1);
|
||||
const int N = positions_.dim_size(2);
|
||||
|
||||
const int K = neighborhood_out_->dim_size(2);
|
||||
|
||||
for (int b = 0; b < B; ++b) {
|
||||
const Dtype* pc_raw = positions_.flat<Dtype>().data() + b * Dp * N;
|
||||
Eigen::Map<const Eigen::Matrix<Dtype, Eigen::Dynamic, Eigen::Dynamic,
|
||||
Eigen::RowMajor>>
|
||||
pc(pc_raw, Dp, N);
|
||||
// pc: [Dp, N]
|
||||
|
||||
for (int n = 0; n < N; ++n) {
|
||||
const auto query = pc.col(n);
|
||||
const auto diffs =
|
||||
(pc.colwise() - query).array() * (pc.colwise() - query).array();
|
||||
const auto distances = (diffs.colwise().sum()).array().sqrt();
|
||||
|
||||
std::vector<Dtype> vec_dist;
|
||||
std::vector<Dtype> vec_ids;
|
||||
|
||||
for (int i = 0; i < N; ++i) {
|
||||
vec_dist.push_back(distances(i));
|
||||
vec_ids.push_back(i);
|
||||
}
|
||||
|
||||
std::sort(std::begin(vec_ids), std::end(vec_ids),
|
||||
[&](int i1, int i2) { return vec_dist[i1] < vec_dist[i2]; });
|
||||
|
||||
for (int k = 0; k < K; ++k) {
|
||||
neighborhood_out(b, n, k) = vec_ids[k];
|
||||
distances_out(b, n, k) = vec_dist[vec_ids[k]];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template struct KnnBruteforceFunctor<CPUDevice, float, int>;
|
||||
template struct KnnBruteforceFunctor<CPUDevice, double, int>;
|
||||
|
||||
} // namespace functor
|
||||
} // namespace tensorflow
|
||||
@@ -0,0 +1,273 @@
|
||||
/* 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.
|
||||
==============================================================================*/
|
||||
// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
|
||||
|
||||
#if GOOGLE_CUDA
|
||||
|
||||
#define EIGEN_USE_GPU
|
||||
|
||||
#include <vector>
|
||||
#include "cuda_utils.h"
|
||||
|
||||
#include <cuda.h>
|
||||
#include <helper_cuda.h>
|
||||
|
||||
#include <cub/cub.cuh>
|
||||
|
||||
#include <curand.h>
|
||||
#include <curand_kernel.h>
|
||||
|
||||
#include <limits>
|
||||
|
||||
#include "knn_bruteforce_op.h"
|
||||
#include "tensorflow/core/util/cuda_kernel_helper.h"
|
||||
|
||||
template <typename Dtype, typename NBtype, int C_THREADS = 256, int C_VPT = 2>
|
||||
struct BlockBFKernel;
|
||||
|
||||
template <typename Dtype, typename NBtype, int C_THREADS, int C_VPT>
|
||||
__global__ void runBlockBFKernel(
|
||||
const BlockBFKernel<Dtype, NBtype, C_THREADS, C_VPT> kernel) {
|
||||
kernel();
|
||||
}
|
||||
|
||||
template <typename Dtype, typename NBtype, int C_THREADS, int C_VPT>
|
||||
struct BlockBFKernel {
|
||||
void launch(int B) {
|
||||
dim3 block(C_THREADS, 1);
|
||||
dim3 grid(N, 1, B);
|
||||
|
||||
size_t shm_size = Dp * sizeof(Dtype);
|
||||
|
||||
runBlockBFKernel<<<grid, block, shm_size>>>((*this));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void operator()() const {
|
||||
// extern __shared__ Dtype s_shm[];
|
||||
Dtype* s_shm = DynamicSharedMemory<Dtype>();
|
||||
Dtype* s_point = (Dtype*)&s_shm[0];
|
||||
|
||||
if (N > C_THREADS * C_VPT) {
|
||||
printf(
|
||||
"<BlockBFKernel> Critical problem!!!!! Not enough resources spend "
|
||||
"for N points!!! %d < %d \n",
|
||||
C_THREADS * C_VPT, N);
|
||||
return;
|
||||
}
|
||||
|
||||
int b = blockIdx.z;
|
||||
int y = blockIdx.x;
|
||||
|
||||
int tid = threadIdx.x;
|
||||
for (int dpi = tid; dpi < Dp; dpi += blockDim.x) {
|
||||
s_point[dpi] = d_data[b * N * Dp + dpi * N + y]; // not aligned with
|
||||
// data!
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
Dtype thread_dists[C_VPT]; // keys
|
||||
NBtype thread_ids[C_VPT]; // values
|
||||
|
||||
typedef cub::BlockRadixSort<Dtype, C_THREADS, C_VPT, NBtype> BlockRadixSort;
|
||||
typedef cub::BlockStore<Dtype, C_THREADS, C_VPT,
|
||||
cub::BLOCK_STORE_WARP_TRANSPOSE>
|
||||
BlockStoreDists;
|
||||
typedef cub::BlockStore<NBtype, C_THREADS, C_VPT,
|
||||
cub::BLOCK_STORE_WARP_TRANSPOSE>
|
||||
BlockStoreIds;
|
||||
|
||||
// Allocate shared memory
|
||||
__shared__ union {
|
||||
typename BlockRadixSort::TempStorage sort;
|
||||
typename BlockStoreDists::TempStorage store_dists;
|
||||
typename BlockStoreIds::TempStorage store_ids;
|
||||
|
||||
} temp_storage;
|
||||
|
||||
for (int vpt_i = 0; vpt_i < C_VPT; ++vpt_i) {
|
||||
int x = vpt_i * C_THREADS + threadIdx.x;
|
||||
|
||||
if (x < N) {
|
||||
Dtype sum = 0.f;
|
||||
for (int dpi = 0; dpi < Dp; ++dpi) {
|
||||
// Dtype val = d_data[b * N * Dp + dpi * N +
|
||||
//x]
|
||||
//- d_data[b * N * Dp + dpi * N + y];
|
||||
Dtype val = d_data[b * N * Dp + dpi * N + x] - s_point[dpi];
|
||||
sum += val * val;
|
||||
}
|
||||
thread_dists[vpt_i] = sqrt(sum);
|
||||
thread_ids[vpt_i] = x;
|
||||
} else {
|
||||
thread_dists[vpt_i] = std::numeric_limits<Dtype>::max();
|
||||
thread_ids[vpt_i] = -1;
|
||||
}
|
||||
}
|
||||
|
||||
BlockRadixSort(temp_storage.sort).Sort(thread_dists, thread_ids);
|
||||
__syncthreads();
|
||||
|
||||
BlockStoreIds(temp_storage.store_ids)
|
||||
.Store(&d_knn_ids[b * N * K + y * K], thread_ids, K);
|
||||
__syncthreads();
|
||||
|
||||
BlockStoreDists(temp_storage.store_dists)
|
||||
.Store(&d_knn_dists[b * N * K + y * K], thread_dists, K);
|
||||
}
|
||||
|
||||
const Dtype* d_data;
|
||||
|
||||
Dtype* d_knn_dists;
|
||||
NBtype* d_knn_ids;
|
||||
|
||||
int N;
|
||||
int Dp;
|
||||
int K;
|
||||
};
|
||||
|
||||
template <typename Dtype, typename NBtype>
|
||||
struct BlockBFKernelAttributesSetter {
|
||||
template <typename T>
|
||||
void setAttributes(T& kernel) {
|
||||
kernel.d_data = d_data;
|
||||
kernel.d_knn_dists = d_knn_dists;
|
||||
kernel.d_knn_ids = d_knn_ids;
|
||||
|
||||
kernel.N = N;
|
||||
kernel.Dp = Dp;
|
||||
kernel.K = K;
|
||||
}
|
||||
|
||||
const Dtype* d_data;
|
||||
|
||||
Dtype* d_knn_dists;
|
||||
NBtype* d_knn_ids;
|
||||
|
||||
int N;
|
||||
int Dp;
|
||||
int K;
|
||||
};
|
||||
|
||||
namespace tensorflow {
|
||||
namespace functor {
|
||||
|
||||
template <typename Dtype, typename NBtype>
|
||||
struct KnnBruteforceFunctor<GPUDevice, Dtype, NBtype> {
|
||||
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& positions,
|
||||
Tensor* neighborhood_out, Tensor* distances,
|
||||
Tensor* timings) {
|
||||
// printf("GPU: KNNBF! \n");
|
||||
|
||||
// printf("return_timings: %d \n",return_timings);
|
||||
|
||||
const int B = positions.dim_size(0);
|
||||
const int D = positions.dim_size(1);
|
||||
const int N = positions.dim_size(2);
|
||||
|
||||
const int K = neighborhood_out->dim_size(2);
|
||||
// printf("B: %d | N: %d | K: %d \n", B, N, K);
|
||||
|
||||
BlockBFKernelAttributesSetter<Dtype, NBtype> attr;
|
||||
attr.d_data = positions.flat<Dtype>().data();
|
||||
attr.d_knn_ids = neighborhood_out->flat<NBtype>().data();
|
||||
attr.d_knn_dists = distances->flat<Dtype>().data();
|
||||
|
||||
attr.N = N;
|
||||
attr.Dp = D;
|
||||
attr.K = K;
|
||||
|
||||
cudaEvent_t start, stop;
|
||||
if (return_timings) {
|
||||
cudaEventCreate(&start);
|
||||
cudaEventCreate(&stop);
|
||||
cudaEventRecord(start, 0);
|
||||
}
|
||||
|
||||
// TODO: beautify
|
||||
if (N <= 32) {
|
||||
BlockBFKernel<Dtype, NBtype, 32, 1> knn;
|
||||
attr.setAttributes(knn);
|
||||
knn.launch(B);
|
||||
} else if (N <= 64) {
|
||||
BlockBFKernel<Dtype, NBtype, 64, 1> knn;
|
||||
attr.setAttributes(knn);
|
||||
knn.launch(B);
|
||||
} else if (N <= 128) {
|
||||
BlockBFKernel<Dtype, NBtype, 128, 1> knn;
|
||||
attr.setAttributes(knn);
|
||||
knn.launch(B);
|
||||
} else if (N <= 256) {
|
||||
BlockBFKernel<Dtype, NBtype, 128, 2> knn;
|
||||
attr.setAttributes(knn);
|
||||
knn.launch(B);
|
||||
} else if (N <= 512) {
|
||||
BlockBFKernel<Dtype, NBtype, 128, 4> knn;
|
||||
attr.setAttributes(knn);
|
||||
knn.launch(B);
|
||||
} else if (N <= 1024) {
|
||||
BlockBFKernel<Dtype, NBtype, 256, 4> knn;
|
||||
attr.setAttributes(knn);
|
||||
knn.launch(B);
|
||||
} else if (N <= 2048) {
|
||||
BlockBFKernel<Dtype, NBtype, 256, 8> knn;
|
||||
attr.setAttributes(knn);
|
||||
knn.launch(B);
|
||||
} else if (N <= 4096) {
|
||||
BlockBFKernel<Dtype, NBtype, 512, 8> knn;
|
||||
attr.setAttributes(knn);
|
||||
knn.launch(B);
|
||||
} else if (N <= 1024 * 8) {
|
||||
BlockBFKernel<Dtype, NBtype, 1024, 8> knn;
|
||||
attr.setAttributes(knn);
|
||||
knn.launch(B);
|
||||
} else {
|
||||
printf(
|
||||
"point sets greater then 8k are note yet supported!! Change to "
|
||||
"knn_graph operation!! \n");
|
||||
}
|
||||
|
||||
if (return_timings) {
|
||||
float time;
|
||||
cudaEventRecord(stop, 0);
|
||||
cudaEventSynchronize(stop);
|
||||
cudaEventElapsedTime(&time, start, stop);
|
||||
// printf("complete elapsed time for KNNBf: %f ms
|
||||
//\n", time);
|
||||
Dtype time2 = time;
|
||||
|
||||
cudaMemcpy(timings->flat<Dtype>().data(), &time2, sizeof(Dtype),
|
||||
cudaMemcpyHostToDevice);
|
||||
|
||||
cudaEventDestroy(start);
|
||||
cudaEventDestroy(stop);
|
||||
}
|
||||
|
||||
cudaDeviceSynchronize();
|
||||
getLastCudaError("KNNBF execution failed");
|
||||
checkCudaErrors(cudaDeviceSynchronize());
|
||||
}
|
||||
|
||||
bool return_timings;
|
||||
};
|
||||
|
||||
template struct KnnBruteforceFunctor<GPUDevice, float, int>;
|
||||
// // too much shared memory
|
||||
// template struct KnnBruteforceFunctor<GPUDevice, double, int>;
|
||||
|
||||
} // namespace functor
|
||||
} // namespace tensorflow
|
||||
|
||||
#endif // GOOGLE_CUDA
|
||||
@@ -0,0 +1,81 @@
|
||||
/* 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.
|
||||
==============================================================================*/
|
||||
// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
|
||||
|
||||
#include "knn_bruteforce_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 KnnBruteforceOp : public OpKernel {
|
||||
public:
|
||||
explicit KnnBruteforceOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
|
||||
OP_REQUIRES_OK(ctx, ctx->GetAttr("K", &K));
|
||||
OP_REQUIRES_OK(ctx, ctx->GetAttr("return_timings", &return_timings));
|
||||
}
|
||||
|
||||
void Compute(OpKernelContext* ctx) override {
|
||||
const Tensor& positions = ctx->input(0);
|
||||
const Tensor& neighborhood_in = ctx->input(1);
|
||||
|
||||
const int B = positions.shape().dim_size(0);
|
||||
const int N = positions.shape().dim_size(2);
|
||||
|
||||
Tensor* neighborhood_out = nullptr;
|
||||
Tensor* distances = nullptr;
|
||||
Tensor* timings = nullptr;
|
||||
|
||||
OP_REQUIRES_OK(ctx, ctx->allocate_output(0, TensorShape({B, N, K}),
|
||||
&neighborhood_out));
|
||||
|
||||
OP_REQUIRES_OK(ctx,
|
||||
ctx->allocate_output(1, TensorShape({B, N, K}), &distances));
|
||||
|
||||
OP_REQUIRES_OK(ctx, ctx->allocate_output(2, TensorShape({1}), &timings));
|
||||
|
||||
::tensorflow::functor::KnnBruteforceFunctor<Device, Dtype, int> knnBFF;
|
||||
knnBFF.return_timings = return_timings;
|
||||
knnBFF(ctx, positions, neighborhood_out, distances, timings);
|
||||
}
|
||||
|
||||
private:
|
||||
TF_DISALLOW_COPY_AND_ASSIGN(KnnBruteforceOp);
|
||||
|
||||
int K;
|
||||
bool return_timings;
|
||||
// int subBatch;
|
||||
};
|
||||
|
||||
#define REGISTER_CUSTOM_OP(NAME, DEVICE, T) \
|
||||
REGISTER_KERNEL_BUILDER( \
|
||||
Name(#NAME).Device(DEVICE_##DEVICE).TypeConstraint<T>("T"), \
|
||||
NAME##Op<DEVICE##Device, T>)
|
||||
|
||||
REGISTER_CUSTOM_OP(KnnBruteforce, CPU, float);
|
||||
|
||||
#ifdef GOOGLE_CUDA
|
||||
REGISTER_CUSTOM_OP(KnnBruteforce, GPU, float);
|
||||
#endif // GOOGLE_CUDA
|
||||
#undef REGISTER_CUSTOM_OP
|
||||
|
||||
} // namespace tensorflow
|
||||
@@ -0,0 +1,44 @@
|
||||
/* 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.
|
||||
==============================================================================*/
|
||||
// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
|
||||
|
||||
#ifndef LIB_KNN_BF_OP_H_
|
||||
#define LIB_KNN_BF_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, typename NBtype>
|
||||
struct KnnBruteforceFunctor {
|
||||
void operator()(::tensorflow::OpKernelContext *ctx, const Tensor &positions,
|
||||
Tensor *neighborhood_out, Tensor *distances, Tensor *timings);
|
||||
|
||||
bool return_timings;
|
||||
};
|
||||
|
||||
} // namespace functor
|
||||
} // namespace tensorflow
|
||||
|
||||
#endif // LIB_KNN_BF_OP_H_
|
||||
@@ -0,0 +1,128 @@
|
||||
#!/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 tabulate import tabulate
|
||||
from scipy.spatial.distance import pdist, squareform
|
||||
|
||||
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=1, 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
|
||||
|
||||
dtype = np.float64
|
||||
|
||||
def find_neighbors(positions, K):
|
||||
# B, Dpos, N
|
||||
all_neighbors = []
|
||||
for batch in positions:
|
||||
distances = squareform(pdist(batch.T, 'euclidean'))
|
||||
all_neighbors.append(np.argsort(distances, axis=1)[:, :K])
|
||||
return np.array(all_neighbors).transpose(0, 2, 1)
|
||||
|
||||
def random_values(shape):
|
||||
# return (np.random.randn(*shape) * 100).astype(np.int32).astype(np.float32)
|
||||
# return (np.random.randn(*shape) * 100).astype(np.int32).astype(np.float32)
|
||||
return np.random.randn(*shape).astype(np.float32)
|
||||
|
||||
self.theta = random_values([1, self.Dp, self.Din, self.Dout]).astype(dtype)
|
||||
self.bias = random_values([self.Din, self.Dout]).astype(dtype)
|
||||
|
||||
self.position = random_values([self.B, self.Dp, self.N]).astype(dtype)
|
||||
self.features = random_values([self.B, self.Din, self.N]).astype(dtype)
|
||||
self.neighborhood = find_neighbors(
|
||||
self.position, self.K).astype(dtype=np.int32)
|
||||
|
||||
def init_ops(self, dtype=np.float32):
|
||||
self.theta_op = tf.convert_to_tensor(self.theta.astype(dtype))
|
||||
self.bias_op = tf.convert_to_tensor(self.bias.astype(dtype))
|
||||
|
||||
self.features_op = tf.convert_to_tensor(self.features.astype(dtype))
|
||||
self.position_op = tf.convert_to_tensor(self.position.astype(dtype))
|
||||
self.neighborhood_op = tf.convert_to_tensor(self.neighborhood)
|
||||
|
||||
def expected_feature_shape(self):
|
||||
return [self.B, self.Din, self.N]
|
||||
|
||||
def expected_output_shape(self):
|
||||
return [self.B, self.Dout, self.N]
|
||||
|
||||
|
||||
class VerboseTestCase(tf.test.TestCase):
|
||||
|
||||
def assertAllClose(self, a, b, rtol=1e-6, atol=1e-6):
|
||||
max_outputs = 20
|
||||
|
||||
def max_tol(b):
|
||||
return atol + rtol * np.abs(b)
|
||||
|
||||
if not np.allclose(a, b, rtol=rtol, atol=atol):
|
||||
cond = np.logical_or(
|
||||
np.abs(a - b) > atol + rtol * np.abs(b),
|
||||
np.isnan(a) != np.isnan(b))
|
||||
|
||||
lines = []
|
||||
if a.ndim:
|
||||
|
||||
shape = a.shape
|
||||
|
||||
a = a.flatten()
|
||||
b = b.flatten()
|
||||
|
||||
cond = np.logical_or(
|
||||
np.abs(a - b) > atol + rtol * np.abs(b),
|
||||
np.isnan(a) != np.isnan(b))
|
||||
|
||||
idxArr = np.arange(a.shape[0])[np.where(cond)]
|
||||
xArr = a[np.where(cond)]
|
||||
yArr = b[np.where(cond)]
|
||||
for idx, x, y in zip(idxArr, xArr, yArr):
|
||||
idx = np.unravel_index(idx, shape)
|
||||
lines.append((idx, x, y, np.abs(x - y), max_tol(y)))
|
||||
max_outputs -= 1
|
||||
if max_outputs == 0:
|
||||
break
|
||||
print(tabulate(lines, headers=["index", "actual", "expected",
|
||||
"diff", "max-tol"]))
|
||||
print("diff (sum): ", np.abs(a - b).sum())
|
||||
print("diff (max): ", np.abs(a - b).max())
|
||||
print("diff (mean): ", np.abs(a - b).mean())
|
||||
else:
|
||||
# np.where is broken for scalars
|
||||
x, y = a, b
|
||||
lines.append((x, y, np.abs(x - y), max_tol(y)))
|
||||
print(tabulate(lines, headers=["actual", "expected",
|
||||
"diff", "max-tol"]))
|
||||
|
||||
assert np.allclose(a, b, rtol=rtol, atol=atol, equal_nan=True), "failed"
|
||||
@@ -0,0 +1,41 @@
|
||||
// ComputerGraphics Tuebingen, 2018
|
||||
|
||||
#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("KnnBruteforce")
|
||||
.Input("position: T") // position: each datapoint in nd space [B,Dp, N].
|
||||
.Output("neighborhood_out: NBtype") // neighborhood_out: all K nearest
|
||||
// neighbors [B, N, K].
|
||||
.Output("distances: T") // distances: all K nearest distances
|
||||
// [B, N, K].
|
||||
.Output("timings: T") // timings:
|
||||
// [1]
|
||||
.Attr("K: int")
|
||||
.Attr("return_timings: bool = false")
|
||||
.Attr("T: realnumbertype")
|
||||
.Attr("NBtype: {int32} = DT_INT32")
|
||||
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
|
||||
const auto position = c->input(0);
|
||||
const auto neighborhood_out = c->input(1);
|
||||
|
||||
int K;
|
||||
c->GetAttr("K", &K);
|
||||
|
||||
auto B = c->Dim(position, 0);
|
||||
auto N = c->Dim(position, 2);
|
||||
|
||||
c->set_output(0, c->MakeShape({B, N, K}));
|
||||
c->set_output(1, c->MakeShape({B, N, K}));
|
||||
|
||||
return Status::OK();
|
||||
});
|
||||
|
||||
} // namespace tensorflow
|
||||
|
||||
// doc: K: number of neighbors.
|
||||
@@ -0,0 +1,29 @@
|
||||
#!/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 unittest
|
||||
import glob
|
||||
|
||||
operations = [fn[:-3] for fn in glob.glob('test_*.py') if not '_all' in fn]
|
||||
suite = unittest.TestSuite()
|
||||
|
||||
for op in operations:
|
||||
suite.addTest(unittest.defaultTestLoader.loadTestsFromName(op))
|
||||
|
||||
unittest.TextTestRunner().run(suite)
|
||||
@@ -18,91 +18,119 @@
|
||||
# Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
|
||||
|
||||
|
||||
from PointTestCase import TPC, PointTestCase, summary
|
||||
from misc import FakePointCloud, VerboseTestCase
|
||||
import tensorflow as tf
|
||||
|
||||
import numpy as np
|
||||
from __init__ import flex_convolution
|
||||
|
||||
case = FakePointCloud(B=2, N=32, K=4, Din=2, Dout=6, Dp=3)
|
||||
|
||||
class FlexConvTest(PointTestCase):
|
||||
|
||||
class FlexConvTest(VerboseTestCase):
|
||||
def __init__(self, methodName="runTest"):
|
||||
super(FlexConvTest, self).__init__(methodName)
|
||||
|
||||
def _forward(self, use_gpu=False, force_gpu=False):
|
||||
self.init_ops()
|
||||
def _forward(self, use_gpu=False, force_gpu=False, dtype=np.float32):
|
||||
case.init_ops(dtype=dtype)
|
||||
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_op = flex_convolution(case.features_op,
|
||||
case.position_op, case.neighborhood_op,
|
||||
case.theta_op, case.bias_op)
|
||||
actual = sess.run(actual_op)
|
||||
return actual
|
||||
|
||||
def test_forward(self):
|
||||
def test_forward(self, dtype=np.float32):
|
||||
cpu = self._forward(use_gpu=False)
|
||||
gpu = self._forward(use_gpu=True)
|
||||
self.assertAllClose(cpu, gpu, 1e-5, 1e-5)
|
||||
self.assertAllClose(cpu, gpu, 1e-4)
|
||||
|
||||
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')
|
||||
def test_forward_features_gpu_floats(self):
|
||||
cpu32 = self._forward(use_gpu=True, dtype=np.float32)
|
||||
cpu64 = self._forward(use_gpu=True, dtype=np.float64)
|
||||
self.assertAllClose(cpu32, cpu64)
|
||||
|
||||
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_features(self, use_gpu=False, dtype=np.float32, numdiff=True):
|
||||
case.init_ops(dtype=dtype)
|
||||
with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
|
||||
actual_op = flex_convolution(case.features_op,
|
||||
case.position_op, case.neighborhood_op,
|
||||
case.theta_op, case.bias_op)
|
||||
if numdiff:
|
||||
return tf.test.compute_gradient(
|
||||
[case.features_op], [case.features.shape], actual_op,
|
||||
case.expected_output_shape())[0]
|
||||
else:
|
||||
return sess.run(tf.gradients(actual_op, [case.features_op]))[0]
|
||||
|
||||
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)
|
||||
def _backward_bias(self, use_gpu=False, dtype=np.float32, numdiff=True):
|
||||
case.init_ops(dtype=dtype)
|
||||
with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
|
||||
actual_op = flex_convolution(case.features_op,
|
||||
case.position_op, case.neighborhood_op,
|
||||
case.theta_op, case.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')
|
||||
if numdiff:
|
||||
return tf.test.compute_gradient(
|
||||
[case.bias_op], [case.bias.shape], actual_op,
|
||||
case.expected_output_shape())[0]
|
||||
else:
|
||||
return sess.run(tf.gradients(actual_op, [case.bias_op]))[0]
|
||||
|
||||
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, dtype=np.float32, numdiff=True):
|
||||
case.init_ops(dtype=dtype)
|
||||
with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
|
||||
actual_op = flex_convolution(case.features_op,
|
||||
case.position_op, case.neighborhood_op,
|
||||
case.theta_op, case.bias_op)
|
||||
|
||||
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)
|
||||
if numdiff:
|
||||
return tf.test.compute_gradient(
|
||||
[case.theta_op], [case.theta.shape], actual_op,
|
||||
case.expected_output_shape())[0]
|
||||
else:
|
||||
return sess.run(tf.gradients(actual_op, [case.theta_op]))[0]
|
||||
|
||||
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')
|
||||
def test_backward_features_cpu_float64(self):
|
||||
actual, expected = self._backward_features(use_gpu=False, dtype=np.float64)
|
||||
self.assertAllClose(actual, expected)
|
||||
|
||||
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_bias_cpu_float64(self):
|
||||
actual, expected = self._backward_bias(use_gpu=False, dtype=np.float64)
|
||||
self.assertAllClose(actual, expected)
|
||||
|
||||
def test_backward_theta_cpu_float64(self):
|
||||
actual, expected = self._backward_theta(use_gpu=False, dtype=np.float64)
|
||||
self.assertAllClose(actual, expected)
|
||||
|
||||
def test_backward_features_gpu_float64(self):
|
||||
actual, expected = self._backward_features(use_gpu=True, dtype=np.float64)
|
||||
self.assertAllClose(actual, expected)
|
||||
|
||||
def test_backward_bias_gpu_float64(self):
|
||||
actual, expected = self._backward_bias(use_gpu=True, dtype=np.float64)
|
||||
self.assertAllClose(actual, expected)
|
||||
|
||||
def test_backward_theta_gpu_float64(self):
|
||||
actual, expected = self._backward_theta(use_gpu=True, dtype=np.float64)
|
||||
self.assertAllClose(actual, expected)
|
||||
|
||||
def test_backward_features_gpu_float32(self, dtype=np.float32):
|
||||
cpu = self._backward_features(use_gpu=False, dtype=dtype, numdiff=False)
|
||||
gpu = self._backward_features(use_gpu=True, dtype=dtype, numdiff=False)
|
||||
self.assertAllClose(cpu, gpu)
|
||||
|
||||
def test_backward_bias_gpu_float32(self, dtype=np.float32):
|
||||
cpu = self._backward_bias(use_gpu=False, dtype=dtype, numdiff=False)
|
||||
gpu = self._backward_bias(use_gpu=True, dtype=dtype, numdiff=False)
|
||||
self.assertAllClose(cpu, gpu, 1e-5)
|
||||
|
||||
def test_backward_theta_gpu_float32(self, dtype=np.float32):
|
||||
cpu = self._backward_theta(use_gpu=False, dtype=dtype, numdiff=False)
|
||||
gpu = self._backward_theta(use_gpu=True, dtype=dtype, numdiff=False)
|
||||
self.assertAllClose(cpu, gpu)
|
||||
|
||||
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__':
|
||||
|
||||
@@ -18,98 +18,115 @@
|
||||
# Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
|
||||
|
||||
|
||||
from PointTestCase import TPC, PointTestCase, summary
|
||||
from misc import FakePointCloud, VerboseTestCase
|
||||
import tensorflow as tf
|
||||
import numpy as np
|
||||
from __init__ import flex_convolution_transpose
|
||||
|
||||
case = FakePointCloud(B=2, N=32, K=4, Din=2, Dout=6, Dp=3)
|
||||
|
||||
class FlexDeconvTest(PointTestCase):
|
||||
|
||||
class FlexConvTest(VerboseTestCase):
|
||||
def __init__(self, methodName="runTest"):
|
||||
super(FlexDeconvTest, self).__init__(methodName)
|
||||
super(FlexConvTest, self).__init__(methodName)
|
||||
|
||||
def _forward(self, use_gpu=False, force_gpu=False):
|
||||
self.init_ops()
|
||||
case.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)
|
||||
|
||||
actual_op = flex_convolution_transpose(case.features_op,
|
||||
case.position_op,
|
||||
case.neighborhood_op,
|
||||
case.theta_op, case.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)
|
||||
summary(cpu, gpu, 'self.features')
|
||||
self.assertAllClose(cpu, gpu, rtol=1e-05)
|
||||
self.assertAllClose(cpu, gpu, 1e-4)
|
||||
|
||||
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)
|
||||
def _backward_features(self, use_gpu=False, dtype=np.float32, numdiff=True):
|
||||
case.init_ops(dtype=dtype)
|
||||
with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
|
||||
actual_op = flex_convolution_transpose(case.features_op,
|
||||
case.position_op,
|
||||
case.neighborhood_op,
|
||||
case.theta_op, case.bias_op)
|
||||
if numdiff:
|
||||
return tf.test.compute_gradient(
|
||||
[case.features_op], [case.features.shape], actual_op,
|
||||
case.expected_output_shape())[0]
|
||||
else:
|
||||
return sess.run(tf.gradients(actual_op, [case.features_op]))[0]
|
||||
|
||||
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)
|
||||
def _backward_bias(self, use_gpu=False, dtype=np.float32, numdiff=True):
|
||||
case.init_ops(dtype=dtype)
|
||||
with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
|
||||
actual_op = flex_convolution_transpose(case.features_op,
|
||||
case.position_op,
|
||||
case.neighborhood_op,
|
||||
case.theta_op, case.bias_op)
|
||||
|
||||
err = tf.test.compute_gradient_error([self.features_op],
|
||||
[self.features.shape],
|
||||
actual_op, TPC.expected_output_shape())
|
||||
self.assertLess(err, 1e-2)
|
||||
if numdiff:
|
||||
return tf.test.compute_gradient(
|
||||
[case.bias_op], [case.bias.shape], actual_op,
|
||||
case.expected_output_shape())[0]
|
||||
else:
|
||||
return sess.run(tf.gradients(actual_op, [case.bias_op]))[0]
|
||||
|
||||
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)
|
||||
def _backward_theta(self, use_gpu=False, dtype=np.float32, numdiff=True):
|
||||
case.init_ops(dtype=dtype)
|
||||
with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
|
||||
actual_op = flex_convolution_transpose(case.features_op,
|
||||
case.position_op,
|
||||
case.neighborhood_op,
|
||||
case.theta_op, case.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')
|
||||
if numdiff:
|
||||
return tf.test.compute_gradient(
|
||||
[case.theta_op], [case.theta.shape], actual_op,
|
||||
case.expected_output_shape())[0]
|
||||
else:
|
||||
return sess.run(tf.gradients(actual_op, [case.theta_op]))[0]
|
||||
|
||||
err = tf.test.compute_gradient_error([self.bias_op],
|
||||
[self.bias.shape], actual_op,
|
||||
TPC.expected_output_shape())
|
||||
self.assertLess(err, 1e-2)
|
||||
def test_backward_features_cpu_float64(self):
|
||||
actual, expected = self._backward_features(use_gpu=False, dtype=np.float64)
|
||||
self.assertAllClose(actual, expected)
|
||||
|
||||
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)
|
||||
def test_backward_bias_cpu_float64(self):
|
||||
actual, expected = self._backward_bias(use_gpu=False, dtype=np.float64)
|
||||
self.assertAllClose(actual, expected)
|
||||
|
||||
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')
|
||||
def test_backward_theta_cpu_float64(self):
|
||||
actual, expected = self._backward_theta(use_gpu=False, dtype=np.float64)
|
||||
self.assertAllClose(actual, expected)
|
||||
|
||||
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_gpu_float64(self):
|
||||
actual, expected = self._backward_features(use_gpu=True, dtype=np.float64)
|
||||
self.assertAllClose(actual, expected)
|
||||
|
||||
def test_backward_features(self):
|
||||
self._backward_features(use_gpu=False)
|
||||
self._backward_features(use_gpu=True)
|
||||
def test_backward_bias_gpu_float64(self):
|
||||
actual, expected = self._backward_bias(use_gpu=True, dtype=np.float64)
|
||||
self.assertAllClose(actual, expected)
|
||||
|
||||
def test_backward_bias(self):
|
||||
self._backward_bias(use_gpu=False)
|
||||
self._backward_bias(use_gpu=True)
|
||||
def test_backward_theta_gpu_float64(self):
|
||||
actual, expected = self._backward_theta(use_gpu=True, dtype=np.float64)
|
||||
self.assertAllClose(actual, expected)
|
||||
|
||||
def test_backward_theta(self):
|
||||
self._backward_theta(use_gpu=False)
|
||||
self._backward_theta(use_gpu=True)
|
||||
def test_backward_features_gpu_float32(self, dtype=np.float32):
|
||||
cpu = self._backward_features(use_gpu=False, dtype=dtype, numdiff=False)
|
||||
gpu = self._backward_features(use_gpu=True, dtype=dtype, numdiff=False)
|
||||
self.assertAllClose(cpu, gpu)
|
||||
|
||||
def test_backward_bias_gpu_float32(self, dtype=np.float32):
|
||||
cpu = self._backward_bias(use_gpu=False, dtype=dtype, numdiff=False)
|
||||
gpu = self._backward_bias(use_gpu=True, dtype=dtype, numdiff=False)
|
||||
self.assertAllClose(cpu, gpu, 1e-5)
|
||||
|
||||
def test_backward_theta_gpu_float32(self, dtype=np.float32):
|
||||
cpu = self._backward_theta(use_gpu=False, dtype=dtype, numdiff=False)
|
||||
gpu = self._backward_theta(use_gpu=True, dtype=dtype, numdiff=False)
|
||||
self.assertAllClose(cpu, gpu)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
@@ -18,70 +18,85 @@
|
||||
# Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
|
||||
|
||||
|
||||
from PointTestCase import PointTestCase
|
||||
from misc import FakePointCloud, VerboseTestCase
|
||||
import tensorflow as tf
|
||||
import numpy as np
|
||||
|
||||
from __init__ import flex_pooling
|
||||
|
||||
case = FakePointCloud(B=2, N=32, K=4, Din=2, Dout=6, Dp=3)
|
||||
|
||||
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
|
||||
class FlexPoolTest(VerboseTestCase):
|
||||
def __init__(self, methodName="runTest"):
|
||||
super(FlexPoolTest, self).__init__(methodName)
|
||||
|
||||
def test_forward(self):
|
||||
cpu = self._forward(use_gpu=False)
|
||||
gpu = self._forward(use_gpu=True)
|
||||
self.assertAllClose(cpu, gpu)
|
||||
def _forward(self, use_gpu=False, dtype=np.float32):
|
||||
case.init_ops(dtype=dtype)
|
||||
with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
|
||||
actual_op, winner_op = flex_pooling(
|
||||
case.features_op, case.neighborhood_op)
|
||||
actual = sess.run(actual_op)
|
||||
return actual
|
||||
|
||||
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 test_forward_same_float32(self):
|
||||
cpu = self._forward(use_gpu=False, dtype=np.float32)
|
||||
gpu = self._forward(use_gpu=True, dtype=np.float32)
|
||||
self.assertAllClose(cpu, gpu)
|
||||
|
||||
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)
|
||||
def test_forward_same_float64(self):
|
||||
cpu = self._forward(use_gpu=False, dtype=np.float64)
|
||||
gpu = self._forward(use_gpu=True, dtype=np.float64)
|
||||
self.assertAllClose(cpu, gpu)
|
||||
|
||||
graph_features_grad = tf.gradients(actual_op, [self.features_op])[0]
|
||||
def test_backward_same_float32(self):
|
||||
cpu, winner_cpu = self._backward(use_gpu=False, dtype=np.float32)
|
||||
gpu, winner_gpu = self._backward(use_gpu=True, dtype=np.float32)
|
||||
self.assertAllClose(winner_cpu, winner_gpu)
|
||||
self.assertAllClose(cpu, gpu)
|
||||
|
||||
dx, winner = sess.run([graph_features_grad, winner_op])
|
||||
return dx, winner
|
||||
def test_backward_same_float64(self):
|
||||
cpu, winner_cpu = self._backward(use_gpu=False, dtype=np.float64)
|
||||
gpu, winner_gpu = self._backward(use_gpu=True, dtype=np.float64)
|
||||
self.assertAllClose(winner_cpu, winner_gpu)
|
||||
self.assertAllClose(cpu, gpu)
|
||||
|
||||
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)
|
||||
def _backward(self, use_gpu=False, dtype=np.float32):
|
||||
case.init_ops(dtype=dtype)
|
||||
with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
|
||||
actual_op, winner_op = flex_pooling(
|
||||
case.features_op,
|
||||
case.neighborhood_op)
|
||||
|
||||
x = tf.convert_to_tensor(x.astype(np.float32))
|
||||
n = tf.convert_to_tensor(n.astype(np.int32))
|
||||
graph_features_grad = tf.gradients(actual_op, [case.features_op])[0]
|
||||
|
||||
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)
|
||||
dx, winner = sess.run([graph_features_grad, winner_op])
|
||||
return dx, winner
|
||||
|
||||
def test_backward_simple(self):
|
||||
cpu = self._simple_backward(use_gpu=False)
|
||||
cpu[0, 0, 2] -= 4
|
||||
self.assertEqual(cpu.sum(), 0)
|
||||
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)
|
||||
|
||||
gpu = self._simple_backward(use_gpu=True)
|
||||
gpu[0, 0, 2] -= 4
|
||||
self.assertEqual(gpu.sum(), 0)
|
||||
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_cpu(self):
|
||||
cpu = self._simple_backward(use_gpu=False)
|
||||
cpu[0, 0, 2] -= 4
|
||||
self.assertEqual(cpu.sum(), 0)
|
||||
|
||||
def test_backward_simple_gpu(self):
|
||||
gpu = self._simple_backward(use_gpu=True)
|
||||
gpu[0, 0, 2] -= 4
|
||||
self.assertEqual(gpu.sum(), 0)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
tf.test.main()
|
||||
tf.test.main()
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
#!/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 misc import FakePointCloud, VerboseTestCase
|
||||
import tensorflow as tf
|
||||
import numpy as np
|
||||
from scipy.spatial.distance import pdist, squareform
|
||||
|
||||
from __init__ import knn_bruteforce
|
||||
|
||||
case = FakePointCloud(B=2, N=32, K=4, Din=2, Dout=6, Dp=3)
|
||||
case = FakePointCloud(B=1, N=4, K=2, Din=1, Dout=1, Dp=3)
|
||||
|
||||
|
||||
def python_bruteforce(positions, K):
|
||||
# B, Dpos, N
|
||||
all_neighbors = []
|
||||
all_distances = []
|
||||
for batch in positions:
|
||||
distances = squareform(pdist(batch.T, 'euclidean'))
|
||||
all_neighbors.append(np.argsort(distances, axis=1)[:, :K])
|
||||
all_distances.append(np.sort(distances, axis=1)[:, :K])
|
||||
return np.array(all_neighbors), np.array(all_distances)
|
||||
|
||||
|
||||
class KnnBruteforceTest(VerboseTestCase):
|
||||
def __init__(self, methodName="runTest"):
|
||||
super(KnnBruteforceTest, self).__init__(methodName)
|
||||
|
||||
def _forward(self, use_gpu=False):
|
||||
case.init_ops()
|
||||
expected_nn, expected_dist = python_bruteforce(case.position, K=4)
|
||||
|
||||
with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
|
||||
actual_nn, actual_dist, _ = knn_bruteforce(case.position_op, K=4)
|
||||
actual_nn, actual_dist = sess.run([actual_nn, actual_dist])
|
||||
|
||||
self.assertAllClose(expected_dist, actual_dist)
|
||||
self.assertAllClose(expected_nn, actual_nn)
|
||||
|
||||
def test_forward_cpu(self):
|
||||
self._forward(use_gpu=False)
|
||||
|
||||
def test_forward_gpu(self):
|
||||
self._forward(use_gpu=True)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
tf.test.main()
|
||||
Reference in New Issue
Block a user