extractParamsFromWeigthMap.js 3.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970
  1. "use strict";
  2. Object.defineProperty(exports, "__esModule", { value: true });
  3. var tslib_1 = require("tslib");
  4. var common_1 = require("../common");
  5. function extractorsFactory(weightMap, paramMappings) {
  6. var extractWeightEntry = common_1.extractWeightEntryFactory(weightMap, paramMappings);
  7. function extractConvParams(prefix) {
  8. var filters = extractWeightEntry(prefix + "/weights", 4, prefix + "/filters");
  9. var bias = extractWeightEntry(prefix + "/bias", 1);
  10. return { filters: filters, bias: bias };
  11. }
  12. function extractFCParams(prefix) {
  13. var weights = extractWeightEntry(prefix + "/weights", 2);
  14. var bias = extractWeightEntry(prefix + "/bias", 1);
  15. return { weights: weights, bias: bias };
  16. }
  17. function extractPReluParams(paramPath) {
  18. return extractWeightEntry(paramPath, 1);
  19. }
  20. function extractSharedParams(prefix) {
  21. var conv1 = extractConvParams(prefix + "/conv1");
  22. var prelu1_alpha = extractPReluParams(prefix + "/prelu1_alpha");
  23. var conv2 = extractConvParams(prefix + "/conv2");
  24. var prelu2_alpha = extractPReluParams(prefix + "/prelu2_alpha");
  25. var conv3 = extractConvParams(prefix + "/conv3");
  26. var prelu3_alpha = extractPReluParams(prefix + "/prelu3_alpha");
  27. return { conv1: conv1, prelu1_alpha: prelu1_alpha, conv2: conv2, prelu2_alpha: prelu2_alpha, conv3: conv3, prelu3_alpha: prelu3_alpha };
  28. }
  29. function extractPNetParams() {
  30. var sharedParams = extractSharedParams('pnet');
  31. var conv4_1 = extractConvParams('pnet/conv4_1');
  32. var conv4_2 = extractConvParams('pnet/conv4_2');
  33. return tslib_1.__assign(tslib_1.__assign({}, sharedParams), { conv4_1: conv4_1, conv4_2: conv4_2 });
  34. }
  35. function extractRNetParams() {
  36. var sharedParams = extractSharedParams('rnet');
  37. var fc1 = extractFCParams('rnet/fc1');
  38. var prelu4_alpha = extractPReluParams('rnet/prelu4_alpha');
  39. var fc2_1 = extractFCParams('rnet/fc2_1');
  40. var fc2_2 = extractFCParams('rnet/fc2_2');
  41. return tslib_1.__assign(tslib_1.__assign({}, sharedParams), { fc1: fc1, prelu4_alpha: prelu4_alpha, fc2_1: fc2_1, fc2_2: fc2_2 });
  42. }
  43. function extractONetParams() {
  44. var sharedParams = extractSharedParams('onet');
  45. var conv4 = extractConvParams('onet/conv4');
  46. var prelu4_alpha = extractPReluParams('onet/prelu4_alpha');
  47. var fc1 = extractFCParams('onet/fc1');
  48. var prelu5_alpha = extractPReluParams('onet/prelu5_alpha');
  49. var fc2_1 = extractFCParams('onet/fc2_1');
  50. var fc2_2 = extractFCParams('onet/fc2_2');
  51. var fc2_3 = extractFCParams('onet/fc2_3');
  52. return tslib_1.__assign(tslib_1.__assign({}, sharedParams), { conv4: conv4, prelu4_alpha: prelu4_alpha, fc1: fc1, prelu5_alpha: prelu5_alpha, fc2_1: fc2_1, fc2_2: fc2_2, fc2_3: fc2_3 });
  53. }
  54. return {
  55. extractPNetParams: extractPNetParams,
  56. extractRNetParams: extractRNetParams,
  57. extractONetParams: extractONetParams
  58. };
  59. }
  60. function extractParamsFromWeigthMap(weightMap) {
  61. var paramMappings = [];
  62. var _a = extractorsFactory(weightMap, paramMappings), extractPNetParams = _a.extractPNetParams, extractRNetParams = _a.extractRNetParams, extractONetParams = _a.extractONetParams;
  63. var pnet = extractPNetParams();
  64. var rnet = extractRNetParams();
  65. var onet = extractONetParams();
  66. common_1.disposeUnusedWeightTensors(weightMap, paramMappings);
  67. return { params: { pnet: pnet, rnet: rnet, onet: onet }, paramMappings: paramMappings };
  68. }
  69. exports.extractParamsFromWeigthMap = extractParamsFromWeigthMap;
  70. //# sourceMappingURL=extractParamsFromWeigthMap.js.map