extractParamsFromWeigthMap.js 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354
  1. "use strict";
  2. Object.defineProperty(exports, "__esModule", { value: true });
  3. var common_1 = require("../common");
  4. var loadConvParamsFactory_1 = require("../common/loadConvParamsFactory");
  5. var utils_1 = require("../utils");
  6. function loadParamsFactory(weightMap, paramMappings) {
  7. var extractWeightEntry = common_1.extractWeightEntryFactory(weightMap, paramMappings);
  8. var extractConvParams = loadConvParamsFactory_1.loadConvParamsFactory(extractWeightEntry);
  9. var extractSeparableConvParams = common_1.loadSeparableConvParamsFactory(extractWeightEntry);
  10. function extractReductionBlockParams(mappedPrefix) {
  11. var separable_conv0 = extractSeparableConvParams(mappedPrefix + "/separable_conv0");
  12. var separable_conv1 = extractSeparableConvParams(mappedPrefix + "/separable_conv1");
  13. var expansion_conv = extractConvParams(mappedPrefix + "/expansion_conv");
  14. return { separable_conv0: separable_conv0, separable_conv1: separable_conv1, expansion_conv: expansion_conv };
  15. }
  16. function extractMainBlockParams(mappedPrefix) {
  17. var separable_conv0 = extractSeparableConvParams(mappedPrefix + "/separable_conv0");
  18. var separable_conv1 = extractSeparableConvParams(mappedPrefix + "/separable_conv1");
  19. var separable_conv2 = extractSeparableConvParams(mappedPrefix + "/separable_conv2");
  20. return { separable_conv0: separable_conv0, separable_conv1: separable_conv1, separable_conv2: separable_conv2 };
  21. }
  22. return {
  23. extractConvParams: extractConvParams,
  24. extractSeparableConvParams: extractSeparableConvParams,
  25. extractReductionBlockParams: extractReductionBlockParams,
  26. extractMainBlockParams: extractMainBlockParams
  27. };
  28. }
  29. function extractParamsFromWeigthMap(weightMap, numMainBlocks) {
  30. var paramMappings = [];
  31. var _a = loadParamsFactory(weightMap, paramMappings), extractConvParams = _a.extractConvParams, extractSeparableConvParams = _a.extractSeparableConvParams, extractReductionBlockParams = _a.extractReductionBlockParams, extractMainBlockParams = _a.extractMainBlockParams;
  32. var entry_flow_conv_in = extractConvParams('entry_flow/conv_in');
  33. var entry_flow_reduction_block_0 = extractReductionBlockParams('entry_flow/reduction_block_0');
  34. var entry_flow_reduction_block_1 = extractReductionBlockParams('entry_flow/reduction_block_1');
  35. var entry_flow = {
  36. conv_in: entry_flow_conv_in,
  37. reduction_block_0: entry_flow_reduction_block_0,
  38. reduction_block_1: entry_flow_reduction_block_1
  39. };
  40. var middle_flow = {};
  41. utils_1.range(numMainBlocks, 0, 1).forEach(function (idx) {
  42. middle_flow["main_block_" + idx] = extractMainBlockParams("middle_flow/main_block_" + idx);
  43. });
  44. var exit_flow_reduction_block = extractReductionBlockParams('exit_flow/reduction_block');
  45. var exit_flow_separable_conv = extractSeparableConvParams('exit_flow/separable_conv');
  46. var exit_flow = {
  47. reduction_block: exit_flow_reduction_block,
  48. separable_conv: exit_flow_separable_conv
  49. };
  50. common_1.disposeUnusedWeightTensors(weightMap, paramMappings);
  51. return { params: { entry_flow: entry_flow, middle_flow: middle_flow, exit_flow: exit_flow }, paramMappings: paramMappings };
  52. }
  53. exports.extractParamsFromWeigthMap = extractParamsFromWeigthMap;
  54. //# sourceMappingURL=extractParamsFromWeigthMap.js.map