loco::Pull *pullLayer()
{
loco::Pull *pull = _graph.nodes()->create<loco::Pull>();
- _graph.inputs()->create()->node(pull);
+
+ auto graph_input = _graph.inputs()->create();
+ graph_input->name("graph_input");
+ graph_input->node(pull);
+
pull->dtype(loco::DataType::FLOAT32);
setSampleShape(pull);
return pull;
loco::Push *pushLayer(loco::Node *input)
{
loco::Push *push = _graph.nodes()->create<loco::Push>();
- _graph.outputs()->create()->node(push);
+
+ auto graph_output = _graph.outputs()->create();
+ graph_output->name("graph_output");
+ graph_output->node(push);
+
push->from(input);
return push;
}
// This version is taken from comment in fbs
constexpr uint32_t version = 3;
+ registerGraphIOName(graph, gd);
+
// parse graph into SerializedModelData structure
exportOpDefinedTensors(graph->nodes(), _builder, gd);
return tflite::Padding_SAME;
}
+void registerGraphIOName(loco::Graph *graph, SerializedModelData &gd)
+{
+ for (uint32_t in = 0; in < graph->inputs()->size(); ++in)
+ {
+ gd._input_names.emplace_back(graph->inputs()->at(in)->name());
+ }
+ for (uint32_t out = 0; out < graph->outputs()->size(); ++out)
+ {
+ gd._output_names.emplace_back(graph->outputs()->at(out)->name());
+ }
+}
+
} // namepsace loco_exporter
#define __LOCO_EXPORTER_UTILS_H__
#include "schema_generated.h"
-#include "loco/IR/Nodes.h"
+#include "loco.h"
#include "loco/IR/PermutingCodec.h"
std::unordered_map<loco::Node *, tflite::TensorType> _node_to_type;
std::unordered_map<loco::Node *, ShapeDescription> _node_to_shape;
+ // Graph input and output names
+ std::vector<std::string> _input_names;
+ std::vector<std::string> _output_names;
+
/**
* @brief if opcode is not registered in table of opcodes add it
* @param builtin_code
tflite::Padding getOpPadding(loco::MaxPool2D *node);
+/// @brief Register graph input and output names to SerializedModelData
+void registerGraphIOName(loco::Graph *graph, SerializedModelData &gd);
+
} // namespace loco_exporter
#endif // __LOCO_EXPORTER_UTILS_H__
// encode and register tensor itself using attributes from previous steps
auto tensor_id = static_cast<uint32_t>(gd._tensors.size());
- std::string name = "t_" + std::to_string(tensor_id);
+ std::string name;
+ if (dynamic_cast<loco::Pull *>(node)) // current node is input
+ {
+ // TODO Now only supports single input. Let's support multi inputs
+ assert(gd._input_names.size() == 1);
+ name = gd._input_names[0];
+ }
+ else if (dynamic_cast<loco::Push *>(*loco::succs(node).begin())) // next node is output
+ {
+ // TODO Now only supports single output. Let's support multi outputs
+ assert(gd._output_names.size() == 1);
+ name = gd._output_names[0];
+ }
+ else
+ {
+ name = "t_" + std::to_string(tensor_id);
+ }
auto name_offset = builder.CreateString(name);
auto tensor_offset = CreateTensor(builder, shape_offset, tensor_type, buffer_id, name_offset,
/*quantization*/ 0, /*is_variable*/ false);