@tensorflow/tfjs-core
Version:
Hardware-accelerated JavaScript library for machine intelligence
344 lines • 12.5 kB
JavaScript
/**
* @license
* Copyright 2017 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 { add } from './engine';
import * as tf from './index';
import { ALL_ENVS, describeWithFlags } from './jasmine_util';
import { zerosLike } from './ops/ops';
import { backpropagateGradients, getFilteredNodesXToY } from './tape';
import { expectArraysClose } from './test_util';
describeWithFlags('getFilteredNodesXToY', ALL_ENVS, () => {
it('no paths from x to y', () => {
const x = tf.scalar(1);
const intermediate1 = tf.scalar(0);
const intermediate2 = tf.scalar(0);
const y = tf.scalar(2);
const tape = [
{
id: 0,
kernelName: 'node0',
inputs: { x },
outputs: [intermediate1],
gradient: null
},
{
id: 1,
kernelName: 'node1',
inputs: { intermediate2 },
outputs: [y],
gradient: null
}
];
const filteredTapeNodes = getFilteredNodesXToY(tape, [x], y);
expect(filteredTapeNodes.length).toBe(0);
expect(filteredTapeNodes).toEqual([]);
});
it('one operation x => y', () => {
const x = tf.scalar(1);
const y = tf.scalar(2);
const tape = [
{ id: 0, kernelName: 'node0', inputs: { x }, outputs: [y], gradient: null }
];
const filteredTapeNodes = getFilteredNodesXToY(tape, [x], y);
expect(filteredTapeNodes.length).toBe(1);
expect(filteredTapeNodes).toEqual(tape);
});
it('1 operation [x0, x1] => y, all input paths', () => {
const x0 = tf.scalar(0);
const x1 = tf.scalar(1);
const y = tf.scalar(2);
const tape = [{
id: 0,
kernelName: 'node0',
inputs: { x0, x1 },
outputs: [y],
gradient: null
}];
const filteredTapeNodes = getFilteredNodesXToY(tape, [x0, x1], y);
expect(filteredTapeNodes.length).toBe(1);
expect(filteredTapeNodes).toEqual(tape);
});
it('one operation [x0, x1] => y, one input paths', () => {
const x0 = tf.scalar(0);
const x1 = tf.scalar(1);
const y = tf.scalar(2);
const tape = [{
id: 0,
kernelName: 'node0',
inputs: { x0, x1 },
outputs: [y],
gradient: null
}];
const filteredTapeNodes = getFilteredNodesXToY(tape, [x0], y);
expect(filteredTapeNodes.length).toBe(1);
// x1 input should be pruned, we don't ask for the gradient of x1.
expect(filteredTapeNodes[0]).toEqual({
id: 0,
kernelName: 'node0',
inputs: { x0 },
outputs: [y],
gradient: null
});
});
it('two operations x => intermediate => y', () => {
const x = tf.scalar(1);
const intermediate = tf.scalar(0);
const y = tf.scalar(2);
const tape = [
{
id: 0,
kernelName: 'node0',
inputs: { x },
outputs: [intermediate],
gradient: null
},
{
id: 1,
kernelName: 'node1',
inputs: { intermediate },
outputs: [y],
gradient: null
}
];
const filteredTapeNodes = getFilteredNodesXToY(tape, [x], y);
expect(filteredTapeNodes.length).toBe(2);
expect(filteredTapeNodes).toEqual(tape);
});
it('two operations [x0, x1], [x2] => ' +
'intermediate => y', () => {
const x0 = tf.scalar(1);
const x1 = tf.scalar(2);
const x2 = tf.scalar(3);
const intermediate = tf.scalar(4);
const y = tf.scalar(2);
const tape = [
{
id: 0,
kernelName: 'node0',
inputs: { x0, x1 },
outputs: [intermediate],
gradient: null
},
{
id: 1,
kernelName: 'node1',
inputs: { x2, intermediate },
outputs: [y],
gradient: null
}
];
const filteredTapeNodes = getFilteredNodesXToY(tape, [x0, x1, x2], y);
expect(filteredTapeNodes.length).toBe(2);
expect(filteredTapeNodes).toEqual(tape);
});
it('x => y and x => orphan', () => {
const x = tf.scalar(1);
const orphan = tf.scalar(0);
const y = tf.scalar(2);
const tape = [
{
id: 0,
kernelName: 'node0',
inputs: { x },
outputs: [orphan],
gradient: null
},
{ id: 1, kernelName: 'node1', inputs: { x }, outputs: [y], gradient: null }
];
const filteredTapeNodes = getFilteredNodesXToY(tape, [x], y);
expect(filteredTapeNodes.length).toBe(1);
// The orphan should be removed.
expect(filteredTapeNodes[0]).toEqual(tape[1]);
});
it('x => y and orphan => y', () => {
const x = tf.scalar(1);
const orphan = tf.scalar(0);
const y = tf.scalar(2);
const tape = [{
id: 0,
kernelName: 'node0',
inputs: { x, orphan },
outputs: [y],
gradient: null
}];
const filteredTapeNodes = getFilteredNodesXToY(tape, [x], y);
expect(filteredTapeNodes.length).toBe(1);
// The orphan should be pruned from the node's input.
expect(filteredTapeNodes[0]).toEqual({
id: 0,
kernelName: 'node0',
inputs: { x },
outputs: [y],
gradient: null
});
});
it('1 op with 3 outputs x => y1, y2, y3', () => {
const x = tf.scalar(1);
const y1 = tf.scalar(2);
const y2 = tf.scalar(2);
const y3 = tf.scalar(2);
const tape = [{
id: 0,
kernelName: 'node0',
inputs: { x },
outputs: [y1, y2, y3],
gradient: null
}];
const filteredNodes1 = getFilteredNodesXToY(tape, [x], y1);
expect(filteredNodes1.length).toBe(1);
expect(filteredNodes1).toEqual(tape);
const filteredNodes2 = getFilteredNodesXToY(tape, [x], y2);
expect(filteredNodes2.length).toBe(1);
expect(filteredNodes2).toEqual(tape);
const filteredNodes3 = getFilteredNodesXToY(tape, [x], y3);
expect(filteredNodes3.length).toBe(1);
expect(filteredNodes3).toEqual(tape);
});
});
describeWithFlags('backpropagateGradients', ALL_ENVS, () => {
it('Throws if gradient is not defined', () => {
const x = tf.scalar(0);
const y = tf.scalar(1);
const dy = tf.scalar(1);
const accumulatedGradientsMap = {};
accumulatedGradientsMap[y.id] = dy;
const tape = [
{ id: 0, kernelName: 'node0', inputs: { x }, outputs: [y], gradient: null }
];
expect(() => backpropagateGradients(accumulatedGradientsMap, tape, f => tf.tidy(f), add))
.toThrowError();
});
it('basic backprop with 1 node', async () => {
const x = tf.scalar(0);
const y = tf.scalar(1);
const dy = tf.scalar(1);
const accumulatedGradientsMap = {};
accumulatedGradientsMap[y.id] = dy;
const tape = [{
id: 0,
kernelName: 'node0',
inputs: { x },
outputs: [y],
gradient: (dys) => {
return { x: () => dys[0].add(tf.scalar(1)) };
}
}];
backpropagateGradients(accumulatedGradientsMap, tape, f => tf.tidy(f), add);
expectArraysClose(await accumulatedGradientsMap[x.id].data(), [2]);
});
it('basic backprop with 2 nodes', async () => {
const x = tf.scalar(0);
const intermediate = tf.scalar(1);
const y = tf.scalar(2);
const dy = tf.scalar(1);
const accumulatedGradientsMap = {};
accumulatedGradientsMap[y.id] = dy;
const tape = [
{
id: 0,
kernelName: 'node0',
inputs: { x },
outputs: [intermediate],
gradient: (dys) => {
return { x: () => dys[0].add(tf.scalar(1)) };
}
},
{
id: 1,
kernelName: 'node1',
inputs: { intermediate },
outputs: [y],
gradient: (dys) => {
return { intermediate: () => dys[0].add(tf.scalar(1)) };
}
}
];
backpropagateGradients(accumulatedGradientsMap, tape, f => tf.tidy(f), add);
// dx = dy + 1 + 1
expectArraysClose(await accumulatedGradientsMap[x.id].data(), [3]);
});
it('basic backprop with a split node accumulates gradients', async () => {
const x = tf.scalar(0);
const intermediate1 = tf.scalar(1);
const intermediate2 = tf.scalar(2);
const y = tf.scalar(3);
const dy = tf.scalar(1);
const accumulatedGradientsMap = {};
accumulatedGradientsMap[y.id] = dy;
const tape = [
{
id: 0,
kernelName: 'node0',
inputs: { x },
outputs: [intermediate1],
gradient: (dys) => {
return { x: () => dys[0].add(tf.scalar(1)) };
}
},
{
id: 1,
kernelName: 'node1',
inputs: { x },
outputs: [intermediate2],
gradient: (dys) => {
return { x: () => dys[0].add(tf.scalar(1)) };
}
},
{
id: 2,
kernelName: 'node2',
inputs: { intermediate1, intermediate2 },
outputs: [y],
gradient: (dys) => {
return {
intermediate1: () => dys[0].add(tf.scalar(1)),
intermediate2: () => dys[0].add(tf.scalar(1))
};
}
}
];
backpropagateGradients(accumulatedGradientsMap, tape, f => tf.tidy(f), add);
// dx = dy + 1 + 1 + 1 + 1 + 1
expectArraysClose(await accumulatedGradientsMap[x.id].data(), [(await dy.data())[0] + 5]);
});
it('backprop over 1 node with 3 outputs, w.r.t to the 2nd output', async () => {
const x = tf.tensor1d([1, 1, 1]);
const y1 = tf.scalar(1);
const y2 = tf.scalar(1);
const y3 = tf.scalar(1);
const accumulatedGradientsMap = {};
// Backproping through the 2nd output.
const dy2 = tf.scalar(5);
accumulatedGradientsMap[y2.id] = dy2;
let dys;
const tape = [{
id: 0,
kernelName: 'node0',
inputs: { x },
outputs: [y1, y2, y3],
gradient: (dys_) => {
dys = dys_.map(dy => dy || zerosLike(y1));
return { x: () => tf.stack(dys) };
}
}];
backpropagateGradients(accumulatedGradientsMap, tape, f => tf.tidy(f), add);
expectArraysClose(await accumulatedGradientsMap[x.id].data(), [0, 5, 0]);
expectArraysClose(await dys[0].data(), [0]);
expectArraysClose(await dys[1].data(), [5]);
expectArraysClose(await dys[2].data(), [0]);
});
});
//# sourceMappingURL=tape_test.js.map