@tensorflow/tfjs-core
Version:
Hardware-accelerated JavaScript library for machine intelligence
108 lines • 4.19 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 { mul } from '../ops/mul';
import { scalar } from '../ops/scalar';
import { zerosLike } from '../ops/zeros_like';
import { registerClass } from '../serialization';
import { SGDOptimizer } from './sgd_optimizer';
/** @doclink Optimizer */
export class MomentumOptimizer extends SGDOptimizer {
constructor(learningRate, momentum, useNesterov = false) {
super(learningRate);
this.learningRate = learningRate;
this.momentum = momentum;
this.useNesterov = useNesterov;
this.accumulations = [];
this.m = scalar(this.momentum);
}
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.accumulations[i] == null) {
const trainable = false;
this.accumulations[i] = {
originalName: `${name}/momentum`,
variable: tidy(() => zerosLike(value).variable(trainable))
};
}
const accumulation = this.accumulations[i].variable;
const gradient = Array.isArray(variableGradients) ?
variableGradients[i].tensor :
variableGradients[name];
if (gradient == null) {
return;
}
tidy(() => {
let newValue;
const newAccumulation = add(mul(this.m, accumulation), gradient);
if (this.useNesterov) {
newValue = add(mul(this.c, add(gradient, mul(newAccumulation, this.m))), value);
}
else {
newValue = add(mul(this.c, newAccumulation), value);
}
accumulation.assign(newAccumulation);
value.assign(newValue);
});
});
this.incrementIterations();
}
dispose() {
this.m.dispose();
if (this.accumulations != null) {
dispose(this.accumulations.map(v => v.variable));
}
}
/**
* Sets the momentum of the optimizer.
*
* @param momentum
*/
setMomentum(momentum) {
this.momentum = momentum;
}
async getWeights() {
// Order matters for Python compatibility.
return [await this.saveIterations()].concat(this.accumulations.map(v => ({ name: v.originalName, tensor: v.variable })));
}
async setWeights(weightValues) {
weightValues = await this.extractIterations(weightValues);
const trainable = false;
this.accumulations = weightValues.map(v => ({ originalName: v.name, variable: v.tensor.variable(trainable) }));
}
getConfig() {
return {
'learningRate': this.learningRate,
'momentum': this.momentum,
'useNesterov': this.useNesterov
};
}
/** @nocollapse */
static fromConfig(cls, config) {
return new cls(config['learningRate'], config['momentum'], config['useNesterov']);
}
}
/** @nocollapse */
MomentumOptimizer.className = 'Momentum'; // Name matters for Python compatibility.
registerClass(MomentumOptimizer);
//# sourceMappingURL=momentum_optimizer.js.map