import { __awaiter, __generator } from "tslib";
|
import * as tf from '@tensorflow/tfjs-core';
|
import { getModelUris } from './common/getModelUris';
|
import { loadWeightMap } from './dom';
|
import { env } from './env';
|
var NeuralNetwork = /** @class */ (function () {
|
function NeuralNetwork(_name) {
|
this._name = _name;
|
this._params = undefined;
|
this._paramMappings = [];
|
}
|
Object.defineProperty(NeuralNetwork.prototype, "params", {
|
get: function () { return this._params; },
|
enumerable: true,
|
configurable: true
|
});
|
Object.defineProperty(NeuralNetwork.prototype, "paramMappings", {
|
get: function () { return this._paramMappings; },
|
enumerable: true,
|
configurable: true
|
});
|
Object.defineProperty(NeuralNetwork.prototype, "isLoaded", {
|
get: function () { return !!this.params; },
|
enumerable: true,
|
configurable: true
|
});
|
NeuralNetwork.prototype.getParamFromPath = function (paramPath) {
|
var _a = this.traversePropertyPath(paramPath), obj = _a.obj, objProp = _a.objProp;
|
return obj[objProp];
|
};
|
NeuralNetwork.prototype.reassignParamFromPath = function (paramPath, tensor) {
|
var _a = this.traversePropertyPath(paramPath), obj = _a.obj, objProp = _a.objProp;
|
obj[objProp].dispose();
|
obj[objProp] = tensor;
|
};
|
NeuralNetwork.prototype.getParamList = function () {
|
var _this = this;
|
return this._paramMappings.map(function (_a) {
|
var paramPath = _a.paramPath;
|
return ({
|
path: paramPath,
|
tensor: _this.getParamFromPath(paramPath)
|
});
|
});
|
};
|
NeuralNetwork.prototype.getTrainableParams = function () {
|
return this.getParamList().filter(function (param) { return param.tensor instanceof tf.Variable; });
|
};
|
NeuralNetwork.prototype.getFrozenParams = function () {
|
return this.getParamList().filter(function (param) { return !(param.tensor instanceof tf.Variable); });
|
};
|
NeuralNetwork.prototype.variable = function () {
|
var _this = this;
|
this.getFrozenParams().forEach(function (_a) {
|
var path = _a.path, tensor = _a.tensor;
|
_this.reassignParamFromPath(path, tensor.variable());
|
});
|
};
|
NeuralNetwork.prototype.freeze = function () {
|
var _this = this;
|
this.getTrainableParams().forEach(function (_a) {
|
var path = _a.path, variable = _a.tensor;
|
var tensor = tf.tensor(variable.dataSync());
|
variable.dispose();
|
_this.reassignParamFromPath(path, tensor);
|
});
|
};
|
NeuralNetwork.prototype.dispose = function (throwOnRedispose) {
|
if (throwOnRedispose === void 0) { throwOnRedispose = true; }
|
this.getParamList().forEach(function (param) {
|
if (throwOnRedispose && param.tensor.isDisposed) {
|
throw new Error("param tensor has already been disposed for path " + param.path);
|
}
|
param.tensor.dispose();
|
});
|
this._params = undefined;
|
};
|
NeuralNetwork.prototype.serializeParams = function () {
|
return new Float32Array(this.getParamList()
|
.map(function (_a) {
|
var tensor = _a.tensor;
|
return Array.from(tensor.dataSync());
|
})
|
.reduce(function (flat, arr) { return flat.concat(arr); }));
|
};
|
NeuralNetwork.prototype.load = function (weightsOrUrl) {
|
return __awaiter(this, void 0, void 0, function () {
|
return __generator(this, function (_a) {
|
switch (_a.label) {
|
case 0:
|
if (weightsOrUrl instanceof Float32Array) {
|
this.extractWeights(weightsOrUrl);
|
return [2 /*return*/];
|
}
|
return [4 /*yield*/, this.loadFromUri(weightsOrUrl)];
|
case 1:
|
_a.sent();
|
return [2 /*return*/];
|
}
|
});
|
});
|
};
|
NeuralNetwork.prototype.loadFromUri = function (uri) {
|
return __awaiter(this, void 0, void 0, function () {
|
var weightMap;
|
return __generator(this, function (_a) {
|
switch (_a.label) {
|
case 0:
|
if (uri && typeof uri !== 'string') {
|
throw new Error(this._name + ".loadFromUri - expected model uri");
|
}
|
return [4 /*yield*/, loadWeightMap(uri, this.getDefaultModelName())];
|
case 1:
|
weightMap = _a.sent();
|
this.loadFromWeightMap(weightMap);
|
return [2 /*return*/];
|
}
|
});
|
});
|
};
|
NeuralNetwork.prototype.loadFromDisk = function (filePath) {
|
return __awaiter(this, void 0, void 0, function () {
|
var readFile, _a, manifestUri, modelBaseUri, fetchWeightsFromDisk, loadWeights, manifest, _b, _c, weightMap;
|
return __generator(this, function (_d) {
|
switch (_d.label) {
|
case 0:
|
if (filePath && typeof filePath !== 'string') {
|
throw new Error(this._name + ".loadFromDisk - expected model file path");
|
}
|
readFile = env.getEnv().readFile;
|
_a = getModelUris(filePath, this.getDefaultModelName()), manifestUri = _a.manifestUri, modelBaseUri = _a.modelBaseUri;
|
fetchWeightsFromDisk = function (filePaths) { return Promise.all(filePaths.map(function (filePath) { return readFile(filePath).then(function (buf) { return buf.buffer; }); })); };
|
loadWeights = tf.io.weightsLoaderFactory(fetchWeightsFromDisk);
|
_c = (_b = JSON).parse;
|
return [4 /*yield*/, readFile(manifestUri)];
|
case 1:
|
manifest = _c.apply(_b, [(_d.sent()).toString()]);
|
return [4 /*yield*/, loadWeights(manifest, modelBaseUri)];
|
case 2:
|
weightMap = _d.sent();
|
this.loadFromWeightMap(weightMap);
|
return [2 /*return*/];
|
}
|
});
|
});
|
};
|
NeuralNetwork.prototype.loadFromWeightMap = function (weightMap) {
|
var _a = this.extractParamsFromWeigthMap(weightMap), paramMappings = _a.paramMappings, params = _a.params;
|
this._paramMappings = paramMappings;
|
this._params = params;
|
};
|
NeuralNetwork.prototype.extractWeights = function (weights) {
|
var _a = this.extractParams(weights), paramMappings = _a.paramMappings, params = _a.params;
|
this._paramMappings = paramMappings;
|
this._params = params;
|
};
|
NeuralNetwork.prototype.traversePropertyPath = function (paramPath) {
|
if (!this.params) {
|
throw new Error("traversePropertyPath - model has no loaded params");
|
}
|
var result = paramPath.split('/').reduce(function (res, objProp) {
|
if (!res.nextObj.hasOwnProperty(objProp)) {
|
throw new Error("traversePropertyPath - object does not have property " + objProp + ", for path " + paramPath);
|
}
|
return { obj: res.nextObj, objProp: objProp, nextObj: res.nextObj[objProp] };
|
}, { nextObj: this.params });
|
var obj = result.obj, objProp = result.objProp;
|
if (!obj || !objProp || !(obj[objProp] instanceof tf.Tensor)) {
|
throw new Error("traversePropertyPath - parameter is not a tensor, for path " + paramPath);
|
}
|
return { obj: obj, objProp: objProp };
|
};
|
return NeuralNetwork;
|
}());
|
export { NeuralNetwork };
|
//# sourceMappingURL=NeuralNetwork.js.map
|
