compiler/tflchef/core/src/ModelChef.cpp

   1 /*
   2  * Copyright (c) 2018 Samsung Electronics Co., Ltd. All Rights Reserved
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *    http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #include "tflchef/ModelChef.h"
  18 #include <souschef/RangedArguments.h>
  19 #include <souschef/Registry.h>
  20
  21 #include "Convert.h"
  22
  23 #include <souschef/DataChefs.h>
  24
  25 #include "OpChef.h"
  26 #include "OpChefs.h"
  27
  28 #include <souschef/Dataset.h>
  29
  30 #include "Log.h"
  31
  32 #include <iterator>
  33 #include <map>
  34 #include <string>
  35 #include <vector>
  36
  37 #include <cassert>
  38 #include <fstream>
  39 #include <iostream>
  40 #include <numeric>
  41 #include <sstream>
  42 #include <stdexcept>
  43
  44 namespace
  45 {
  46
  47 using namespace souschef;
  48
  49 template <typename T> std::vector<T> as_vector(const ::google::protobuf::RepeatedPtrField<T> &field)
  50 {
  51   std::vector<T> res;
  52   for (const auto &elem : field)
  53   {
  54     res.emplace_back(elem);
  55   }
  56   return res;
  57 }
  58
  59 template <typename T> Dataset<T> as_dataset(const ::google::protobuf::RepeatedPtrField<T> &field)
  60 {
  61   return Dataset<T>(as_vector<T>(field));
  62 }
  63
  64 } // namespace
  65
  66 namespace
  67 {
  68
  69 template <typename T> using Dims = std::vector<T>;
  70
  71 Dims<int32_t> as_dims(const tflchef::TensorShape &shape)
  72 {
  73   std::vector<int32_t> res;
  74
  75   for (auto &dim : shape.dim())
  76   {
  77     res.emplace_back(static_cast<int32_t>(dim));
  78   }
  79
  80   return res;
  81 }
  82
  83 int32_t element_count(const Dims<int32_t> &dims)
  84 {
  85   return std::accumulate(dims.begin(), dims.end(), 1, std::multiplies<int32_t>());
  86 }
  87
  88 } // namespace
  89
  90 namespace
  91 {
  92
  93 class GeneratedModelImpl final : public tflchef::GeneratedModel::Impl
  94 {
  95 public:
  96   GeneratedModelImpl(std::unique_ptr<flatbuffers::FlatBufferBuilder> &&builder)
  97       : _builder{std::move(builder)}
  98   {
  99     // DO NOTHING
 100   }
 101
 102 public:
 103   const char *base(void) const override
 104   {
 105     // Return the base address of generated flatbuffer model
 106     return reinterpret_cast<const char *>(_builder->GetBufferPointer());
 107   }
 108
 109 public:
 110   size_t size(void) const override
 111   {
 112     // Return the size of generated flatbuffer model
 113     return _builder->GetSize();
 114   }
 115
 116 private:
 117   std::unique_ptr<flatbuffers::FlatBufferBuilder> _builder;
 118 };
 119
 120 } // namespace
 121
 122 namespace
 123 {
 124
 125 struct DataChefRegistry final : public Registry<DataChefFactory>
 126 {
 127 };
 128
 129 DataChefRegistry &data_chef_registry(const tflchef::TensorType &type)
 130 {
 131   static DataChefRegistry s32;
 132   static DataChefRegistry s64;
 133   static DataChefRegistry fp32;
 134   static DataChefRegistry u8;
 135   static DataChefRegistry boolean;
 136
 137   switch (type)
 138   {
 139     case tflchef::INT32:
 140       return s32;
 141     case tflchef::INT64:
 142       return s64;
 143     case tflchef::FLOAT32:
 144       return fp32;
 145     case tflchef::UINT8:
 146       return u8;
 147     case tflchef::BOOL:
 148       return boolean;
 149     default:
 150       break;
 151   }
 152
 153   throw std::runtime_error{"Unknown tensor type"};
 154 }
 155
 156 struct OpChefRegistry final : public Registry<OpChefFactory>
 157 {
 158 };
 159
 160 OpChefRegistry &op_chef_registry(void)
 161 {
 162   static OpChefRegistry registry;
 163   return registry;
 164 }
 165
 166 /// @brief This will prepare a map of unique builtin codes in the model recipe
 167 std::map<tflite::BuiltinOperator, int32_t>
 168 gather_builtincode_map(const ::tflchef::ModelRecipe &model_recipe)
 169 {
 170   // Key and value of the map are BuiltinOperator and operator version
 171   std::map<tflite::BuiltinOperator, int32_t> builtin_map;
 172
 173   for (const auto &operation : model_recipe.operation())
 174   {
 175     auto op_chef = op_chef_registry().lookup(operation.type()).create(&operation);
 176     if (op_chef->code() == tflite::BuiltinOperator_CUSTOM)
 177       continue;
 178
 179     // Various operation version is unified as the highest version among them
 180     if (builtin_map.find(op_chef->code()) == builtin_map.end() ||
 181         builtin_map[op_chef->code()] < operation.version())
 182       builtin_map[op_chef->code()] = operation.version();
 183   }
 184
 185   // Add ops used in Graphs(subgraphs)
 186   for (int g = 0; g < model_recipe.graph_size(); ++g)
 187   {
 188     const auto &graph = model_recipe.graph(g);
 189     for (const auto &operation : graph.operation())
 190     {
 191       auto op_chef = op_chef_registry().lookup(operation.type()).create(&operation);
 192       if (op_chef->code() == tflite::BuiltinOperator_CUSTOM)
 193         continue;
 194
 195       // Various operation version is unified as the highest version among them
 196       if (builtin_map.find(op_chef->code()) == builtin_map.end() ||
 197           builtin_map[op_chef->code()] < operation.version())
 198         builtin_map[op_chef->code()] = operation.version();
 199     }
 200   }
 201
 202   return builtin_map;
 203 }
 204
 205 /// @brief This will prepare a set of unique custom codes in the mode recipe
 206 std::set<std::string> gather_customcode_set(const ::tflchef::ModelRecipe &model_recipe)
 207 {
 208   std::set<std::string> customcode_set;
 209   for (const auto &operation : model_recipe.operation())
 210   {
 211     auto op_chef = op_chef_registry().lookup(operation.type()).create(&operation);
 212     if (op_chef->code() == tflite::BuiltinOperator_CUSTOM)
 213       customcode_set.insert(operation.type());
 214   }
 215
 216   // Add ops used in Graphs(subgraphs)
 217   for (int g = 0; g < model_recipe.graph_size(); ++g)
 218   {
 219     const auto &graph = model_recipe.graph(g);
 220     for (const auto &operation : graph.operation())
 221     {
 222       auto op_chef = op_chef_registry().lookup(operation.type()).create(&operation);
 223       if (op_chef->code() == tflite::BuiltinOperator_CUSTOM)
 224         customcode_set.insert(operation.type());
 225     }
 226   }
 227
 228   return customcode_set;
 229 }
 230
 231 } // namespace
 232
 233 namespace
 234 {
 235
 236 struct CookParams
 237 {
 238   std::vector<flatbuffers::Offset<::tflite::Buffer>> &buffer_vec;
 239   std::vector<flatbuffers::Offset<::tflite::OperatorCode>> &code_vec;
 240   std::vector<flatbuffers::Offset<::tflite::SubGraph>> &subgraph_vec;
 241   std::unique_ptr<flatbuffers::FlatBufferBuilder> &flatbuffer_builder;
 242   std::map<tflite::BuiltinOperator, int32_t> &builtin_code_map;
 243   std::string noname;
 244 };
 245
 246 template <typename T> void cook_graph(const T &graph, CookParams &cp)
 247 {
 248   LOGGER(l);
 249
 250   std::vector<flatbuffers::Offset<::tflite::Buffer>> &buffer_vec = cp.buffer_vec;
 251   std::vector<flatbuffers::Offset<::tflite::OperatorCode>> &code_vec = cp.code_vec;
 252   std::vector<flatbuffers::Offset<::tflite::SubGraph>> &subgraph_vec = cp.subgraph_vec;
 253   std::unique_ptr<flatbuffers::FlatBufferBuilder> &flatbuffer_builder = cp.flatbuffer_builder;
 254   std::map<tflite::BuiltinOperator, int32_t> &builtin_code_map = cp.builtin_code_map;
 255
 256   // Operand-related
 257   std::vector<flatbuffers::Offset<::tflite::Tensor>> tensor_vec;
 258
 259   // Operation-related
 260   std::vector<flatbuffers::Offset<::tflite::Operator>> operator_vec;
 261
 262   // default name for graph
 263   std::string graph_name = cp.noname;
 264   if (graph.has_name())
 265     graph_name = graph.name();
 266
 267   // Tensor Name -> Tensor ID mapping (per Graph)
 268   std::map<std::string, int32_t> symbol_table;
 269
 270   auto lookup = [&symbol_table, &graph_name](const std::string &name) {
 271     if (symbol_table.find(name) != symbol_table.end())
 272       return symbol_table.at(name);
 273     else if (name == "")
 274       return -1; // -1 in TFLite means that optional input tensor is empty.
 275     else
 276     {
 277       std::string msg = "tflchef : input not found in " + graph_name + " graph";
 278       throw std::runtime_error(msg.c_str());
 279     }
 280   };
 281
 282   int32_t buffer_start = buffer_vec.size();
 283   int32_t buffer_index = 0;
 284
 285   // Create buffer(s) 1~n(I) for input(s)
 286   const auto size_input = graph.input_size();
 287   for (int ci = 0; ci < size_input; ++ci)
 288   {
 289     tflite::BufferBuilder buffer_builder{*flatbuffer_builder};
 290     buffer_vec.emplace_back(buffer_builder.Finish());
 291   }
 292   // Create buffer(s) n(I)+1~n(I)+n(O) for output(s)
 293   const auto size_output = graph.output_size();
 294   for (int co = 0; co < size_output; ++co)
 295   {
 296     tflite::BufferBuilder buffer_builder{*flatbuffer_builder};
 297     buffer_vec.emplace_back(buffer_builder.Finish());
 298   }
 299
 300   auto input_names = as_dataset(graph.input()).vectorize();
 301   auto output_names = as_dataset(graph.output()).vectorize();
 302
 303   for (const auto &operand : graph.operand())
 304   {
 305     assert(operand.has_name());
 306
 307     assert(operand.has_type());
 308
 309     flatbuffers::Offset<flatbuffers::Vector<int32_t>> shape;
 310     std::vector<int32_t> dims;
 311     if (operand.has_shape())
 312     {
 313       dims = as_dims(operand.shape());
 314       shape = flatbuffer_builder->CreateVector(dims);
 315     }
 316
 317     auto name = flatbuffer_builder->CreateString(operand.name());
 318
 319     buffer_index = 0;
 320
 321     // Create Buffer if filler is specified
 322     if (operand.has_filler())
 323     {
 324       const auto &filler = operand.filler();
 325
 326       assert(filler.has_tag());
 327
 328       auto args = ranged_arguments(filler.arg().begin(), filler.arg().end());
 329       auto chef = data_chef_registry(operand.type()).lookup(filler.tag()).create(args);
 330
 331       assert(chef != nullptr);
 332
 333       // Create Data
 334       int32_t count = (element_count(dims) > 0) ? element_count(dims) : filler.arg_size();
 335       auto data_vec = chef->generate(count);
 336       auto data = flatbuffer_builder->CreateVector(data_vec);
 337
 338       // Create Buffer
 339       tflite::BufferBuilder buffer_builder{*flatbuffer_builder};
 340       buffer_builder.add_data(data);
 341       auto buffer = buffer_builder.Finish();
 342
 343       // Update Buffer Index & Vector
 344       buffer_index = buffer_vec.size();
 345       buffer_vec.emplace_back(buffer);
 346     }
 347     else
 348     {
 349       // if this is input or output, assign to that buffer_index
 350       int idx = 0;
 351       for (auto it = input_names.begin(); it != input_names.end(); ++it, ++idx)
 352       {
 353         if (*it == operand.name())
 354         {
 355           buffer_index = buffer_start + idx;
 356           break;
 357         }
 358       }
 359       if (buffer_index == 0)
 360       {
 361         idx = 0;
 362         for (auto it = output_names.begin(); it != output_names.end(); ++it, ++idx)
 363         {
 364           if (*it == operand.name())
 365           {
 366             buffer_index = buffer_start + size_input + idx;
 367             break;
 368           }
 369         }
 370       }
 371       if (buffer_index == 0)
 372       {
 373         // we couldn't find the buffer; create an empty buffer for this tensor
 374         buffer_index = buffer_vec.size();
 375
 376         tflite::BufferBuilder buffer_builder{*flatbuffer_builder};
 377         buffer_vec.emplace_back(buffer_builder.Finish());
 378       }
 379     }
 380     assert(buffer_index != 0);
 381
 382     flatbuffers::Offset<tflite::QuantizationParameters> quant_index;
 383
 384     // Create QuantizationParameters if quant is specified
 385     if (operand.has_quant())
 386     {
 387       const auto &quant = operand.quant();
 388
 389       // Create each parameters
 390       // NOTE if some parameters are not given, those will be set to default value
 391       std::vector<float> quant_max_vec(quant.max_size());
 392       std::vector<float> quant_min_vec(quant.min_size());
 393       std::vector<float> quant_scale_vec(quant.scale_size());
 394       std::vector<int64_t> quant_zero_point_vec(quant.zero_point_size());
 395
 396       for (uint32_t i = 0; i < quant.max_size(); ++i)
 397         quant_max_vec.at(i) = quant.max(i);
 398       for (uint32_t i = 0; i < quant.min_size(); ++i)
 399         quant_min_vec.at(i) = quant.min(i);
 400       for (uint32_t i = 0; i < quant.scale_size(); ++i)
 401         quant_scale_vec.at(i) = quant.scale(i);
 402       for (uint32_t i = 0; i < quant.zero_point_size(); ++i)
 403         quant_zero_point_vec.at(i) = quant.zero_point(i);
 404
 405       auto quant_max = flatbuffer_builder->CreateVector(quant_max_vec);
 406       auto quant_min = flatbuffer_builder->CreateVector(quant_min_vec);
 407       auto quant_scale = flatbuffer_builder->CreateVector(quant_scale_vec);
 408       auto quant_zero_point = flatbuffer_builder->CreateVector(quant_zero_point_vec);
 409
 410       // Create QuantizationParameters
 411       tflite::QuantizationParametersBuilder quant_builder{*flatbuffer_builder};
 412       quant_builder.add_max(quant_max);
 413       quant_builder.add_min(quant_min);
 414       quant_builder.add_scale(quant_scale);
 415       quant_builder.add_zero_point(quant_zero_point);
 416       quant_builder.add_quantized_dimension(quant.quantized_dimension());
 417
 418       // Update QuantizationParameters Index
 419       quant_index = quant_builder.Finish();
 420     }
 421
 422     // Create Tensor
 423     tflite::TensorBuilder tensor_builder{*flatbuffer_builder};
 424
 425     tensor_builder.add_shape(shape);
 426     tensor_builder.add_type(as_tflite_tensortype(operand.type()));
 427     tensor_builder.add_buffer(buffer_index);
 428     tensor_builder.add_name(name);
 429     if (operand.has_quant())
 430       tensor_builder.add_quantization(quant_index);
 431
 432     // Append!
 433     tensor_vec.emplace_back(tensor_builder.Finish());
 434
 435     // Update Tensor Name -> Tensor Index Map
 436     int32_t tensor_index = symbol_table.size();
 437     const auto &tensor_name = operand.name();
 438
 439     INFO(l) << "Symbol [" << tensor_name << "] = Tensor " << tensor_index << std::endl;
 440
 441     symbol_table[tensor_name] = tensor_index;
 442   }
 443
 444   // Create Operator
 445   for (const auto &operation : graph.operation())
 446   {
 447     assert(operation.has_type());
 448
 449     auto op_chef = op_chef_registry().lookup(operation.type()).create(&operation);
 450
 451     // Create 'inputs'
 452     std::vector<int32_t> input_vec = as_dataset(operation.input()).map(lookup).vectorize();
 453     auto inputs = flatbuffer_builder->CreateVector(input_vec);
 454
 455     // Create 'outputs'
 456     std::vector<int32_t> output_vec = as_dataset(operation.output()).map(lookup).vectorize();
 457     auto outputs = flatbuffer_builder->CreateVector(output_vec);
 458
 459     // Create Option
 460     auto options = op_chef->value(*flatbuffer_builder);
 461
 462     // Create Custom option
 463     auto circle_custom_options = op_chef->custom_value(*flatbuffer_builder);
 464
 465     // Create Operator
 466     tflite::OperatorBuilder op_builder{*flatbuffer_builder};
 467
 468     // Get operator code index from builtin_code_set with assumption, order of
 469     // builtin_code_set is same as that of code_vec
 470     auto op_it = builtin_code_map.find(op_chef->code());
 471     assert(op_it != builtin_code_map.end());
 472     uint32_t opcode_index = std::distance(builtin_code_map.begin(), op_it);
 473
 474     op_builder.add_opcode_index(opcode_index);
 475     op_builder.add_inputs(inputs);
 476     op_builder.add_outputs(outputs);
 477     op_builder.add_builtin_options_type(op_chef->type());
 478     op_builder.add_builtin_options(options);
 479     op_builder.add_custom_options(circle_custom_options);
 480     op_builder.add_custom_options_format(tflite::CustomOptionsFormat_FLEXBUFFERS);
 481     // Append Operator
 482     operator_vec.emplace_back(op_builder.Finish());
 483   }
 484
 485   // Create network input/output vector
 486   std::vector<int32_t> input_vec = as_dataset(graph.input()).map(lookup).vectorize();
 487   std::vector<int32_t> output_vec = as_dataset(graph.output()).map(lookup).vectorize();
 488
 489   // Create "SubGraph" arguments
 490   auto tensors = flatbuffer_builder->CreateVector(tensor_vec);
 491   auto inputs = flatbuffer_builder->CreateVector(input_vec);
 492   auto outputs = flatbuffer_builder->CreateVector(output_vec);
 493   auto operators = flatbuffer_builder->CreateVector(operator_vec);
 494   auto name = flatbuffer_builder->CreateString(graph_name);
 495
 496   tflite::SubGraphBuilder subgraph_builder{*flatbuffer_builder};
 497
 498   subgraph_builder.add_tensors(tensors);
 499   subgraph_builder.add_inputs(inputs);
 500   subgraph_builder.add_outputs(outputs);
 501   subgraph_builder.add_operators(operators);
 502   subgraph_builder.add_name(name);
 503
 504   subgraph_vec.emplace_back(subgraph_builder.Finish());
 505 }
 506
 507 } // namespace
 508
 509 namespace tflchef
 510 {
 511
 512 /**
 513  * @brief Generate a (in-memory) TensorFlow Lite model from a given model recipe
 514  */
 515 GeneratedModel cook(const ::tflchef::ModelRecipe &model_recipe)
 516 {
 517 // Initialize Op Chef Registry
 518 #define OP_CHEF(NAME, FACTORY_CLASS) \
 519   op_chef_registry().add(#NAME, std::unique_ptr<FACTORY_CLASS>(new FACTORY_CLASS()));
 520 #include "OpChef.def"
 521 #undef OP_CHEF
 522
 523 // Initialize Data Chef Registry
 524 #define DATA_CHEF(TYPE, NAME, FACTORY_CLASS) \
 525   data_chef_registry(::tflchef::TYPE)        \
 526       .add(#NAME, std::unique_ptr<FACTORY_CLASS>(new FACTORY_CLASS()));
 527 #include <souschef/DataChef.def>
 528 #undef DATA_CHEF
 529
 530   //
 531   // Create FlatBufferBuilder
 532   //
 533   auto flatbuffer_builder =
 534       std::unique_ptr<flatbuffers::FlatBufferBuilder>(new flatbuffers::FlatBufferBuilder(1024));
 535
 536   // Operand-related
 537   std::vector<flatbuffers::Offset<::tflite::Buffer>> buffer_vec;
 538
 539   // Operation-related
 540   std::vector<flatbuffers::Offset<::tflite::OperatorCode>> code_vec;
 541
 542   // Graphs-related
 543   std::vector<flatbuffers::Offset<::tflite::SubGraph>> subgraph_vec;
 544
 545   // Create OperatorCode with Builtin Operator
 546   auto builtin_code_map = gather_builtincode_map(model_recipe);
 547   for (auto const &opcode : builtin_code_map)
 548   {
 549     tflite::OperatorCodeBuilder code_builder{*flatbuffer_builder};
 550     code_builder.add_builtin_code(opcode.first);
 551     code_builder.add_version(opcode.second);
 552     auto code = code_builder.Finish();
 553     // Update OperatorCode vector
 554     code_vec.emplace_back(code);
 555   }
 556
 557   // Create OperatorCode with Custom Operator
 558   std::set<std::string> custom_code_set = gather_customcode_set(model_recipe);
 559   if (custom_code_set.size() &&
 560       builtin_code_map.find(tflite::BuiltinOperator_CUSTOM) == builtin_code_map.end())
 561     builtin_code_map[tflite::BuiltinOperator_CUSTOM] = 1;
 562
 563   for (auto opcode : custom_code_set)
 564   {
 565     auto custom_code = flatbuffer_builder->CreateString(opcode);
 566     tflite::OperatorCodeBuilder code_builder{*flatbuffer_builder};
 567     code_builder.add_builtin_code(tflite::BuiltinOperator_CUSTOM);
 568     code_builder.add_custom_code(custom_code);
 569     auto code = code_builder.Finish();
 570     // Update OperatorCode vector
 571     code_vec.emplace_back(code);
 572   }
 573
 574   // Create an Empty Buffer
 575   //
 576   // Buffer 0 SHOULD be an empty buffer in TensorFlow Lite model file
 577   // (Please refer to the comment for Tensor.buffer field in schema)
 578   {
 579     tflite::BufferBuilder buffer_builder{*flatbuffer_builder};
 580     buffer_vec.emplace_back(buffer_builder.Finish());
 581   }
 582
 583   //
 584   // Create Main graph
 585   //
 586   CookParams cp{buffer_vec, code_vec, subgraph_vec, flatbuffer_builder, builtin_code_map, "main"};
 587
 588   cook_graph<::tflchef::ModelRecipe>(model_recipe, cp);
 589
 590   //
 591   // Create subgraphs if exist
 592   //
 593   for (int g = 0; g < model_recipe.graph_size(); ++g)
 594   {
 595     const auto &graph = model_recipe.graph(g);
 596
 597     std::ostringstream stringStream;
 598     stringStream << "sub_" << (g + 1);
 599
 600     CookParams cp{buffer_vec,         code_vec,         subgraph_vec,
 601                   flatbuffer_builder, builtin_code_map, stringStream.str()};
 602
 603     cook_graph<::tflchef::Graph>(graph, cp);
 604   }
 605
 606   // Create "Model" arguments
 607   auto buffers = flatbuffer_builder->CreateVector(buffer_vec);
 608   auto operator_codes = flatbuffer_builder->CreateVector(code_vec);
 609   auto subgraphs = flatbuffer_builder->CreateVector(subgraph_vec);
 610   auto description = flatbuffer_builder->CreateString("Generated by tflchef");
 611
 612   // Create "Model"
 613   tflite::ModelBuilder model_builder{*flatbuffer_builder};
 614
 615   model_builder.add_version(3);
 616   model_builder.add_operator_codes(operator_codes);
 617   model_builder.add_subgraphs(subgraphs);
 618   model_builder.add_description(description);
 619   model_builder.add_buffers(buffers);
 620
 621   auto model = model_builder.Finish();
 622
 623   // Finalize
 624   ::tflite::FinishModelBuffer(*flatbuffer_builder, model);
 625
 626   // Return "GenerateModel"
 627   return GeneratedModel{
 628       std::unique_ptr<GeneratedModelImpl>(new GeneratedModelImpl(std::move(flatbuffer_builder)))};
 629 }
 630
 631 } // namespace tflchef