diff --git a/tools/tensorflow/tensorflow2ncnn.cpp b/tools/tensorflow/tensorflow2ncnn.cpp index 669316d7e..4cf5ed583 100644 --- a/tools/tensorflow/tensorflow2ncnn.cpp +++ b/tools/tensorflow/tensorflow2ncnn.cpp @@ -169,21 +169,38 @@ int main(int argc, char** argv) weights[output_name] = tensorflow::TensorProto(); continue; } - else if (node.op() == "Add" || node.op() == "BiasAdd" - || node.op() == "Max" || node.op() == "Maximum" || node.op() == "Mul" - || node.op() == "RealDiv" || node.op() == "Sub") + else { - // check weights - for (int j=0; j::iterator it = weights.find(input_name); - if (it != weights.end()) + if (isBinaryOp) + { + // check weights + for (int j=0; jsecond; - weights.erase(it); + const std::string& input_name = node.input(j); + + std::map::iterator it = weights.find(input_name); + if (it != weights.end()) + { + // binary op with const, insert MemoryData layer and const blob + binaryop_consts[input_name] = it->second; + weights.erase(it); + } } } }