@tensorflow/tfjs-core
Version:
Hardware-accelerated JavaScript library for machine intelligence
93 lines • 3.91 kB
JavaScript
/**
* @license
* Copyright 2018 Google LLC. All Rights Reserved.
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
* =============================================================================
*/
import { ENGINE } from '../engine';
import { dispose, tidy } from '../globals';
import { add } from '../ops/add';
import { div } from '../ops/div';
import { fill } from '../ops/fill';
import { mul } from '../ops/mul';
import { sqrt } from '../ops/sqrt';
import { square } from '../ops/square';
import { registerClass } from '../serialization';
import { Optimizer } from './optimizer';
/** @doclink Optimizer */
export class AdagradOptimizer extends Optimizer {
constructor(learningRate, initialAccumulatorValue = 0.1) {
super();
this.learningRate = learningRate;
this.initialAccumulatorValue = initialAccumulatorValue;
this.accumulatedGrads = [];
}
applyGradients(variableGradients) {
const variableNames = Array.isArray(variableGradients) ?
variableGradients.map(item => item.name) :
Object.keys(variableGradients);
variableNames.forEach((name, i) => {
const value = ENGINE.registeredVariables[name];
if (this.accumulatedGrads[i] == null) {
const trainable = false;
this.accumulatedGrads[i] = {
originalName: `${name}/accumulator`,
variable: tidy(() => fill(value.shape, this.initialAccumulatorValue)
.variable(trainable))
};
}
const gradient = Array.isArray(variableGradients) ?
variableGradients[i].tensor :
variableGradients[name];
if (gradient == null) {
return;
}
const accumulatedGrad = this.accumulatedGrads[i].variable;
tidy(() => {
const newAccumulatedGrad = add(accumulatedGrad, square(gradient));
accumulatedGrad.assign(newAccumulatedGrad);
const newValue = add(mul(div(gradient, sqrt(add(newAccumulatedGrad, ENGINE.backend.epsilon()))), -this.learningRate), value);
value.assign(newValue);
});
});
this.incrementIterations();
}
dispose() {
if (this.accumulatedGrads != null) {
dispose(this.accumulatedGrads.map(v => v.variable));
}
}
async getWeights() {
// Order matters for Python compatibility.
return [await this.saveIterations()].concat(this.accumulatedGrads.map(v => ({ name: v.originalName, tensor: v.variable })));
}
async setWeights(weightValues) {
weightValues = await this.extractIterations(weightValues);
const trainable = false;
this.accumulatedGrads = weightValues.map(v => ({ originalName: v.name, variable: v.tensor.variable(trainable) }));
}
getConfig() {
return {
'learningRate': this.learningRate,
'initialAccumulatorValue': this.initialAccumulatorValue,
};
}
/** @nocollapse */
static fromConfig(cls, config) {
return new cls(config['learningRate'], config['initialAccumulatorValue']);
}
}
/** @nocollapse */
AdagradOptimizer.className = 'Adagrad'; // Note: Name matters for Python compatibility.
registerClass(AdagradOptimizer);
//# sourceMappingURL=adagrad_optimizer.js.map