| 36 | }; |
| 37 | |
| 38 | void SampleNodeSplit::Compute(const DAGNodeProto& node_def, |
| 39 | OpKernelContext* ctx) { |
| 40 | // get split num |
| 41 | int32_t split_num = QueryProxy::GetInstance()->GetShardNum(); |
| 42 | // get count tensor |
| 43 | Tensor* count_tensor = nullptr; |
| 44 | ctx->tensor(node_def.inputs(0), &count_tensor); |
| 45 | int32_t count = count_tensor->Raw<int32_t>()[0]; |
| 46 | // get node_type tensor |
| 47 | Tensor* node_type_tensor = nullptr; |
| 48 | ctx->tensor(node_def.inputs(1), &node_type_tensor); |
| 49 | std::vector<int32_t> node_types(node_type_tensor->NumElements()); |
| 50 | std::copy(node_type_tensor->Raw<int32_t>(), |
| 51 | node_type_tensor->Raw<int32_t>() + node_types.size(), |
| 52 | node_types.begin()); |
| 53 | |
| 54 | const std::vector<std::vector<float>>& shard_node_weight = |
| 55 | QueryProxy::GetInstance()->GetShardNodeWeight(); |
| 56 | |
| 57 | std::vector<int32_t> split_cnt(split_num); |
| 58 | int32_t remain_cnt = count; |
| 59 | std::vector<size_t> no_zero_idxs; no_zero_idxs.reserve(split_num); |
| 60 | for (int32_t i = 0; i < split_num; ++i) { |
| 61 | float sum_weight0 = 0, sum_weight1 = 0; |
| 62 | for (int32_t node_type : node_types) { |
| 63 | if (node_type == -1) { |
| 64 | node_type = shard_node_weight.size() - 1; |
| 65 | if (node_types.size() > 1) { |
| 66 | EULER_LOG(FATAL) << "ERROR node type!"; |
| 67 | } |
| 68 | } |
| 69 | sum_weight0 += shard_node_weight[node_type][i]; |
| 70 | sum_weight1 += shard_node_weight[node_type][split_num]; |
| 71 | } |
| 72 | if (fabs(sum_weight1 - 0.0) < 0.0000001) { |
| 73 | EULER_LOG(FATAL) << "node type sum weight is zero!"; |
| 74 | } |
| 75 | split_cnt[i] = floor(count * sum_weight0 / sum_weight1); |
| 76 | if (sum_weight0 > 0) { |
| 77 | no_zero_idxs.push_back(i); |
| 78 | } |
| 79 | remain_cnt -= split_cnt[i]; |
| 80 | } |
| 81 | for (size_t i = 0; remain_cnt > 0; ++i, --remain_cnt) { |
| 82 | size_t no_zero_idx = no_zero_idxs[ |
| 83 | floor(common::ThreadLocalRandom() * no_zero_idxs.size())]; |
| 84 | split_cnt[no_zero_idx] += 1; |
| 85 | } |
| 86 | |
| 87 | /* output */ |
| 88 | int32_t output_idx = 0; |
| 89 | int32_t begin_id = 0; |
| 90 | for (int32_t i = 0; i < split_num; ++i) { |
| 91 | // output split intput |
| 92 | std::string split_count_out_name = |
| 93 | OutputName(node_def, output_idx); |
| 94 | ++output_idx; |
| 95 | Tensor* split_count_t = nullptr; |
nothing calls this directly
no test coverage detected