gx

chenyc

2025-02-12 ea42ff3ebee1eeb3fb29423aa848a249441db81c

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
import * as tf from '../../dist/tfjs.esm';
 
import { disposeUnusedWeightTensors, extractWeightEntryFactory, FCParams, ParamMapping } from '../common/index';
import { NetParams } from './types';
 
export function extractParamsFromWeightMap(
  weightMap: tf.NamedTensorMap,
): { params: NetParams, paramMappings: ParamMapping[] } {
  const paramMappings: ParamMapping[] = [];
 
  const extractWeightEntry = extractWeightEntryFactory(weightMap, paramMappings);
 
  function extractFcParams(prefix: string): FCParams {
    const weights = extractWeightEntry(`${prefix}/weights`, 2);
    const bias = extractWeightEntry(`${prefix}/bias`, 1);
    return { weights, bias };
  }
 
  const params = {
    fc: extractFcParams('fc'),
  };
 
  disposeUnusedWeightTensors(weightMap, paramMappings);
 
  return { params, paramMappings };
}