@tensorflow/tfjs-core
Version:
Hardware-accelerated JavaScript library for machine intelligence
270 lines • 13.2 kB
JavaScript
"use strict";
/**
* @license
* Copyright 2019 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.
* =============================================================================
*/
Object.defineProperty(exports, "__esModule", { value: true });
var engine_1 = require("../engine");
var conv_1 = require("../ops/conv");
var conv_util = require("../ops/conv_util");
var operation_1 = require("../ops/operation");
var tensor_util_1 = require("../tensor_util");
var tensor_util_env_1 = require("../tensor_util_env");
var util = require("../util");
var broadcast_util = require("./broadcast_util");
/**
* Computes the dot product of two matrices with optional activation and bias.
*
* ```js
* const a = tf.tensor2d([-1, -2], [1, 2]);
* const b = tf.tensor2d([1, 2, 3, 4], [2, 2]);
* const bias = tf.tensor2d([1, 2], [1, 2]);
*
* tf.fused.matMul(a, b, false, false, bias, 'relu').print();
* ```
*
* @param a First matrix in dot product operation.
* @param b Second matrix in dot product operation.
* @param transposeA If true, `a` is transposed before multiplication.
* @param transposeB If true, `b` is transposed before multiplication.
* @param bias Matrix to be added to the result.
* @param activation Name of activation kernel (defaults to `linear`).
*/
/** @doc {heading: 'Operations', subheading: 'Matrices', namespace: 'fused'} */
function matMul_(a, b, transposeA, transposeB, bias, activation) {
if (transposeA === void 0) { transposeA = false; }
if (transposeB === void 0) { transposeB = false; }
if (activation === void 0) { activation = 'linear'; }
var _a;
var $a = tensor_util_env_1.convertToTensor(a, 'a', 'fused matMul');
var $b = tensor_util_env_1.convertToTensor(b, 'b', 'fused matMul');
_a = tensor_util_1.makeTypesMatch($a, $b), $a = _a[0], $b = _a[1];
var innerShapeA = transposeA ? $a.shape[$a.rank - 2] : $a.shape[$a.rank - 1];
var innerShapeB = transposeB ? $b.shape[$b.rank - 1] : $b.shape[$b.rank - 2];
var outerShapeA = transposeA ? $a.shape[$a.rank - 1] : $a.shape[$a.rank - 2];
var outerShapeB = transposeB ? $b.shape[$b.rank - 2] : $b.shape[$b.rank - 1];
var outerDimsA = $a.shape.slice(0, -2);
var outerDimsB = $b.shape.slice(0, -2);
var batchDimA = util.sizeFromShape(outerDimsA);
var batchDimB = util.sizeFromShape(outerDimsB);
util.assert($a.rank >= 2 && $b.rank >= 2 && $a.rank === $b.rank, function () {
return "Error in fused matMul: inputs must have the same rank of at least " +
("2, got ranks " + $a.rank + " and " + $b.rank + ".");
});
util.assert(util.arraysEqual(outerDimsA, outerDimsB), function () { return "Error in fused matMul: outer dimensions (" + outerDimsA + ") and (" +
(outerDimsB + ") of Tensors with shapes " + $a.shape + " and ") +
($b.shape + " must match."); });
util.assert(innerShapeA === innerShapeB, function () { return "Error in fused matMul: inner shapes (" + innerShapeA + ") and (" +
(innerShapeB + ") of Tensors with shapes " + $a.shape + " and ") +
($b.shape + " and transposeA=" + transposeA) +
(" and transposeB=" + transposeB + " must match."); });
var outShape = $a.shape.slice(0, -2).concat([outerShapeA, outerShapeB]);
var a3D = transposeA ? $a.as3D(batchDimA, innerShapeA, outerShapeA) :
$a.as3D(batchDimA, outerShapeA, innerShapeA);
var b3D = transposeB ? $b.as3D(batchDimB, outerShapeB, innerShapeB) :
$b.as3D(batchDimB, innerShapeB, outerShapeB);
var $bias;
if (bias != null) {
$bias = tensor_util_env_1.convertToTensor(bias, 'bias', 'fused matMul');
$bias = tensor_util_1.makeTypesMatch($bias, $a)[0];
broadcast_util.assertAndGetBroadcastShape(outShape, $bias.shape);
}
var grad = function (dy, saved) {
var a3D = saved[0], b3D = saved[1], y = saved[2];
var dyActivation;
if (activation == null || activation === 'linear') {
dyActivation = dy;
}
else if (activation === 'relu') {
dyActivation = dy.mul(y.step());
}
else {
throw new Error("Gradient for activation " + activation + " has not been " +
"implemented yet.");
}
var biasGradient = {};
if (bias != null) {
biasGradient = {
$bias: function () {
var res = dyActivation;
// Using dyActivation as reference shape because outputShape does not
// account for the fact that we temporarily reshape inputs to 3D as
// part of batched matMul.
var reduceAxes = broadcast_util.getReductionAxes($bias.shape, dyActivation.shape);
if (reduceAxes.length > 0) {
res = res.sum(reduceAxes);
}
return res.reshape($bias.shape);
}
};
}
if (!transposeA && !transposeB) {
return Object.assign({
$a: function () { return dyActivation.matMul(b3D, false, true); },
$b: function () { return a3D.matMul(dyActivation, true, false); }
}, biasGradient);
}
else if (!transposeA && transposeB) {
return Object.assign({
$a: function () { return dyActivation.matMul(b3D, false, false); },
$b: function () { return dyActivation.matMul(a3D, true, false); }
}, biasGradient);
}
else if (transposeA && !transposeB) {
return Object.assign({
$a: function () { return b3D.matMul(dyActivation, false, true); },
$b: function () { return a3D.matMul(dyActivation, false, false); }
}, biasGradient);
}
else {
return Object.assign({
$a: function () { return b3D.matMul(dyActivation, true, true); },
$b: function () { return dyActivation.matMul(a3D, true, true); }
}, biasGradient);
}
};
var inputs = { $a: a3D, $b: b3D };
if (bias != null) {
inputs.$bias = $bias;
}
var res = engine_1.ENGINE.runKernel(function (backend, save) {
var y = backend.fusedBatchMatMul(a3D, b3D, transposeA, transposeB, $bias, activation);
save([a3D, b3D, y]);
return y;
}, inputs, grad);
return res.reshape(outShape);
}
/**
* Computes a 2D convolution over the input x, optionally fused with adding a
* bias and applying an activation.
*
* @param x The input tensor, of rank 4 or rank 3, of shape
* `[batch, height, width, inChannels]`. If rank 3, batch of 1 is
* assumed.
* @param filter The filter, rank 4, of shape
* `[filterHeight, filterWidth, inDepth, outDepth]`.
* @param strides The strides of the convolution: `[strideHeight,
* strideWidth]`.
* @param pad The type of padding algorithm.
* - `same` and stride 1: output will be of same size as input,
* regardless of filter size.
* - `valid`: output will be smaller than input if filter is larger
* than 1x1.
* - For more info, see this guide:
* [https://www.tensorflow.org/api_guides/python/nn#Convolution](
* https://www.tensorflow.org/api_guides/python/nn#Convolution)
* @param dataFormat: An optional string from: "NHWC", "NCHW". Defaults to
* "NHWC". Specify the data format of the input and output data. With the
* default format "NHWC", the data is stored in the order of: [batch,
* height, width, channels]. Only "NHWC" is currently supported.
* @param dilations The dilation rates: `[dilationHeight, dilationWidth]`
* in which we sample input values across the height and width dimensions
* in atrous convolution. Defaults to `[1, 1]`. If `dilations` is a single
* number, then `dilationHeight == dilationWidth`. If it is greater than
* 1, then all values of `strides` must be 1.
* @param dimRoundingMode The rounding mode used when computing output
* dimensions if pad is a number. If none is provided, it will not round
* and error if the output is of fractional size.
* @param bias Tensor to be added to the result.
* @param activation Name of activation kernel (defaults to `linear`).
*/
/** @doc {heading: 'Operations', subheading: 'Convolution'} */
function conv2d_(x, filter, strides, pad, dataFormat, dilations, dimRoundingMode, bias, activation) {
if (dataFormat === void 0) { dataFormat = 'NHWC'; }
if (dilations === void 0) { dilations = [1, 1]; }
if (activation === void 0) { activation = 'linear'; }
var $x = tensor_util_env_1.convertToTensor(x, 'x', 'conv2d');
var $filter = tensor_util_env_1.convertToTensor(filter, 'filter', 'conv2d');
var x4D = $x;
var reshapedTo4D = false;
if ($x.rank === 3) {
reshapedTo4D = true;
x4D = $x.as4D(1, $x.shape[0], $x.shape[1], $x.shape[2]);
}
util.assert(x4D.rank === 4, function () { return "Error in fused conv2d: input must be rank 4, but got rank " +
(x4D.rank + "."); });
util.assert($filter.rank === 4, function () { return "Error in fused conv2d: filter must be rank 4, but got rank " +
($filter.rank + "."); });
if (dimRoundingMode != null) {
util.assert(util.isInt(pad), function () { return "Error in fused conv2d: pad must be an integer when using, " +
("dimRoundingMode " + dimRoundingMode + " but got pad " + pad + "."); });
}
util.assert(x4D.shape[3] === $filter.shape[2], function () { return "Error in conv2d: depth of input (" + x4D.shape[3] + ") must match " +
("input depth for filter " + $filter.shape[2] + "."); });
util.assert(conv_util.eitherStridesOrDilationsAreOne(strides, dilations), function () { return 'Error in conv2D: Either strides or dilations must be 1. ' +
("Got strides " + strides + " and dilations '" + dilations + "'"); });
util.assert(dataFormat === 'NHWC', function () { return "Error in conv2d: got dataFormat of " + dataFormat + " but only NHWC is currently supported."; });
var convInfo = conv_util.computeConv2DInfo(x4D.shape, $filter.shape, strides, dilations, pad, dimRoundingMode);
var $bias;
if (bias != null) {
$bias = tensor_util_env_1.convertToTensor(bias, 'bias', 'fused conv2d');
$bias = tensor_util_1.makeTypesMatch($bias, $x)[0];
broadcast_util.assertAndGetBroadcastShape(convInfo.outShape, $bias.shape);
}
var grad = function (dy, saved) {
var _a = saved, $filter = _a[0], x4D = _a[1], y = _a[2];
var dyActivation;
if (activation == null || activation === 'linear') {
dyActivation = dy;
}
else if (activation === 'relu') {
dyActivation = dy.mul(y.step());
}
else {
throw new Error("Gradient for activation " + activation + " has not been " +
"implemented yet.");
}
util.assert(conv_util.tupleValuesAreOne(dilations), function () { return 'Error in gradient of fused conv2D: ' +
"dilation rates greater than 1 " +
("are not yet supported in gradients. Got dilations '" + dilations + "'"); });
var biasGradient = {};
if (bias != null) {
biasGradient = {
$bias: function () {
var res = dyActivation;
var reduceAxes = broadcast_util.getReductionAxes($bias.shape, dyActivation.shape);
if (reduceAxes.length > 0) {
res = res.sum(reduceAxes);
}
return res.reshape($bias.shape);
}
};
}
return Object.assign({
x: function () {
return conv_1.conv2dDerInput(x4D.shape, dyActivation, $filter, strides, pad);
},
$filter: function () {
return conv_1.conv2dDerFilter(x4D, dyActivation, $filter.shape, strides, pad);
}
}, biasGradient);
};
var inputs = { x: x4D, $filter: $filter };
if (bias != null) {
inputs.$bias = $bias;
}
var res = engine_1.ENGINE.runKernel(function (backend, save) {
var res = backend.fusedConv2d(x4D, $filter, convInfo, $bias, activation);
save([$filter, x4D, res]);
return res;
}, inputs, grad);
if (reshapedTo4D) {
return res.as3D(res.shape[1], res.shape[2], res.shape[3]);
}
return res;
}
exports.matMul = operation_1.op({ matMul_: matMul_ });
exports.conv2d = operation_1.op({ conv2d_: conv2d_ });
//# sourceMappingURL=fused_ops.js.map