diteix / lil-square-of-eight Goto Github PK

const model = tf.sequential();

model.add(tf.layers.dense({ units: 16, inputShape: [8] }));
model.add(tf.layers.dense({ units: 16, inputShape: [16], activation: 'sigmoid' }));
model.add(tf.layers.dense({ units: 1, inputShape: [16] }));
model.compile({ optimizer: tf.train.adam(0.1), loss: 'meanSquaredError', lr: 0.3 });

const xs1 = tf.tensor2d([
                [0, 1, 1, 1, 2, 0, 3, 0], [0, 1, 1, 0, 2, 1, 3, 0], [0, 1, 1, 0, 2, 0, 3, 1], [0, 0, 1, 1, 2, 1, 3, 0], [0, 0, 1, 1, 2, 0, 3, 1], [0, 0, 1, 0, 2, 1, 3, 1],
                [0, 1, 1, 1, 2, 1, 3, 1],
                [0, 0, 1, 0, 2, 0, 3, 0], [0, 1, 1, 0, 2, 0, 3, 0], [0, 0, 1, 1, 2, 0, 3, 0], [0, 0, 1, 0, 2, 1, 3, 0], [0, 0, 1, 0, 2, 0, 3, 1],
                [0, 0, 1, 1, 2, 1, 3, 1], [0, 1, 1, 0, 2, 1, 3, 1], [0, 1, 1, 1, 2, 0, 3, 1], [0, 1, 1, 1, 2, 1, 3, 0]]);

await model.fit(xs1, tf.tensor2d([[0], [0], [0], [0], [0], [0], [0], [1], [1], [1], [1], [1], [2], [2], [2], [2]]), {epochs: 550});

console.log(model.predict(tf.tensor2d([[0, 1, 1, 0, 2, 0, 3, 0]])));