| 808 | } |
| 809 | |
| 810 | void Model::OptimizeGradientDescent(int64 cpu_budget, int64 ram_budget) { |
| 811 | std::shared_ptr<Node> snapshot; |
| 812 | { |
| 813 | tf_shared_lock lock(mu_); |
| 814 | snapshot = output_->Snapshot(nullptr); |
| 815 | } |
| 816 | VLOG(2) << "Starting optimization of tunable parameters with GradientDescent"; |
| 817 | auto parameters = CollectTunableParameters(snapshot); |
| 818 | auto essential_parameters = CollectEssentialParallelism(snapshot); |
| 819 | // We add the number of model's buffered bytes because it is excluded from the |
| 820 | // memory budget, but it is included in the maximum number of buffered bytes. |
| 821 | ram_budget += TotalBufferedBytes(snapshot); |
| 822 | for (auto& pair : parameters) { |
| 823 | pair.second->value = pair.second->min; |
| 824 | } |
| 825 | // Gradient descent step size. |
| 826 | constexpr double kDescentStep = 0.1L; |
| 827 | |
| 828 | // Optimization is stopped once the `OutputTime` improvement is smaller than |
| 829 | // this value. |
| 830 | constexpr double kOptimizationPrecision = 100.0L; |
| 831 | |
| 832 | // Maximum number of iterations for optimization. |
| 833 | constexpr int64 kMaxIterations = 1000; |
| 834 | |
| 835 | double output_time = 0; |
| 836 | double new_output_time; |
| 837 | double new_value; |
| 838 | for (int i = 0; i < kMaxIterations; ++i) { |
| 839 | std::map<string, double> gradient; |
| 840 | new_output_time = OutputTime(snapshot, &gradient); |
| 841 | int64 model_parallelism = 0; |
| 842 | for (auto& pair : essential_parameters) { |
| 843 | model_parallelism += std::round(pair.second->value); |
| 844 | } |
| 845 | // We terminate once the improvement of the output latency is too small or |
| 846 | // the essential transformations' parallelism reaches the CPU budget or the |
| 847 | // worst-case total buffer size exceeds the memory budget. |
| 848 | if (std::abs(output_time - new_output_time) < kOptimizationPrecision || |
| 849 | model_parallelism > cpu_budget || |
| 850 | TotalMaximumBufferedBytes(snapshot) > ram_budget) { |
| 851 | break; |
| 852 | } |
| 853 | double max_abs_derivative = 1.0; |
| 854 | for (auto& pair : parameters) { |
| 855 | if (pair.second->value != pair.second->max) { |
| 856 | max_abs_derivative = |
| 857 | std::max(max_abs_derivative, std::abs(gradient[pair.first])); |
| 858 | } |
| 859 | } |
| 860 | for (auto& pair : parameters) { |
| 861 | new_value = pair.second->value - |
| 862 | kDescentStep * gradient[pair.first] / max_abs_derivative; |
| 863 | // Projection on a feasible interval. |
| 864 | if (new_value > pair.second->max) { |
| 865 | pair.second->value = pair.second->max; |
| 866 | } else if (new_value < pair.second->min) { |
| 867 | pair.second->value = pair.second->min; |