wassname/Flex-Convolution
1
0
Fork 0
mirror of https://github.com/wassname/Flex-Convolution.git synced 2026-09-11 11:52:15 +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_conv_op.cc
+16 -26
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
#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