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_deconv_op.cc
+17 -11
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_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