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

Add more operation and some refactoring of the code to ease usage.

Browse Source
This commit is contained in:
Fabian Groh
2018-11-06 14:00:16 +01:00
committed by PatWie
parent 32e91a688f
commit 2cbeb7e699
27 changed files with 1264 additions and 452 deletions
Download Patch File Download Diff File Expand all files Collapse all files
user_ops/kernels/flex_pool_kernel_gpu.cu.cc
+12 -18
View File
@@ -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