Class EBatchNormGradOp

en_ops/e_batch_norm.cc:106–173

REGISTER_KERNEL_BUILDER(Name("EBatchNorm").Device(DEVICE_CPU), EBatchNormOp);

class EBatchNormGradOp : public OpKernel
{
public:
    explicit EBatchNormGradOp(OpKernelConstruction *context) : OpKernel(context)
    {
        context->GetAttr("epsilon", &epsilon_);
        OP_REQUIRES_OK(context, context->GetAttr("eid_low", &eid_low_));
        OP_REQUIRES_OK(context, context->GetAttr("eid_high", &eid_high_));
        OP_REQUIRES_OK(context, context->GetAttr("times", &times_));
        lib = dlopen("/home/zhangyan/jhrsgx/privacy_test/privacy_tf/sgx_tf_ops/sgx.so", RTLD_LAZY);
        OP_REQUIRES(context, lib != NULL, errors::Unknown("Unable to load sgx.so!"));
    }
    void Compute(OpKernelContext *context) override
    {

        const Tensor &grad = context->input(0);
        const Tensor &input = context->input(1);
        const Tensor &scale = context->input(2);
        const Tensor &offset = context->input(3);
        const Tensor &mean = context->input(4);
        const Tensor &variance = context->input(5);

        Tensor *grad_input = NULL;
        Tensor *grad_scale = NULL;
        Tensor *grad_offset = NULL;
        Tensor *grad_mean = NULL;
        Tensor *grad_var = NULL;
        OP_REQUIRES_OK(context, context->allocate_output(0, input.shape(), &grad_input));
        OP_REQUIRES_OK(context, context->allocate_output(1, scale.shape(), &grad_scale));
        OP_REQUIRES_OK(context, context->allocate_output(2, offset.shape(), &grad_offset));
        OP_REQUIRES_OK(context, context->allocate_output(3, mean.shape(), &grad_mean));
        OP_REQUIRES_OK(context, context->allocate_output(4, variance.shape(), &grad_var));

        auto grad_flat = grad.flat<float>();
        auto input_flat = input.flat<float>();
        auto scale_flat = scale.flat<float>();
        auto offset_flat = offset.flat<float>();
        auto mean_flat = mean.flat<float>();
        auto variance_flat = variance.flat<float>();

        auto grad_input_flat = grad_input->flat<float>();
        auto grad_scale_flat = grad_scale->flat<float>();
        auto grad_offset_flat = grad_offset->flat<float>();
        auto grad_mean_flat = grad_mean->flat<float>();
        auto grad_var_flat = grad_var->flat<float>();

        const TensorShape &input_shape = input.shape();
        int M = input_shape.dim_size(0) / times_;
        int H = input_shape.dim_size(1);
        int W = input_shape.dim_size(2);
        const int N = input_flat.size() / times_;

        unsigned long int eid_ = (eid_high_ << 32) + eid_low_;
        typedef void (*function)(unsigned long int eid, float *input, float *grad, float *scale, float *offset, float *mean_x, float *variance, int N, int M, int H, int W, float *grad_in, float *grad_sc, float *grad_off, float *grad_mean, float *grad_var, float epsilon);
        dlerror();
        function batchnorm_grad_kernel = (function)dlsym(lib, "batchnorm_grad");
        const char *dlsym_error = dlerror();
        OP_REQUIRES(context, !dlsym_error, errors::Unknown("loading of batchnorm_grad failed: ", dlsym_error));