ecsimsw

TF_source ) REGISTER_KERNELS 본문

TF_source ) REGISTER_KERNELS

JinHwan Kim 2019. 4. 7. 07:02

"CODE"

 

///// tensorflow/core/kernels/training_ops.cc

#define REGISTER_KERNELS(T, Tindices)                                 \
  REGISTER_KERNEL_BUILDER(Name("SparseApplyProximalAdagrad")          \
                              .Device(DEVICE_CPU)                     \
                              .TypeConstraint("T")                 \
                              .TypeConstraint("Tindices"),  \
                          SparseApplyProximalAdagradOp<T, Tindices>); \
  REGISTER_KERNEL_BUILDER(Name("ResourceSparseApplyProximalAdagrad")  \
                              .Device(DEVICE_CPU)                     \
                              .TypeConstraint("T")                 \
                              .TypeConstraint("Tindices"),  \
                          SparseApplyProximalAdagradOp<T, Tindices>);

 

//// tensorflow/core/framework/op_kernel.h

 

#define REGISTER_KERNEL_BUILDER(kernel_builder, ...) \
  REGISTER_KERNEL_BUILDER_UNIQ_HELPER(__COUNTER__, kernel_builder, __VA_ARGS__)

#define REGISTER_KERNEL_BUILDER_UNIQ_HELPER(ctr, kernel_builder, ...) \
  REGISTER_KERNEL_BUILDER_UNIQ(ctr, kernel_builder, __VA_ARGS__)

#define REGISTER_KERNEL_BUILDER_UNIQ(ctr, kernel_builder, ...)        \
  constexpr bool should_register_##ctr##__flag =                      \
      SHOULD_REGISTER_OP_KERNEL(#__VA_ARGS__);                        \
  static ::tensorflow::kernel_factory::OpKernelRegistrar              \
      registrar__body__##ctr##__object(                               \
          should_register_##ctr##__flag                               \
              ? ::tensorflow::register_kernel::kernel_builder.Build() \
              : nullptr,                                              \
          #__VA_ARGS__,                                               \
          [](::tensorflow::OpKernelConstruction* context)             \
              -> ::tensorflow::OpKernel* {                            \
            return new __VA_ARGS__(context);                          \
          });

 

"NEXT"

- OpKernel

'Machine Learning > tf_source' 카테고리의 다른 글

TF_source ) KernelRegistry  (0) 2019.04.20
TF_source) Opkernel  (0) 2019.04.12
TF_source) OpkernelConstruction  (0) 2019.04.12
TF_source ) Factory  (0) 2019.04.08
TF_source ) OpKernelRegistrar  (0) 2019.04.07
Comments