extractParams.js 4.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788
  1. import * as tf from '@tensorflow/tfjs-core';
  2. import { extractWeightsFactory } from '../common';
  3. import { isFloat } from '../utils';
  4. function extractorsFactory(extractWeights, paramMappings) {
  5. function extractFilterValues(numFilterValues, numFilters, filterSize) {
  6. var weights = extractWeights(numFilterValues);
  7. var depth = weights.length / (numFilters * filterSize * filterSize);
  8. if (isFloat(depth)) {
  9. throw new Error("depth has to be an integer: " + depth + ", weights.length: " + weights.length + ", numFilters: " + numFilters + ", filterSize: " + filterSize);
  10. }
  11. return tf.tidy(function () { return tf.transpose(tf.tensor4d(weights, [numFilters, depth, filterSize, filterSize]), [2, 3, 1, 0]); });
  12. }
  13. function extractConvParams(numFilterValues, numFilters, filterSize, mappedPrefix) {
  14. var filters = extractFilterValues(numFilterValues, numFilters, filterSize);
  15. var bias = tf.tensor1d(extractWeights(numFilters));
  16. paramMappings.push({ paramPath: mappedPrefix + "/filters" }, { paramPath: mappedPrefix + "/bias" });
  17. return { filters: filters, bias: bias };
  18. }
  19. function extractScaleLayerParams(numWeights, mappedPrefix) {
  20. var weights = tf.tensor1d(extractWeights(numWeights));
  21. var biases = tf.tensor1d(extractWeights(numWeights));
  22. paramMappings.push({ paramPath: mappedPrefix + "/weights" }, { paramPath: mappedPrefix + "/biases" });
  23. return {
  24. weights: weights,
  25. biases: biases
  26. };
  27. }
  28. function extractConvLayerParams(numFilterValues, numFilters, filterSize, mappedPrefix) {
  29. var conv = extractConvParams(numFilterValues, numFilters, filterSize, mappedPrefix + "/conv");
  30. var scale = extractScaleLayerParams(numFilters, mappedPrefix + "/scale");
  31. return { conv: conv, scale: scale };
  32. }
  33. function extractResidualLayerParams(numFilterValues, numFilters, filterSize, mappedPrefix, isDown) {
  34. if (isDown === void 0) { isDown = false; }
  35. var conv1 = extractConvLayerParams((isDown ? 0.5 : 1) * numFilterValues, numFilters, filterSize, mappedPrefix + "/conv1");
  36. var conv2 = extractConvLayerParams(numFilterValues, numFilters, filterSize, mappedPrefix + "/conv2");
  37. return { conv1: conv1, conv2: conv2 };
  38. }
  39. return {
  40. extractConvLayerParams: extractConvLayerParams,
  41. extractResidualLayerParams: extractResidualLayerParams
  42. };
  43. }
  44. export function extractParams(weights) {
  45. var _a = extractWeightsFactory(weights), extractWeights = _a.extractWeights, getRemainingWeights = _a.getRemainingWeights;
  46. var paramMappings = [];
  47. var _b = extractorsFactory(extractWeights, paramMappings), extractConvLayerParams = _b.extractConvLayerParams, extractResidualLayerParams = _b.extractResidualLayerParams;
  48. var conv32_down = extractConvLayerParams(4704, 32, 7, 'conv32_down');
  49. var conv32_1 = extractResidualLayerParams(9216, 32, 3, 'conv32_1');
  50. var conv32_2 = extractResidualLayerParams(9216, 32, 3, 'conv32_2');
  51. var conv32_3 = extractResidualLayerParams(9216, 32, 3, 'conv32_3');
  52. var conv64_down = extractResidualLayerParams(36864, 64, 3, 'conv64_down', true);
  53. var conv64_1 = extractResidualLayerParams(36864, 64, 3, 'conv64_1');
  54. var conv64_2 = extractResidualLayerParams(36864, 64, 3, 'conv64_2');
  55. var conv64_3 = extractResidualLayerParams(36864, 64, 3, 'conv64_3');
  56. var conv128_down = extractResidualLayerParams(147456, 128, 3, 'conv128_down', true);
  57. var conv128_1 = extractResidualLayerParams(147456, 128, 3, 'conv128_1');
  58. var conv128_2 = extractResidualLayerParams(147456, 128, 3, 'conv128_2');
  59. var conv256_down = extractResidualLayerParams(589824, 256, 3, 'conv256_down', true);
  60. var conv256_1 = extractResidualLayerParams(589824, 256, 3, 'conv256_1');
  61. var conv256_2 = extractResidualLayerParams(589824, 256, 3, 'conv256_2');
  62. var conv256_down_out = extractResidualLayerParams(589824, 256, 3, 'conv256_down_out');
  63. var fc = tf.tidy(function () { return tf.transpose(tf.tensor2d(extractWeights(256 * 128), [128, 256]), [1, 0]); });
  64. paramMappings.push({ paramPath: "fc" });
  65. if (getRemainingWeights().length !== 0) {
  66. throw new Error("weights remaing after extract: " + getRemainingWeights().length);
  67. }
  68. var params = {
  69. conv32_down: conv32_down,
  70. conv32_1: conv32_1,
  71. conv32_2: conv32_2,
  72. conv32_3: conv32_3,
  73. conv64_down: conv64_down,
  74. conv64_1: conv64_1,
  75. conv64_2: conv64_2,
  76. conv64_3: conv64_3,
  77. conv128_down: conv128_down,
  78. conv128_1: conv128_1,
  79. conv128_2: conv128_2,
  80. conv256_down: conv256_down,
  81. conv256_1: conv256_1,
  82. conv256_2: conv256_2,
  83. conv256_down_out: conv256_down_out,
  84. fc: fc
  85. };
  86. return { params: params, paramMappings: paramMappings };
  87. }
  88. //# sourceMappingURL=extractParams.js.map