| 1145 | } |
| 1146 | |
| 1147 | Status Converter::ConvertNode(const NodeDef& node_def) { |
| 1148 | std::vector<TRT_TensorOrWeights> inputs, outputs; |
| 1149 | TF_RETURN_IF_ERROR(this->GetInputs(node_def, &inputs)); |
| 1150 | |
| 1151 | OpConverterParams params(this, node_def, inputs, &outputs, &weight_store_); |
| 1152 | const string& op = node_def.op(); |
| 1153 | auto itr = op_registry_.find(op); |
| 1154 | if (itr == op_registry_.end()) { |
| 1155 | return errors::Unimplemented("No converter registered for op: ", op); |
| 1156 | } |
| 1157 | OpConverter op_converter = itr->second; |
| 1158 | TF_RETURN_IF_ERROR(op_converter(¶ms)); |
| 1159 | |
| 1160 | for (size_t i = 0; i < outputs.size(); ++i) { |
| 1161 | TRT_TensorOrWeights& output = outputs[i]; |
| 1162 | string output_name = node_def.name(); |
| 1163 | if (i != 0) absl::StrAppend(&output_name, ":", i); |
| 1164 | // We need to check the name before setting it. If the input is one of the |
| 1165 | // engine input, setting the name here will overwrite engine input |
| 1166 | // bindings which will cause runtime error. |
| 1167 | // TODO(tmorris): Remove this work-around once we use TRT's IIdentityLayer |
| 1168 | // in ConvertIdentity. |
| 1169 | if (output.is_tensor()) { |
| 1170 | const char* tensor_name = output.tensor()->getName(); |
| 1171 | if (!IsEngineInput(tensor_name)) { |
| 1172 | // TRT initializes tensor names as "(Unnamed ITensor* N)". We rename |
| 1173 | // them to match their corresponding TensorFlow name. |
| 1174 | // Note: ITensors that we create internally within TF-TRT which are |
| 1175 | // not inputs or outputs of a node will not be renamed. This is a |
| 1176 | // potential cause of confusion if an error message or warning |
| 1177 | // mentions the unnamed tensor. |
| 1178 | output.tensor()->setName(output_name.c_str()); |
| 1179 | } |
| 1180 | } |
| 1181 | VLOG(2) << "Adding out tensor " << output_name << ": " |
| 1182 | << output.DebugString(); |
| 1183 | Status status = AddTensorOrWeights(output_name, output); |
| 1184 | if (!status.ok()) { |
| 1185 | return Status(status.code(), |
| 1186 | StrCat("Failed to add output for node ", node_def.name(), |
| 1187 | ": ", status.error_message())); |
| 1188 | } |
| 1189 | } |
| 1190 | return Status::OK(); |
| 1191 | } |
| 1192 | |
| 1193 | Status Converter::AddInputTensor(const string& name, nvinfer1::DataType dtype, |
| 1194 | const nvinfer1::Dims& dims, int batch_size) { |