javascripttensorflowmachine-learningtensorflow.js

Best model type in TensorFlow.js for Color Prediction?

I was creating a color predictor when I realized a problem arise. I got the model to work successfully, but the predictions are always within the same median range of about 2.5 to about 5.5. The model is supposed to output a 0 thru 8 corresponding to each color, and I have an even amount of data point for each color for training. Is there a better model I can use so that it will predict something to be 0 or 7? I'm assuming it won't because it thinks they are some kind of outliers.

Here's my model

const model = tf.sequential();

const hidden = tf.layers.dense({
  units: 3,
  inputShape: [3] //Each input has 3 values r, g, and b
});
const output = tf.layers.dense({
  units: 1 //only one output (the color that corresponds to the rgb values
    });
model.add(hidden);
model.add(output);

model.compile({
  activation: 'sigmoid',
  loss: "meanSquaredError",
  optimizer: tf.train.sgd(0.005)
});

Is this a good model for my problem?

Solution

  • The model is lacking non-linearity, for there is no activation functions. Given a rgb input, the model should predict the most likely color in 8 possible values. It is a classification problem. The model as defined in the question is doing a regression, i.e it is trying to predict a numerical value given the input.

    For classification problem, the last layer should predict probabilities. softmax activation function is mostly used for the last layer in that case. The loss function should be categoricalCrossentropy or binaryCrossEntropy (if there were only two colors to predict).

    Consider the following model predicting 3 classes of colors: red, green and blue

    const model = tf.sequential();
model.add(tf.layers.dense({units: 10, inputShape: [3], activation: 'sigmoid' }));
model.add(tf.layers.dense({units: 10, activation: 'sigmoid' }));
model.add(tf.layers.dense({units: 3, activation: 'softmax' }));

model.compile({ loss: 'categoricalCrossentropy', optimizer: 'adam' });

const xs = tf.tensor([
  [255, 23, 34],
  [255, 23, 43],
  [12, 255, 56],
  [13, 255, 56],
  [12, 23, 255],
  [12, 56, 255]
]);

// Labels
const label = ['red', 'red', 'green', 'green', 'blue', 'blue']
const setLabel = Array.from(new Set(label))
const ys = tf.oneHot(tf.tensor1d(label.map((a) => setLabel.findIndex(e => e === a)), 'int32'), 3)

// Train the model using the data.
  model.fit(xs, ys, {epochs: 100}).then((loss) => {
  const t = model.predict(xs);
  pred = t.argMax(1).dataSync(); // get the class of highest probability
  labelsPred = Array.from(pred).map(e => setLabel[e])
  console.log(labelsPred)
}).catch((e) => {
  console.log(e.message);
})
    <html>
  <head>
    <!-- Load TensorFlow.js -->
    <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@0.13.3/dist/tf.min.js"> </script>
  </head>

  <body>
  </body>
</html>