Skip to content

feat: Change the model build, create separated encoder and decoder #8

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: Change the model build, create separated encoder and decoder #8

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
151 changes: 54 additions & 97 deletions index.js
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
console.log("Hello Autoencoder 🚂");

import * as tf from "@tensorflow/tfjs-node";
import * as tf from "@tensorflow/tfjs-node-gpu";
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do I need to install CUDA or anything like that to run tf.js with node GPU?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes, TensorFlow GPU only supports cuda.

From my Python experience it needs an extremely specific version of cuda, but for JS it worked with out of the box 11.2 version. I only used the GPU version to speed up training, the CPU version works fine.

// import canvas from "canvas";
// const { loadImage } = canvas;
import Jimp from "jimp";
Expand All @@ -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/");
}

Expand Down Expand Up @@ -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",
Expand All @@ -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"});
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I haven't used tf.model before, is it required for this methodology, or can I keep tf.sequential? Is tf.model preferred?

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

from https://js.tensorflow.org/api/latest/#sequential

The key difference between tf.model() and tf.sequential() is that tf.model() is more generic, supporting an arbitrary graph (without cycles) of layers. tf.sequential() is less generic and supports only a linear stack of layers.

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({
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was wondering when I built this if the decoder has to be "compiled" with an optimizer and loss given I'm not training it, only using it for inference?

optimizer: "adam",
loss: "meanSquaredError"
});
return { decoder, autoencoder};
}

async function trainModel(autoencoder, x_train, epochs) {
Expand All @@ -171,6 +127,7 @@ async function trainModel(autoencoder, x_train, epochs) {
batch_size: 32,
shuffle: true,
verbose: true,
validation_split: 0.1
});
}

Expand All @@ -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`
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there a reason you changed this? I like being able to pull the training data directly from the Processing sketch.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

no specific reason, i didn't have processing installed, so i created a Python script to generate the images, i only changed to the output of that

);

let rawData = [];
Expand Down