tflite_metadata.js 2.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  1. const path = require('path');
  2. const flatc = require('./flatc');
  3. const fs = require('fs');
  4. const schema = path.join(__dirname, '..', 'third_party', 'source', 'tensorflow', 'tensorflow', 'lite', 'schema', 'schema.fbs');
  5. const file = path.join(__dirname, '..', 'source', 'tflite-metadata.json');
  6. const input = fs.readFileSync(file, 'utf-8');
  7. const json = JSON.parse(input);
  8. const operators = new Map();
  9. const attributes = new Map();
  10. for (const operator of json) {
  11. if (operators.has(operator.name)) {
  12. throw new Error("Duplicate operator '" + operator.name + "'.");
  13. }
  14. operators.set(operator.name, operator);
  15. if (operator && operator.attributes) {
  16. for (const attribute of operator.attributes) {
  17. const name = operator.name + ':' + attribute.name;
  18. attributes.set(name, attribute);
  19. }
  20. }
  21. }
  22. const root = new flatc.Root('tflite', [], [ schema ]);
  23. const namespace = root.find('tflite', flatc.Namespace);
  24. const builtOperator = namespace.find('tflite.BuiltinOperator', flatc.Type);
  25. const upperCase = new Set([ '2D', 'LSH', 'SVDF', 'RNN', 'L2', 'LSTM' ]);
  26. for (const op of builtOperator.values.keys()) {
  27. let op_key = op === 'BATCH_MATMUL' ? 'BATCH_MAT_MUL' : op;
  28. op_key = op_key.split('_').map((s) => (s.length < 1 || upperCase.has(s)) ? s : s[0] + s.substring(1).toLowerCase()).join('');
  29. const table = namespace.find('tflite.' + op_key + 'Options', flatc.Type);
  30. if (table && table.fields.size > 0) {
  31. if (!operators.has(op_key)) {
  32. const operator = { name: op_key };
  33. operators.set(op_key, operator);
  34. json.push(operator);
  35. }
  36. const operator = operators.get(op_key);
  37. operator.attributes = operator.attributes || [];
  38. for (const field of table.fields.values()) {
  39. const attr_key = op_key + ':' + field.name;
  40. if (!attributes.has(attr_key)) {
  41. const attribute = { name: field.name };
  42. attributes.set(attr_key, attribute);
  43. operator.attributes.push(attribute);
  44. }
  45. const attribute = attributes.get(attr_key);
  46. const type = field.type;
  47. let defaultValue = field.defaultValue;
  48. if (type instanceof flatc.Enum) {
  49. if (!type.keys.has(defaultValue)) {
  50. throw new Error("Invalid '" + type.name + "' default value '" + defaultValue + "'.");
  51. }
  52. defaultValue = type.keys.get(defaultValue);
  53. }
  54. attribute.type = type.name === 'bool' ? 'boolean' : type.name + (field.repeated ? '[]' : '');
  55. attribute.default = defaultValue;
  56. }
  57. }
  58. }
  59. json.sort((a, b) => a.name.localeCompare(b.name));
  60. let output = JSON.stringify(json, null, 2);
  61. output = output.replace(/\s {8}/g, ' ');
  62. output = output.replace(/,\s {8}/g, ', ');
  63. output = output.replace(/\s {6}}/g, ' }');
  64. fs.writeFileSync(file, output, 'utf-8');