Annotate the op_info inputs with extra information when possible (e.g. the input value if it's known statically).
| 78 | // Annotate the op_info inputs with extra information when possible (e.g. the |
| 79 | // input value if it's known statically). |
| 80 | static void ExtractExtraProperties( |
| 81 | const NodeDef& node, |
| 82 | const std::unordered_map<string, const NodeDef*>& name_to_node, |
| 83 | OpInfo* op_info) { |
| 84 | OpRegistry* op_registry = OpRegistry::Global(); |
| 85 | const OpDef* op_def = nullptr; |
| 86 | auto s = op_registry->LookUpOpDef(node.op(), &op_def); |
| 87 | if (!s.ok()) { |
| 88 | op_def = nullptr; |
| 89 | } |
| 90 | |
| 91 | for (int i = 0; i < node.input_size(); ++i) { |
| 92 | const string input_name = node.input(i); |
| 93 | CHECK(!input_name.empty()); |
| 94 | if (IsControlInput(input_name)) { |
| 95 | continue; |
| 96 | } |
| 97 | TensorId input_tensor_id = ParseTensorName(input_name); |
| 98 | const string input_node_name(input_tensor_id.first); |
| 99 | |
| 100 | auto iter = name_to_node.find(input_node_name); |
| 101 | if (iter == name_to_node.end()) continue; |
| 102 | const NodeDef* input_node = iter->second; |
| 103 | |
| 104 | if (i >= op_info->inputs_size()) { |
| 105 | LOG(ERROR) << "OpInfo's inputs doesn't match the graph! OpInfo: " |
| 106 | << op_info->DebugString() |
| 107 | << "\nCurrent node: " << node.DebugString() |
| 108 | << "\nInput node: " << input_node->DebugString(); |
| 109 | } |
| 110 | |
| 111 | // The value attribute in Const input is useful for cost prediction. |
| 112 | if (input_node->op() == "Const" && i < op_info->inputs_size()) { |
| 113 | auto it = input_node->attr().find("value"); |
| 114 | if (it == input_node->attr().end()) continue; |
| 115 | |
| 116 | const AttrValue& attr_value = it->second; |
| 117 | std::vector<TensorProto> tensors = ExtractTensors(attr_value); |
| 118 | if (tensors.empty()) continue; |
| 119 | |
| 120 | const TensorProto& t = tensors[0]; |
| 121 | OpInfo::TensorProperties* input = op_info->mutable_inputs(i); |
| 122 | *(input->mutable_value()) = t; |
| 123 | |
| 124 | // For filename input, the file size can also be useful. |
| 125 | if (op_def && i < op_def->input_arg_size() && |
| 126 | op_def->input_arg(i).name().find("filename") != string::npos) { |
| 127 | Tensor tensor; |
| 128 | if (!tensor.FromProto(t)) { |
| 129 | continue; |
| 130 | } |
| 131 | if (tensor.NumElements() != 1) { |
| 132 | continue; |
| 133 | } |
| 134 | const string& filename = tensor.scalar<tstring>()(); |
| 135 | |
| 136 | Env* env = Env::Default(); |
| 137 | FileStatistics stat; |
no test coverage detected