extractParamsFromWeigthMap.js 6.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103
  1. "use strict";
  2. Object.defineProperty(exports, "__esModule", { value: true });
  3. var common_1 = require("../common");
  4. var utils_1 = require("../utils");
  5. function extractorsFactory(weightMap, paramMappings) {
  6. var extractWeightEntry = common_1.extractWeightEntryFactory(weightMap, paramMappings);
  7. function extractPointwiseConvParams(prefix, idx, mappedPrefix) {
  8. var filters = extractWeightEntry(prefix + "/Conv2d_" + idx + "_pointwise/weights", 4, mappedPrefix + "/filters");
  9. var batch_norm_offset = extractWeightEntry(prefix + "/Conv2d_" + idx + "_pointwise/convolution_bn_offset", 1, mappedPrefix + "/batch_norm_offset");
  10. return { filters: filters, batch_norm_offset: batch_norm_offset };
  11. }
  12. function extractConvPairParams(idx) {
  13. var mappedPrefix = "mobilenetv1/conv_" + idx;
  14. var prefixDepthwiseConv = "MobilenetV1/Conv2d_" + idx + "_depthwise";
  15. var mappedPrefixDepthwiseConv = mappedPrefix + "/depthwise_conv";
  16. var mappedPrefixPointwiseConv = mappedPrefix + "/pointwise_conv";
  17. var filters = extractWeightEntry(prefixDepthwiseConv + "/depthwise_weights", 4, mappedPrefixDepthwiseConv + "/filters");
  18. var batch_norm_scale = extractWeightEntry(prefixDepthwiseConv + "/BatchNorm/gamma", 1, mappedPrefixDepthwiseConv + "/batch_norm_scale");
  19. var batch_norm_offset = extractWeightEntry(prefixDepthwiseConv + "/BatchNorm/beta", 1, mappedPrefixDepthwiseConv + "/batch_norm_offset");
  20. var batch_norm_mean = extractWeightEntry(prefixDepthwiseConv + "/BatchNorm/moving_mean", 1, mappedPrefixDepthwiseConv + "/batch_norm_mean");
  21. var batch_norm_variance = extractWeightEntry(prefixDepthwiseConv + "/BatchNorm/moving_variance", 1, mappedPrefixDepthwiseConv + "/batch_norm_variance");
  22. return {
  23. depthwise_conv: {
  24. filters: filters,
  25. batch_norm_scale: batch_norm_scale,
  26. batch_norm_offset: batch_norm_offset,
  27. batch_norm_mean: batch_norm_mean,
  28. batch_norm_variance: batch_norm_variance
  29. },
  30. pointwise_conv: extractPointwiseConvParams('MobilenetV1', idx, mappedPrefixPointwiseConv)
  31. };
  32. }
  33. function extractMobilenetV1Params() {
  34. return {
  35. conv_0: extractPointwiseConvParams('MobilenetV1', 0, 'mobilenetv1/conv_0'),
  36. conv_1: extractConvPairParams(1),
  37. conv_2: extractConvPairParams(2),
  38. conv_3: extractConvPairParams(3),
  39. conv_4: extractConvPairParams(4),
  40. conv_5: extractConvPairParams(5),
  41. conv_6: extractConvPairParams(6),
  42. conv_7: extractConvPairParams(7),
  43. conv_8: extractConvPairParams(8),
  44. conv_9: extractConvPairParams(9),
  45. conv_10: extractConvPairParams(10),
  46. conv_11: extractConvPairParams(11),
  47. conv_12: extractConvPairParams(12),
  48. conv_13: extractConvPairParams(13)
  49. };
  50. }
  51. function extractConvParams(prefix, mappedPrefix) {
  52. var filters = extractWeightEntry(prefix + "/weights", 4, mappedPrefix + "/filters");
  53. var bias = extractWeightEntry(prefix + "/biases", 1, mappedPrefix + "/bias");
  54. return { filters: filters, bias: bias };
  55. }
  56. function extractBoxPredictorParams(idx) {
  57. var box_encoding_predictor = extractConvParams("Prediction/BoxPredictor_" + idx + "/BoxEncodingPredictor", "prediction_layer/box_predictor_" + idx + "/box_encoding_predictor");
  58. var class_predictor = extractConvParams("Prediction/BoxPredictor_" + idx + "/ClassPredictor", "prediction_layer/box_predictor_" + idx + "/class_predictor");
  59. return { box_encoding_predictor: box_encoding_predictor, class_predictor: class_predictor };
  60. }
  61. function extractPredictionLayerParams() {
  62. return {
  63. conv_0: extractPointwiseConvParams('Prediction', 0, 'prediction_layer/conv_0'),
  64. conv_1: extractPointwiseConvParams('Prediction', 1, 'prediction_layer/conv_1'),
  65. conv_2: extractPointwiseConvParams('Prediction', 2, 'prediction_layer/conv_2'),
  66. conv_3: extractPointwiseConvParams('Prediction', 3, 'prediction_layer/conv_3'),
  67. conv_4: extractPointwiseConvParams('Prediction', 4, 'prediction_layer/conv_4'),
  68. conv_5: extractPointwiseConvParams('Prediction', 5, 'prediction_layer/conv_5'),
  69. conv_6: extractPointwiseConvParams('Prediction', 6, 'prediction_layer/conv_6'),
  70. conv_7: extractPointwiseConvParams('Prediction', 7, 'prediction_layer/conv_7'),
  71. box_predictor_0: extractBoxPredictorParams(0),
  72. box_predictor_1: extractBoxPredictorParams(1),
  73. box_predictor_2: extractBoxPredictorParams(2),
  74. box_predictor_3: extractBoxPredictorParams(3),
  75. box_predictor_4: extractBoxPredictorParams(4),
  76. box_predictor_5: extractBoxPredictorParams(5)
  77. };
  78. }
  79. return {
  80. extractMobilenetV1Params: extractMobilenetV1Params,
  81. extractPredictionLayerParams: extractPredictionLayerParams
  82. };
  83. }
  84. function extractParamsFromWeigthMap(weightMap) {
  85. var paramMappings = [];
  86. var _a = extractorsFactory(weightMap, paramMappings), extractMobilenetV1Params = _a.extractMobilenetV1Params, extractPredictionLayerParams = _a.extractPredictionLayerParams;
  87. var extra_dim = weightMap['Output/extra_dim'];
  88. paramMappings.push({ originalPath: 'Output/extra_dim', paramPath: 'output_layer/extra_dim' });
  89. if (!utils_1.isTensor3D(extra_dim)) {
  90. throw new Error("expected weightMap['Output/extra_dim'] to be a Tensor3D, instead have " + extra_dim);
  91. }
  92. var params = {
  93. mobilenetv1: extractMobilenetV1Params(),
  94. prediction_layer: extractPredictionLayerParams(),
  95. output_layer: {
  96. extra_dim: extra_dim
  97. }
  98. };
  99. common_1.disposeUnusedWeightTensors(weightMap, paramMappings);
  100. return { params: params, paramMappings: paramMappings };
  101. }
  102. exports.extractParamsFromWeigthMap = extractParamsFromWeigthMap;
  103. //# sourceMappingURL=extractParamsFromWeigthMap.js.map