* @param protocols protocols of embeddings * @param shapes shapes of embeddings * @param sampler_protocol protocol of negative sampling * @param num_moment number of moment statistics in optimizers * @param num_partition numebr of partitions * @param num_negative number of negative samples per postive sample * @param batch_size batch size of samples in CPU-GPU transfe
| 825 | * @return GPU memory cost |
| 826 | */ |
| 827 | static size_t gpu_memory_demand(const std::vector<Protocol> &protocols, const std::vector<Index> &shapes, |
| 828 | Protocol sampler_protocol = kTailPartition, int num_moment = 0, |
| 829 | int num_partition = 4, int num_negative = 1, int batch_size = 10000) { |
| 830 | auto partition_shapes = shapes; |
| 831 | int num_embedding = protocols.size(); |
| 832 | Index num_vertex; |
| 833 | for (int i = 0; i < num_embedding; i++) { |
| 834 | Protocol protocol = protocols[i]; |
| 835 | if (protocol & (kHeadPartition | kTailPartition)) { |
| 836 | num_vertex = shapes[i]; |
| 837 | partition_shapes[i] = (shapes[i] + num_partition - 1) / num_partition; |
| 838 | } |
| 839 | if (protocol & kSharedWithPredecessor) { |
| 840 | if (num_partition == 1 || (protocol & kGlobal)) |
| 841 | partition_shapes[i] = 0; |
| 842 | if ((protocol & (kHeadPartition | kTailPartition)) == |
| 843 | (protocols[i - 1] & (kHeadPartition | kTailPartition))) |
| 844 | partition_shapes[i] = 0; |
| 845 | } |
| 846 | } |
| 847 | Index sampler_size = 0; |
| 848 | if (sampler_protocol | kHeadPartition) |
| 849 | sampler_size += (num_vertex + num_partition - 1) / num_partition; |
| 850 | if (sampler_protocol | kTailPartition) |
| 851 | sampler_size += (num_vertex + num_partition - 1) / num_partition; |
| 852 | if (sampler_protocol | kGlobal) |
| 853 | sampler_size += num_vertex; |
| 854 | |
| 855 | size_t demand = 0; |
| 856 | demand += Sampler::gpu_memory_demand(); |
| 857 | demand += Worker::gpu_memory_demand(protocols, partition_shapes, sampler_size, num_moment, num_negative, |
| 858 | batch_size); |
| 859 | return demand; |
| 860 | } |
| 861 | |
| 862 | static size_t cpu_memory_demand() { |
| 863 | size_t demand = 0; |
nothing calls this directly
no outgoing calls
no test coverage detected