From cc302a2d9cafe04081714ccd143582d29a6e184a Mon Sep 17 00:00:00 2001 From: Douglas A Date: Mon, 13 Dec 2021 18:56:51 -0300 Subject: [PATCH] feat: Change the model build, create separated encoder and decoder You dont need to keep track of the decoder layers in order to build it later. Just craete them separatedelly and add to a everything to a single model. --- index.js | 151 ++++++++++++++++++++----------------------------------- 1 file changed, 54 insertions(+), 97 deletions(-) diff --git a/index.js b/index.js index 6ee0bb1..f97a345 100644 --- a/index.js +++ b/index.js @@ -1,6 +1,6 @@ console.log("Hello Autoencoder 🚂"); -import * as tf from "@tensorflow/tfjs-node"; +import * as tf from "@tensorflow/tfjs-node-gpu"; // import canvas from "canvas"; // const { loadImage } = canvas; import Jimp from "jimp"; @@ -11,24 +11,25 @@ const W = 28; main(); async function main() { + + console.log(`Using backend: ${tf.getBackend()}`) + // Build the model - const { decoderLayers, autoencoder } = buildModel(); + const { decoder, autoencoder } = buildModel(); // load all image data - const images = await loadImages(5100); + const images = await loadImages(7000); // train the model - const x_train = tf.tensor2d(images.slice(0, 5000)); - await trainModel(autoencoder, x_train, 200); + const x_train = tf.tensor2d(images.slice(0, 5600)); + await trainModel(autoencoder, x_train, 15); const saveResults = await autoencoder.save("file://public/model/"); - console.log(autoencoder.summary()); - // const autoencoder = await tf.loadLayersModel("file://public/model/model.json"); // test the model const x_test = tf.tensor2d(images.slice(5000)); await generateTests(autoencoder, x_test); // Create a new model with just the decoder - const decoder = createDecoder(decoderLayers); + //const decoder = createDecoder(decoder); const saveDecoder = await decoder.save("file://public/decoder/model/"); } @@ -61,19 +62,10 @@ async function generateTests(autoencoder, x_test) { } } -function createDecoder(decoderLayers) { - const decoder = tf.sequential(); - for (let layer of decoderLayers) { - const newLayer = tf.layers.dense({ - units: layer.units, - activation: layer.activation, - inputShape: [layer.kernel.shape[0]], - }); - decoder.add(newLayer); - newLayer.setWeights(layer.getWeights()); - } - // const learningRate = 0.000001; - // const optimizer = tf.train.adam(learningRate); +function createDecoder(decoder) { + + const learningRate = 0.0001; + const optimizer = tf.train.adam(learningRate, 0.0000001); // adam(learning_rate, decay) decoder.compile({ optimizer: "adam", loss: "meanSquaredError", @@ -82,87 +74,51 @@ function createDecoder(decoderLayers) { } function buildModel() { - const autoencoder = tf.sequential(); - // Build the model - // Encoder - autoencoder.add( - tf.layers.dense({ - units: 256, - inputShape: [W * W], - activation: "relu", - }) - ); - autoencoder.add( - tf.layers.dense({ - units: 128, - activation: "relu", - }) - ); - autoencoder.add( - tf.layers.dense({ - units: 64, - activation: "relu", - }) - ); - autoencoder.add( - tf.layers.dense({ - units: 16, - activation: "relu", - }) - ); - autoencoder.add( - tf.layers.dense({ - units: 4, - activation: "sigmoid", - }) - ); - // How do I start from here? - // Decoder - - let decoderLayers = []; - decoderLayers.push( - tf.layers.dense({ - units: 16, - activation: "relu", - }) - ); - decoderLayers.push( - tf.layers.dense({ - units: 64, - activation: "relu", - }) - ); - decoderLayers.push( - tf.layers.dense({ - units: 128, - activation: "relu", - }) - ); - decoderLayers.push( - tf.layers.dense({ - units: 256, - activation: "relu", - }) - ); - decoderLayers.push( - tf.layers.dense({ - units: W * W, - activation: "sigmoid", - }) - ); - - for (let layer of decoderLayers) { - autoencoder.add(layer); - } - // const learningRate = 0.001; - // const optimizer = tf.train.adam(learningRate); +// encoder + const encoder_input = tf.input({shape: [W*W], name: "encoder_input"}); + // const l0 = tf.layers.flatten().apply(encoder_input); + const l1 = tf.layers.dense({units: 128, activation: "relu"}).apply(encoder_input); + const l2 = tf.layers.dense({units: 64, activation: "relu"}).apply(l1); + const l3 = tf.layers.dense({units: 16, activation: "relu"}).apply(l2); + const l4 = tf.layers.dense({units: 4, activation: "relu"}).apply(l3); + let encoded = tf.layers.dense({units: 2, activation: "relu", name: "encoder_output"}).apply(l4); + + let encoder = tf.model({inputs: encoder_input, outputs: encoded, name: "encoder"}); + console.log(`Encoder Summary: ${encoder.summary()}`); + + const decoder_input = tf.input({shape: [2]}); + let decoder = tf.layers.dense({units: 4, activation: "relu", name: "decoder_input"}).apply(decoder_input); + const l6 = tf.layers.dense({units: 16, activation: "relu"}).apply(decoder); + const l7 = tf.layers.dense({units: 64, activation: "relu"}).apply(l6); + const l8 = tf.layers.dense({units: 128, activation: "relu"}).apply(l7); + let decoded = tf.layers.dense({units: 784, activation: "sigmoid", name: "decoder_output"}).apply(l8); + + decoder = tf.model({inputs: decoder_input, outputs: decoded, name: "decoder"}); + console.log(`Decoder Summary: ${decoder.summary()}`); + + + const auto = tf.input({shape: [W*W]}); + encoded = encoder.apply(auto); + decoded = decoder.apply(encoded); + + const autoencoder = tf.model({inputs: auto, outputs: decoded, name: "autoencoder"}); + console.log(`Autoencoder Summary: ${autoencoder.summary()}`); + + + const learningRate = 0.0001; + const optimizer = tf.train.adam(learningRate, 0.000001); // adam(learning_rate, decay) + autoencoder.compile({ optimizer: "adam", loss: "meanSquaredError", }); - return { decoderLayers, autoencoder }; + decoder.compile({ + optimizer: "adam", + loss: "meanSquaredError" + }); + return { decoder, autoencoder}; } async function trainModel(autoencoder, x_train, epochs) { @@ -171,6 +127,7 @@ async function trainModel(autoencoder, x_train, epochs) { batch_size: 32, shuffle: true, verbose: true, + validation_split: 0.1 }); } @@ -179,7 +136,7 @@ async function loadImages(total) { for (let i = 0; i < total; i++) { const num = numeral(i).format("0000"); const img = await Jimp.read( - `AutoEncoder_TrainingData/data/square${num}.png` + `data/square${num}.png` ); let rawData = [];