Modularise resizeBilinear and resizeNearestNeighbor

tfjs-backend-wasm/src/kernels/ResizeBilinear.ts

@@ -15,22 +15,12 @@
  * =============================================================================
  */
 
-import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo, util} from '@tensorflow/tfjs-core';
+import {NamedAttrMap, NamedTensorInfoMap, registerKernel, ResizeBilinear, ResizeBilinearAttrs, ResizeBilinearInputs, TensorInfo, util} from '@tensorflow/tfjs-core';
 
 import {BackendWasm} from '../backend_wasm';
 
 import {cast} from './Cast';
 
-interface ResizeBilinearInputs extends NamedTensorInfoMap {
-  x: TensorInfo;
-}
-
-interface ResizeBilinearAttrs extends NamedAttrMap {
-  newWidth: number;
-  newHeight: number;
-  alignCorners: boolean;
-}
-
 let wasmResizeBilinear: (
     xId: number, batch: number, oldHeight: number, oldWidth: number,
     numChannels: number, newHeight: number, newWidth: number,
@@ -52,26 +42,29 @@ function setup(backend: BackendWasm): void {
 
 function resizeBilinear(args: {
   backend: BackendWasm,
-  inputs: ResizeBilinearInputs,
-  attrs: ResizeBilinearAttrs
+  inputs: NamedTensorInfoMap,
+  attrs: NamedAttrMap
 }): TensorInfo {
   const {backend, inputs, attrs} = args;
-  const {x} = inputs;
-  const {alignCorners, newHeight, newWidth} = attrs;
 
-  const [batch, oldHeight, oldWidth, numChannels] = x.shape;
+  const {images} = inputs as ResizeBilinearInputs;
+  const {alignCorners, size} = attrs as {} as ResizeBilinearAttrs;
+  const [newHeight, newWidth] = size;
+
+  const [batch, oldHeight, oldWidth, numChannels] = images.shape;
   const outShape = [batch, newHeight, newWidth, numChannels];
 
-  let xData = backend.dataIdMap.get(x.dataId);
+  let xData = backend.dataIdMap.get(images.dataId);
   let castedData;
   if (xData.dtype !== 'float32') {
-    castedData = cast({backend, inputs: {x}, attrs: {dtype: 'float32'}});
+    castedData =
+        cast({backend, inputs: {x: images}, attrs: {dtype: 'float32'}});
     xData = backend.dataIdMap.get(castedData.dataId);
   }
   const xId = xData.id;
 
   const out = backend.makeOutput(outShape, 'float32');
-  if (util.sizeFromShape(x.shape) === 0) {
+  if (util.sizeFromShape(images.shape) === 0) {
     return out;
   }
   const outId = backend.dataIdMap.get(out.dataId).id;
@@ -88,7 +81,7 @@ function resizeBilinear(args: {
 }
 
 registerKernel({
-  kernelName: 'ResizeBilinear',
+  kernelName: ResizeBilinear,
   backendName: 'wasm',
   setupFunc: setup,
   kernelFunc: resizeBilinear

tfjs-core/src/gradients/ResizeBilinear_grad.ts

@@ -0,0 +1,43 @@
+/**
+ * @license
+ * Copyright 2020 Google Inc. 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, ForwardFunc} from '../engine';
+import {ResizeBilinear, ResizeBilinearAttrs, ResizeBilinearGrad, ResizeBilinearGradInputs} from '../kernel_names';
+import {GradConfig} from '../kernel_registry';
+import {NamedAttrMap} from '../kernel_registry';
+import {Tensor, Tensor4D} from '../tensor';
+import {NamedTensorMap} from '../tensor_types';
+
+export const resizeBilinearGradConfig: GradConfig = {
+  kernelName: ResizeBilinear,
+  inputsToSave: ['images'],
+  gradFunc: (dy: Tensor4D, saved: Tensor[], attrs: NamedAttrMap) => {
+    const [images] = saved;
+
+    const backPropKernelFunc: ForwardFunc<Tensor> = (backend) => {
+      const {alignCorners} = attrs as {} as ResizeBilinearAttrs;
+      return backend.resizeBilinearBackprop(
+          dy, images as Tensor4D, alignCorners);
+    };
+
+    const inputs: ResizeBilinearGradInputs = {images};
+    const imagesDer = () => ENGINE.runKernelFunc(
+        backPropKernelFunc, inputs as {} as NamedTensorMap, null /* gradient */,
+        ResizeBilinearGrad, attrs);
+
+    return {images: imagesDer};
+  }
+};

tfjs-core/src/gradients/ResizeNearestNeighbor_grad.ts

@@ -0,0 +1,43 @@
+/**
+ * @license
+ * Copyright 2020 Google Inc. 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, ForwardFunc} from '../engine';
+import {ResizeNearestNeighbor, ResizeNearestNeighborAttrs, ResizeNearestNeighborGrad, ResizeNearestNeighborGradInputs} from '../kernel_names';
+import {GradConfig} from '../kernel_registry';
+import {NamedAttrMap} from '../kernel_registry';
+import {Tensor, Tensor4D} from '../tensor';
+import {NamedTensorMap} from '../tensor_types';
+
+export const resizeNearestNeighborGradConfig: GradConfig = {
+  kernelName: ResizeNearestNeighbor,
+  inputsToSave: ['images'],
+  gradFunc: (dy: Tensor4D, saved: Tensor[], attrs: NamedAttrMap) => {
+    const [images] = saved;
+
+    const backPropKernelFunc: ForwardFunc<Tensor> = (backend) => {
+      const {alignCorners} = attrs as {} as ResizeNearestNeighborAttrs;
+      return backend.resizeNearestNeighborBackprop(
+          dy, images as Tensor4D, alignCorners);
+    };
+
+    const inputs: ResizeNearestNeighborGradInputs = {images};
+    const imagesDer = () => ENGINE.runKernelFunc(
+        backPropKernelFunc, inputs as {} as NamedTensorMap, null /* gradient */,
+        ResizeNearestNeighborGrad, attrs);
+
+    return {images: imagesDer};
+  }
+};

tfjs-core/src/kernel_names.ts

@@ -381,6 +381,27 @@ export type RealInputs = Pick<NamedTensorInfoMap, 'input'>;
 export const Relu = 'Relu';
 export type ReluInputs = Pick<NamedTensorInfoMap, 'x'>;
 
+export const ResizeNearestNeighbor = 'ResizeNearestNeighbor';
+export type ResizeNearestNeighborInputs = Pick<NamedTensorInfoMap, 'images'>;
+export interface ResizeNearestNeighborAttrs {
+  alignCorners: boolean;
+  size: [number, number];
+}
+
+export const ResizeNearestNeighborGrad = 'ResizeNearestNeighborGrad';
+export type ResizeNearestNeighborGradInputs =
+    Pick<NamedTensorInfoMap, 'images'>;
+
+export const ResizeBilinear = 'ResizeBilinear';
+export type ResizeBilinearInputs = Pick<NamedTensorInfoMap, 'images'>;
+export interface ResizeBilinearAttrs {
+  alignCorners: boolean;
+  size: [number, number];
+}
+
+export const ResizeBilinearGrad = 'ResizeBilinearGrad';
+export type ResizeBilinearGradInputs = Pick<NamedTensorInfoMap, 'images'>;
+
 export const Relu6 = 'Relu6';
 export type Relu6Inputs = Pick<NamedTensorInfoMap, 'x'>;
 

tfjs-core/src/ops/image_ops.ts

@@ -17,7 +17,7 @@
 
 import {nonMaxSuppressionV3, nonMaxSuppressionV5} from '../backends/non_max_suppression_impl';
 import {ENGINE, ForwardFunc} from '../engine';
-import {Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D} from '../tensor';
+import {Tensor, Tensor1D, Tensor2D, Tensor4D} from '../tensor';
 import {NamedTensorMap} from '../tensor_types';
 import {convertToTensor} from '../tensor_util_env';
 import {TensorLike} from '../types';
@@ -27,125 +27,6 @@ import {nonMaxSuppSanityCheck} from './nonmax_util';
 import {op} from './operation';
 export {nonMaxSuppression} from './non_max_suppression';
 
-/**
- * Bilinear resize a batch of 3D images to a new shape.
- *
- * @param images The images, of rank 4 or rank 3, of shape
- *     `[batch, height, width, inChannels]`. If rank 3, batch of 1 is assumed.
- * @param size The new shape `[newHeight, newWidth]` to resize the
- *     images to. Each channel is resized individually.
- * @param alignCorners Defaults to False. If true, rescale
- *     input by `(new_height - 1) / (height - 1)`, which exactly aligns the 4
- *     corners of images and resized images. If false, rescale by
- *     `new_height / height`. Treat similarly the width dimension.
- */
-/** @doc {heading: 'Operations', subheading: 'Images', namespace: 'image'} */
-function resizeBilinear_<T extends Tensor3D|Tensor4D>(
-    images: T|TensorLike, size: [number, number], alignCorners = false): T {
-  const $images = convertToTensor(images, 'images', 'resizeBilinear');
-  util.assert(
-      $images.rank === 3 || $images.rank === 4,
-      () => `Error in resizeBilinear: x must be rank 3 or 4, but got ` +
-          `rank ${$images.rank}.`);
-  util.assert(
-      size.length === 2,
-      () => `Error in resizeBilinear: new shape must 2D, but got shape ` +
-          `${size}.`);
-
-  let batchImages = $images as Tensor4D;
-  let reshapedTo4D = false;
-  if ($images.rank === 3) {
-    reshapedTo4D = true;
-    batchImages =
-        $images.as4D(1, $images.shape[0], $images.shape[1], $images.shape[2]);
-  }
-
-  const [newHeight, newWidth] = size;
-  const forward: ForwardFunc<Tensor4D> = (backend, save) => {
-    save([batchImages]);
-    return backend.resizeBilinear(
-        batchImages, newHeight, newWidth, alignCorners);
-  };
-
-  const backward = (dy: Tensor4D, saved: Tensor[]) => {
-    return {
-      x: () => ENGINE.runKernelFunc(
-          backend => backend.resizeBilinearBackprop(
-              dy, saved[0] as Tensor4D, alignCorners),
-          {})
-    };
-  };
-
-  const res = ENGINE.runKernelFunc(
-      forward, {x: batchImages}, backward, 'ResizeBilinear',
-      {alignCorners, newHeight, newWidth});
-  if (reshapedTo4D) {
-    return res.as3D(res.shape[1], res.shape[2], res.shape[3]) as T;
-  }
-  return res as T;
-}
-
-/**
- * NearestNeighbor resize a batch of 3D images to a new shape.
- *
- * @param images The images, of rank 4 or rank 3, of shape
- *     `[batch, height, width, inChannels]`. If rank 3, batch of 1 is assumed.
- * @param size The new shape `[newHeight, newWidth]` to resize the
- *     images to. Each channel is resized individually.
- * @param alignCorners Defaults to False. If true, rescale
- *     input by `(new_height - 1) / (height - 1)`, which exactly aligns the 4
- *     corners of images and resized images. If false, rescale by
- *     `new_height / height`. Treat similarly the width dimension.
- */
-/** @doc {heading: 'Operations', subheading: 'Images', namespace: 'image'} */
-function resizeNearestNeighbor_<T extends Tensor3D|Tensor4D>(
-    images: T|TensorLike, size: [number, number], alignCorners = false): T {
-  const $images = convertToTensor(images, 'images', 'resizeNearestNeighbor');
-  util.assert(
-      $images.rank === 3 || $images.rank === 4,
-      () => `Error in resizeNearestNeighbor: x must be rank 3 or 4, but got ` +
-          `rank ${$images.rank}.`);
-  util.assert(
-      size.length === 2,
-      () =>
-          `Error in resizeNearestNeighbor: new shape must 2D, but got shape ` +
-          `${size}.`);
-  util.assert(
-      $images.dtype === 'float32' || $images.dtype === 'int32',
-      () => '`images` must have `int32` or `float32` as dtype');
-
-  let batchImages = $images as Tensor4D;
-  let reshapedTo4D = false;
-  if ($images.rank === 3) {
-    reshapedTo4D = true;
-    batchImages =
-        $images.as4D(1, $images.shape[0], $images.shape[1], $images.shape[2]);
-  }
-  const [newHeight, newWidth] = size;
-
-  const forward: ForwardFunc<Tensor4D> = (backend, save) => {
-    save([batchImages]);
-    return backend.resizeNearestNeighbor(
-        batchImages, newHeight, newWidth, alignCorners);
-  };
-
-  const backward = (dy: Tensor4D, saved: Tensor[]) => {
-    return {
-      batchImages: () => ENGINE.runKernelFunc(
-          backend => backend.resizeNearestNeighborBackprop(
-              dy, saved[0] as Tensor4D, alignCorners),
-          {})
-    };
-  };
-
-  const res = ENGINE.runKernelFunc(forward, {batchImages}, backward);
-
-  if (reshapedTo4D) {
-    return res.as3D(res.shape[1], res.shape[2], res.shape[3]) as T;
-  }
-  return res as T;
-}
-
 /**
  * Performs non maximum suppression of bounding boxes based on
  * iou (intersection over union).
@@ -350,8 +231,6 @@ function cropAndResize_(
   return res;
 }
 
-export const resizeBilinear = op({resizeBilinear_});
-export const resizeNearestNeighbor = op({resizeNearestNeighbor_});
 export const nonMaxSuppressionAsync = nonMaxSuppressionAsync_;
 export const nonMaxSuppressionWithScore = op({nonMaxSuppressionWithScore_});
 export const nonMaxSuppressionWithScoreAsync = nonMaxSuppressionWithScoreAsync_;

tfjs-core/src/ops/ops.ts

@@ -127,9 +127,21 @@ export {op} from './operation';
 // Second level exports.
 import * as losses from './loss_ops';
 import * as linalg from './linalg_ops';
-import * as image from './image_ops';
 import * as spectral from './spectral_ops';
 import * as fused from './fused_ops';
 import * as signal from './signal_ops';
 
+import {cropAndResize, nonMaxSuppression, nonMaxSuppressionAsync, nonMaxSuppressionWithScore, nonMaxSuppressionWithScoreAsync} from './image_ops';
+import {resizeBilinear} from './resize_bilinear';
+import {resizeNearestNeighbor} from './resize_nearest_neighbor';
+const image = {
+  resizeNearestNeighbor,
+  resizeBilinear,
+  cropAndResize,
+  nonMaxSuppression,
+  nonMaxSuppressionAsync,
+  nonMaxSuppressionWithScore,
+  nonMaxSuppressionWithScoreAsync
+};
+
 export {image, linalg, losses, spectral, fused, signal};