| 74 | } |
| 75 | |
| 76 | Status AutoParallel::Initialize(const GrapplerItem& item) { |
| 77 | num_gpus_ = GetNumAvailableGPUs(); |
| 78 | LOG(INFO) << "Number of GPUs: " << num_gpus_; |
| 79 | item_ = &item; |
| 80 | graph_ = item.graph; |
| 81 | LOG(INFO) << "Original graph size: " << graph_.node_size(); |
| 82 | if (item.fetch.empty()) { |
| 83 | return Status(error::INVALID_ARGUMENT, "No fetch nodes provided."); |
| 84 | } |
| 85 | |
| 86 | if (item.MainVariables().empty()) { |
| 87 | return Status(error::INVALID_ARGUMENT, "No variables provided."); |
| 88 | } |
| 89 | |
| 90 | for (const auto& init : item.init_ops) { |
| 91 | VLOG(1) << "Init node: " << init; |
| 92 | } |
| 93 | |
| 94 | for (const auto& fetch : item.fetch) { |
| 95 | VLOG(1) << "Fetch node: " << fetch; |
| 96 | } |
| 97 | |
| 98 | for (const auto& var : item.MainVariables()) { |
| 99 | VLOG(2) << "Variable: " << var->name(); |
| 100 | } |
| 101 | |
| 102 | const std::set<string> apply_gradients_ops = {"ApplyGradientDescent", |
| 103 | "ApplyProximalGradientDescent", |
| 104 | "ApplyAdadelta", |
| 105 | "ApplyAdagrad", |
| 106 | "ApplyProximalAdagrad", |
| 107 | "ApplyAdagradDA", |
| 108 | "ApplyFtrl", |
| 109 | "ApplyMomentum", |
| 110 | "ApplyAdam", |
| 111 | "ApplyRMSProp", |
| 112 | "ApplyCenteredRMSProp"}; |
| 113 | for (int i = 0; i < graph_.node_size(); i++) { |
| 114 | all_nodes_.insert( |
| 115 | std::make_pair(graph_.node(i).name(), graph_.mutable_node(i))); |
| 116 | if (apply_gradients_ops.find(graph_.node(i).op()) != |
| 117 | apply_gradients_ops.end()) { |
| 118 | apply_gradients_nodes_.insert(graph_.node(i).name()); |
| 119 | VLOG(2) << "Apply gradients node: " << graph_.node(i).name(); |
| 120 | } |
| 121 | } |
| 122 | |
| 123 | auto div_const_node = AddNodeDivConst(); |
| 124 | all_nodes_.insert(std::make_pair(div_const_node->name(), div_const_node)); |
| 125 | std::map<string, int> gradient_pos = {{"ApplyGradientDescent", 2}, |
| 126 | {"ApplyProximalGradientDescent", 4}, |
| 127 | {"ApplyAdadelta", 6}, |
| 128 | {"ApplyAdagrad", 3}, |
| 129 | {"ApplyProximalAdagrad", 5}, |
| 130 | {"ApplyAdagradDA", 3}, |
| 131 | {"ApplyFtrl", 3}, |
| 132 | {"ApplyMomentum", 3}, |
| 133 | {"ApplyAdam", 9}, |
no test coverage detected