node.jsmachine-learningkerastensorflow.jstfjs-node

How to train tensorflow.js on medical data


I'm doing a POC script with tfjs and some (fake) medical data (breast cancer). The data looks like this:

[206, 293, 140, 126, 117, 27, 35, 152, 239, 79] with results (ys) is [1] where 1 is malignant and 0 is benign.

The script appears to train, but the accuracy/loss never changes, and I get the same result regardless of the data. I've already verified the data/format. Script below:

  const transformedData = _.shuffle(data).map(util.transformRow);

// Define the model.
const model = tf.sequential();
// Set up the network layers
model.add(tf.layers.dense({units: 10, activation: 'relu', inputShape: [10]}));
model.add(tf.layers.dense({units: 150, activation: 'relu'}));
model.add(tf.layers.dense({units: 250, activation: 'relu'}));
model.add(tf.layers.dense({units: 250, activation: 'relu'}));
model.add(tf.layers.dense({units: 250, activation: 'relu'}));
model.add(tf.layers.dense({units: 1, activation: 'softmax', outputShape: [1]}));
// Define the optimizer
const optimizer = tf.train.adam(LEARNING_RATE);
// Init the model
model.compile({
    optimizer: optimizer,
    loss: 'meanSquaredError',
    metrics: ['accuracy'],
});

  const ys = transformedData.map(d => [d.ys]);
  const xs = transformedData.map(d => d.xs);

  let xTrain = tf.tensor2d(xs.slice(0,500), [xs.slice(0,500).length, 10]); // [[123,234,345...], [...]...]
  let yTrain = tf.tensor2d(ys.slice(0,500), [ys.slice(0,500).length, 1]); // [[1], [0]...]

  console.log('ready to start training model');
  const history = await model.fit(xTrain, yTrain, {
    epochs: EPOCHS,
    validationData: [xTrain, yTrain],
  }) 

To be clear, I don't care about super accurate results or optimization right now, I just want the script to actually train the model.


Solution

  • Two issues:

    1. Since you are in a binary classification setting, you should use loss: 'binaryCrossentropy' (MSE is for regression problems).

    2. In such a setting, a softmax activation of the output layer does not make any sense; change it to sigmoid.

    Also, you don't share the actual value of your LEARNING_RATE; try removing the parameter altogether - Adam is known to usually work (very) well with its default learning rate.