Dice(x) = [gamma + (1-gamma)p(x)]*x p(x) = sigmoid[(x-mean)*rvar] Let t = (x-mean) * rvar p(x) = 1/(1 + exp(-t)) Dice(x) = [gamma + (1-gamma)/(1 + exp(-t)) ] * x Dice(x) = [1 + gamma * exp(-t)] / [1 + exp(-t)] * x Dice(x) = [gamma + exp(t)] / [1 + exp(t)] * x
| 24 | // Dice(x) = [1 + gamma * exp(-t)] / [1 + exp(-t)] * x |
| 25 | // Dice(x) = [gamma + exp(t)] / [1 + exp(t)] * x |
| 26 | void Compute(OpKernelContext* context) override { |
| 27 | // Grab the input |
| 28 | const Tensor* x_tensor = &context->input(0); |
| 29 | const Tensor* mean_tensor = &context->input(1); |
| 30 | const Tensor* rvar_tensor = &context->input(2); |
| 31 | const Tensor* gamma_tensor = &context->input(3); |
| 32 | |
| 33 | const T* x = x_tensor->flat<T>().data(); |
| 34 | const T* mean = mean_tensor->flat<T>().data(); |
| 35 | const T* rvar = rvar_tensor->flat<T>().data(); |
| 36 | const T* gamma = gamma_tensor->flat<T>().data(); |
| 37 | |
| 38 | int64 cols = x_tensor->dim_size(x_tensor->dims() - 1); |
| 39 | int64 rows = 1; |
| 40 | for (int64 i = 0; i < x_tensor->dims() - 1; ++i) { |
| 41 | rows *= x_tensor->dim_size(i); |
| 42 | } |
| 43 | |
| 44 | int64 gamma_size = 1; |
| 45 | for (int64 i = 0; i < gamma_tensor->dims(); ++i) { |
| 46 | gamma_size *= gamma_tensor->dim_size(i); |
| 47 | } |
| 48 | |
| 49 | // To check the input |
| 50 | OP_REQUIRES(context, (mean_tensor->dims() == 1), |
| 51 | errors::InvalidArgument("dims(mean) != 1")); |
| 52 | OP_REQUIRES(context, (rvar_tensor->dims() == 1), |
| 53 | errors::InvalidArgument("dims(rvar) != 1")); |
| 54 | |
| 55 | OP_REQUIRES( |
| 56 | context, (mean_tensor->dim_size(0) == cols), |
| 57 | errors::InvalidArgument("size(mean) != last_dim_size_of_input")); |
| 58 | OP_REQUIRES( |
| 59 | context, (rvar_tensor->dim_size(0) == cols), |
| 60 | errors::InvalidArgument("size(rvar) != last_dim_size_of_input")); |
| 61 | OP_REQUIRES( |
| 62 | context, (gamma_size == cols), |
| 63 | errors::InvalidArgument("size(gamma) != last_dim_size_of_input")); |
| 64 | |
| 65 | // Create output tensors |
| 66 | Tensor* y_tensor = NULL; |
| 67 | OP_REQUIRES_OK(context, |
| 68 | context->allocate_output(0, x_tensor->shape(), &y_tensor)); |
| 69 | T* y = y_tensor->flat<T>().data(); |
| 70 | |
| 71 | // Do it |
| 72 | // Let every thread compute 16 rows to avoid false sharing |
| 73 | const int64 total_unit = (rows + 15) / 16; |
| 74 | const int64 unit_cost = |
| 75 | 16 * cols * 50; // assume every element consumes 50 cycles |
| 76 | |
| 77 | auto& worker_threads = |
| 78 | *(context->device()->tensorflow_cpu_worker_threads()); |
| 79 | thread::ThreadPool* thread_pool = worker_threads.workers; |
| 80 | |
| 81 | thread_pool->ParallelFor( |
| 82 | total_unit, unit_cost, [&](int64 begin_unit, int64 end_unit) { |
| 83 | auto begin_row = begin_unit * 16; |
nothing calls this directly
no test coverage detected