@stellarapp/tfjs-stellar 1.0.0

package/src/models/u_net.ts

@@ -0,0 +1,240 @@
+import * as tf from "@tensorflow/tfjs";
+import { type ActivationIdentifier } from "@tensorflow/tfjs-layers/dist/keras_format/activation_config";
+
+
+export interface UNetArgs {
+    /**
+     * The starting number of filters.
+     */
+    filters: number;
+    /**
+     * The number of categories. For binary segmentation, `units=1`.
+     */
+    units: number;
+    /**
+     * The activation of the final output convolution layer. Defaults to `sigmoid` if `categories=1`, else `softmax`.
+     */
+    activation?: ActivationIdentifier;
+    /**
+     * The depth of the U-Net or the number of contractions and the number of expansions.
+     */
+    depth: number;
+    /**
+     * Adds residual connections to transform the model into a ResUNet. Defaults to `false`.
+     */
+    residual?: boolean;
+    /**
+     * Adds batch normalization to convolutions. Best used for batched inputs. Defaults to `false`.
+     */
+    batchNorm?: boolean;
+    /**
+     * Set the unbatched input shape of the U-Net in the format `[height, width, channels]`. Defaults to `[null, null, 3]`. If set, only channels is mandatory.
+     */
+    inputShape?: [number | null, number | null, number];
+}
+
+
+export type UNetModelArgs = UNetArgs & Omit<tf.SequentialArgs, "layers">;
+
+
+export class UNetModel extends tf.Sequential {
+
+    constructor(args: UNetModelArgs) {
+        const {
+            filters,
+            units,
+            activation = units == 1 ? "sigmoid" : "softmax",
+            depth,
+            residual = false,
+            batchNorm = false,
+            inputShape = [null, null, 3],
+            ...sequentialArgs
+        } = args;
+
+        sequentialArgs.name = sequentialArgs.name ?? "unet_model";
+
+        super({
+            ...sequentialArgs,
+            // calling user should not modify the layers after instantiation
+            layers: [createUNet({ filters, units, activation, depth, residual, batchNorm, inputShape })]
+        });
+    }
+
+
+    override summary(lineLength?: number, positions?: number[], printFn?: (message?: any, ...optionalParams: any[]) => void): void {
+        super.summary(lineLength, positions, printFn);
+        (this.layers[0] as tf.LayersModel).summary(lineLength, positions, printFn);
+    }
+}
+
+
+export function createUNet({ filters, depth, units, activation, residual = false, batchNorm = false, inputShape = [null, null, 3] }: UNetModelArgs) {
+    if (units < 1) {
+        throw Error(`createUNet: units should be >= 1, got ${units}`);
+    }
+
+    const [image_height, image_width] = inputShape;
+    const divisble_by = 2 ** depth;
+
+    if ((image_height != null && image_height % divisble_by != 0) ||
+        image_width != null && image_width % divisble_by != 0) {
+        throw Error(`createUNet: the input height and width must be divisible by 2^depth (${divisble_by})`)
+    }
+
+    const input = tf.input({ shape: inputShape });
+
+    const skip_connection: tf.SymbolicTensor[] = [];
+
+    let x = input;
+
+    // calculate the filter sizes for each level
+    const filter_sizes = Array.from({ length: depth }, (_, i) => filters * (2 ** i));
+
+    for (const filter_size of filter_sizes) {
+        const contraction = contractionBlock(x, filter_size, residual, batchNorm, `contraction-f${filter_size}`);
+
+        x = tf.layers.maxPooling2d({ poolSize: 2, strides: 2, name: `pool-f${filter_size}` }).apply(contraction) as tf.SymbolicTensor;
+        skip_connection.push(contraction);
+    }
+
+    x = contractionBlock(x, filter_sizes.at(-1)! * 2, residual, batchNorm, "bottleneck");
+
+    for (let i = filter_sizes.length - 1; i >= 0; i--) {
+        x = expansionBlock(x, skip_connection[i], filter_sizes[i], residual, batchNorm, `expansion-f${filter_sizes[i]}`);
+    }
+
+    const output = tf.layers.conv2d({
+        filters: units,
+        kernelSize: 1,
+        padding: "same",
+        activation: activation ?? (units == 1 ? "sigmoid" : "softmax"),
+        name: "output-conv"
+    }).apply(x) as tf.SymbolicTensor;
+
+    return tf.model({ inputs: input, outputs: output, name: "u_net" });
+}
+
+
+export async function loadUNetModel(pathOrIOHandler: string | tf.io.IOHandler, options?: tf.io.LoadOptions) {
+    const model = await tf.loadLayersModel(pathOrIOHandler, options);
+    const unet = createUNet({ depth: 1, filters: 4, units: 1 }); // these are dummy args that are overwritten
+    const { name, ...rest } = model;
+    Object.assign(unet, rest);
+
+    return unet;
+}
+
+
+/**
+ * The contraction block of a U-Net
+ * 
+ * Conv > BN > ReLU > Conv > BN + residual > ReLU
+ * 
+ * TODO: for residual, change order to (BN > ReLU > Conv)x2 + residual
+ * 
+ * @param x a previous layer's symbolic output
+ * @param filters the number of filters, usually half the previous expansion block's
+ * @param residual includes a residual connection
+ * @param batchNorm applies batch normalization before ReLU activation
+ * @param name a unique name for the contraction block
+ */
+function contractionBlock(x: tf.SymbolicTensor, filters: number, residual: boolean, batchNorm: boolean, name: string) {
+
+    const conv1 = tf.layers.conv2d({
+        filters,
+        kernelSize: 3,
+        padding: "same",
+        useBias: !batchNorm,
+        kernelInitializer: "heNormal",
+        name: `${name}-1-conv2d`
+    });
+    const relu1 = tf.layers.reLU({ name: `${name}-1-relu` });
+
+    const conv2 = tf.layers.conv2d({
+        filters,
+        kernelSize: 3,
+        padding: "same",
+        useBias: !batchNorm,
+        kernelInitializer: "heNormal",
+        name: `${name}-2-conv2d`
+    });
+    const relu2 = tf.layers.reLU({ name: `${name}-2-relu` });
+
+    let forward = conv1.apply(x);
+
+    if (batchNorm) {
+        forward = tf.layers.batchNormalization({ name: `${name}-1-batchnorm` }).apply(forward);
+    }
+
+    forward = relu1.apply(forward);
+
+    forward = conv2.apply(forward);
+
+    if (batchNorm) {
+        forward = tf.layers.batchNormalization({ name: `${name}-2-batchnorm` }).apply(forward);
+    }
+
+    if (residual) {
+        let residual_skip = x;
+
+        if (x.shape[x.shape.length - 1] != filters) {
+            // a 1x1 convolution on the input to ensure the residual connection's
+            // channels/filters dim matches the convolution output
+            residual_skip = tf.layers.conv2d({
+                filters,
+                kernelSize: 1,
+                padding: "same",
+                useBias: !batchNorm,
+                kernelInitializer: "heNormal",
+                name: `${name}-residual`
+            }).apply(x) as tf.SymbolicTensor;
+        }
+
+        if (batchNorm) {
+            residual_skip = tf.layers.batchNormalization({
+                name: `${name}-residual-batchnorm`
+            }).apply(residual_skip) as tf.SymbolicTensor;
+        }
+
+        forward = tf.layers.add().apply([
+            residual_skip as tf.SymbolicTensor,
+            forward as tf.SymbolicTensor
+        ])
+    }
+
+    forward = relu2.apply(forward);
+
+    return forward as tf.SymbolicTensor;
+}
+
+
+/**
+ * The expansion block of a U-Net
+ * 
+ * Upconv + skip > contraction block
+ *  
+ * @param x a previous layer's symbolic output
+ * @param skip the corresponding contraction block's output (before pool), shape matches `x`
+ * @param filters the number of filters, usually half the previous expansion block's
+ * @param residual includes a residual connection
+ * @param batchNorm apply batch normalization, should be `false` when batch size is `1`
+ * @param name a unique name for the contraction block
+ */
+function expansionBlock(x: tf.SymbolicTensor, skip: tf.SymbolicTensor, filters: number, residual: boolean, batchNorm: boolean, name: string) {
+
+    const upconv = tf.layers.conv2dTranspose({
+        filters,
+        padding: "same",
+        kernelSize: 2,
+        strides: 2,
+        kernelInitializer: "heNormal",
+        name: `${name}-upconv`
+    });
+
+    const concat = tf.layers.concatenate({ axis: -1, name: `${name}-concat-upconv-skip` });
+
+    let forward = upconv.apply(x) as tf.SymbolicTensor;
+    forward = concat.apply([forward, skip]) as tf.SymbolicTensor;
+
+    return contractionBlock(forward, filters, residual, batchNorm, name);
+}

package/src/packing_mask.ts

@@ -0,0 +1,28 @@
+import * as tf from "@tensorflow/tfjs";
+
+/**
+ * Generate a self-attention mask that prevents packed sequences from cross document
+ * boundaries and attending to each other. The result is a tensor of diagonally
+ * positioned blocks of zeroes (allow attention) and -1e7 values (prevent attention).
+ * The latter is scored zero during the scaled dot product attention's softmax operation.
+ * 
+ * @param boundaries an array of ones (denotes start of a new sample or docment) and zeroes
+ * 
+ * Example boundary of 3 samples are packed into one:
+ * `[1, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0]`
+ */
+export function generatePackingSelfAttentionMask(boundaries: Int32Array) {
+    // see images at
+    // https://reddit.com/r/LocalLLaMA/comments/197efaz/training_llama_mistral_and_mixtralmoe_faster_with/
+    return tf.tidy(() => {
+        // cumsum transforms the tensor such that each sequence in the pack gets its own id,
+        const partitions = tf.tensor1d(boundaries).cumsum();
+
+        return partitions.expandDims(1)
+            .equal(partitions.expandDims(0))
+            .sub(1)
+            .mul(1e7)
+            // introduce a head dimension so it can be broadcasted
+            .expandDims(0);
+    })
+}

package/src/testing.ts

@@ -0,0 +1 @@
+console.log("test")

package/src/tfjs_types.ts

@@ -0,0 +1,15 @@
+import type { Tensor } from "@tensorflow/tfjs";
+
+
+export declare abstract class LazyIterator<T> {
+    abstract next(): Promise<IteratorResult<T>>;
+}
+
+
+export declare abstract class Dataset<T> {
+    abstract iterator(): Promise<LazyIterator<T>>;
+    size: number;
+}
+
+
+export type LossOrMetricFn = (yTrue: Tensor, yPred: Tensor) => Tensor;

package/src/utils.test.ts

@@ -0,0 +1,101 @@
+import * as tf from "@tensorflow/tfjs";
+import { getScaleShape, getRandomCropStart, generateCausalAttentionMask } from "@/utils";
+
+// avoid TFJS node message during Jest testing
+tf.env().set('IS_NODE', false);
+
+
+describe("test custom TFJS utility functions", () => {
+
+    test("crop an image using the same shape, results in same shape", async () => {
+        // cropping an image of the same shape
+        const img_size = [133, 84] as [number, number];
+        const target_size = [133, 84] as [number, number];
+
+        expect(getRandomCropStart(img_size, target_size)).toEqual([0, 0, 0]);
+    });
+
+
+    it("should throw when crop is larger than image", async () => {
+        expect(() => getRandomCropStart([128, 128], [1000, 2000])).toThrow();
+    })
+
+
+    test("cropped image shape", async () => {
+        // cropping from wide to tall image
+        for (let i = 0; i < 100; i++) {
+            const img_size = [4923, 832] as [number, number];
+            const target_size = [333, 739] as [number, number];
+
+            const [crop_start_h, crop_start_w, channels] = getRandomCropStart(img_size, target_size)
+
+            expect(crop_start_h).toBeLessThanOrEqual(img_size[0] - target_size[0]);
+            expect(crop_start_w).toBeLessThanOrEqual(img_size[1] - target_size[1]);
+        }
+
+        // cropping from tall to wide image
+        for (let i = 0; i < 100; i++) {
+            const img_size = [381, 999] as [number, number];
+            const target_size = [300, 157] as [number, number];
+
+            const [crop_start_h, crop_start_w, channels] = getRandomCropStart(img_size, target_size)
+
+            expect(crop_start_h).toBeLessThanOrEqual(img_size[0] - target_size[0]);
+            expect(crop_start_w).toBeLessThanOrEqual(img_size[1] - target_size[1]);
+        }
+    });
+
+
+    test("scale 1:1, results in the same shape", async () => {
+        const scale = getScaleShape([256, 256], [256, 256])
+        expect(scale).toEqual([256, 256]);
+    });
+
+
+    test("scaled image shape", async () => {
+        // scaling squares result in squares
+        const scale1 = getScaleShape([256, 256], [128, 128])
+        expect(scale1).toEqual([128, 128]);
+
+        const scale2 = getScaleShape([128, 128], [256, 256])
+        expect(scale2).toEqual([256, 256]);
+
+        const scale3 = getScaleShape([123, 123], [321, 321])
+        expect(scale3).toEqual([321, 321]);
+
+        const scale4 = getScaleShape([321, 321], [123, 123])
+        expect(scale4).toEqual([123, 123]);
+
+        // scaling rectangles result in rectangles
+        const scale5 = getScaleShape([640, 480], [1280, 960])
+        expect(scale5).toEqual([1280, 960]);
+
+        const scale6 = getScaleShape([480, 640], [960, 1280])
+        expect(scale6).toEqual([960, 1280]);
+
+        const [scale7_h, scale7_w] = getScaleShape([777, 555], [555, 333])
+        expect(scale7_h).toBeGreaterThan(scale7_w);
+
+        const [scale8_h, scale8_w] = getScaleShape([555, 777], [333, 555])
+        expect(scale8_h).toBeLessThan(scale8_w);
+    });
+
+
+    test("causal attention map", async () => {
+        const seq_len = 4;
+        const causal_mask = generateCausalAttentionMask(seq_len, seq_len);
+
+        const _ = -1e7;
+        const expected_mask = tf.tensor([
+            [0, _, _, _],
+            [0, 0, _, _],
+            [0, 0, 0, _],
+            [0, 0, 0, 0]
+        ]);
+
+        // this might fail due to precision issues on the masked positions,
+        // in which case use less <= to 6 or 12 (number of masked positions x2)
+        expect((await causal_mask.sub(expected_mask).sum().data())[0]).toEqual(0);
+    });
+
+});

package/src/utils.ts

@@ -0,0 +1,86 @@
+import * as tf from "@tensorflow/tfjs";
+
+
+/**
+ * Calculate the desired scaled image's height and width. The shortest edge will
+ * be scaled to match its corresponding target shape's edge. The longer
+ * edge might end up larger than intended.
+ * 
+ * @param image_shape the `[height, width]` of the image
+ * @param target_shape the intended `[height, width]` of the final scaled image
+ */
+export function getScaleShape(image_shape: tf.Shape, target_shape: [number, number]): [scaled_height: number, scaled_width: number] {
+    const [img_height, img_width] = image_shape as [number, number, number];
+    const [target_height, target_width] = target_shape;
+
+    // scale based on whichever target_edge / original_edge is largest,
+    // we need the following to be true (1)
+    // height * scale >= target_height
+    // width * scale >= target_height
+    // rearranging to get an equivalent requirement (2)
+    // scale >= target_height / height
+    // scale >= target_width / width
+    // by picking the scale value that's largest of the two, we satisfy (2), and therefore (1)
+    // it may be more intuitive to think of scale as scale_h and scale_w
+    const scale_factor = Math.max(target_height / img_height, target_width / img_width);
+    return [Math.round(img_height * scale_factor), Math.round(img_width * scale_factor)];
+}
+
+
+/**
+ * Calculate the starting point for a crop (slice) operation
+ * on an image tensor with the shape `[height, width, channels]`.
+ * 
+ * @param image_shape the `[height, width]` of the image
+ * @param target_shape the intended `[height, width]` of the final cropped image
+ */
+export function getRandomCropStart(
+    image_shape: [height: number, width: number],
+    target_shape: [height: number, width: number]
+): [number, number, number] {
+    const [img_height, img_width] = image_shape;
+    const [crop_x, crop_y] = target_shape;
+
+    if (img_height < crop_x || img_width < crop_y) {
+        throw Error(`getRandomCropShape: cannot crop with a size that's bigger than,` +
+            ` the image. Original [${img_height}, ${img_width}], crop [${crop_x}, ${crop_y}].`);
+    }
+
+    // there's a +1 because Math.random()'s range is [0, 1), excluding 1,
+    // hence +1 to ensure the full range of possible crop starting points
+    return [
+        // TODO: revisit the +1
+        Math.floor(Math.random() * (img_height - crop_x + 1)),
+        Math.floor(Math.random() * (img_width - crop_y + 1)),
+        0 // not cropping channels, so it starts at the first index
+    ]
+}
+
+
+/**
+ * Calculate the height and width padding such that the image is
+ * divisible by 2^depth.
+ * 
+ * In U-Net image segmentation, the contraction and concatenate
+ * operations requires  the input image's height and width
+ * dimensions to be divisible by 2^depth.
+ */
+export function getPaddingForSegmentation(image: tf.Tensor3D, depth: number): [height: number, width: number] {
+    const divisible = Math.pow(2, depth);
+
+    const [height, width] = image.shape;
+
+    return [
+        (Math.ceil(height / divisible)) * divisible - height,
+        (Math.ceil(width / divisible)) * divisible - width,
+    ]
+}
+
+
+export function generateCausalAttentionMask(query_seq_length: number, key_seq_length: number) {
+    return tf.tidy(() => {
+        return tf.linalg.bandPart(tf.ones([query_seq_length, key_seq_length]), -1, 0)
+            .sub(1)
+            .mul(1e7);
+    })
+}

package/tsconfig.json

@@ -0,0 +1,49 @@
+{
+  // Visit https://aka.ms/tsconfig to read more about this file
+  "compilerOptions": {
+    // File Layout
+    // "rootDir": "./src",
+    // "outDir": "./dist",
+    // Environment Settings
+    // See also https://aka.ms/tsconfig/module
+    "module": "es2022",
+    "target": "esnext",
+    "types": ["jest"],
+    // For nodejs:
+    // "lib": ["esnext"],
+    // "types": ["node"],
+    // and npm install -D @types/node
+    // Other Outputs
+    "sourceMap": true,
+    "declaration": true,
+    "declarationMap": true,
+    // Stricter Typechecking Options
+    //"noUncheckedIndexedAccess": true,
+    "exactOptionalPropertyTypes": true,
+    // Style Options
+    // "noImplicitReturns": true,
+    // "noImplicitOverride": true,
+    // "noUnusedLocals": true,
+    // "noUnusedParameters": true,
+    // "noFallthroughCasesInSwitch": true,
+    // "noPropertyAccessFromIndexSignature": true,
+    // Recommended Options
+    "strict": true,
+    "jsx": "react-jsx",
+    //"verbatimModuleSyntax": true,
+    "isolatedModules": true,
+    "noUncheckedSideEffectImports": true,
+    "moduleDetection": "force",
+    "skipLibCheck": true,
+    "paths": {
+      "@/*": [
+        "./src/*"
+      ],
+      "e2e/*": [
+        "./e2e/*"
+      ]
+    },
+    "moduleResolution": "bundler",
+    "esModuleInterop": true
+  }
+}