plunkerでtensorflow.js その15

#概要

plunkerでtensorflow.jsやってみた。
argMax使ってみた。

#サンプルコード

tf.setBackend('cpu');
let index2alphabet = [];
let alphabet2vector = {};
let trainXArray = [], 
  trainYArray = [];
for (let i = 0; i < 26; i++) 
{
  index2alphabet[i] = String.fromCharCode('a'.charCodeAt(0) + i);
}
for (let i = 0; i < 26; i++) 
{
  alphabet2vector[index2alphabet[i]] = Array(26).fill(0);
  alphabet2vector[index2alphabet[i]][i] = 1;
}
trainXArray.push([alphabet2vector['y'], alphabet2vector['z']]);
trainYArray.push(alphabet2vector['a']);
trainXArray.push([alphabet2vector['z'], alphabet2vector['a']]);
trainYArray.push(alphabet2vector['b']);
for (let i = 0; i < 26 - 2; i++)
{
  trainXArray.push([alphabet2vector[index2alphabet[i]], alphabet2vector[index2alphabet[i + 1]]]);
  trainYArray.push(alphabet2vector[index2alphabet[i + 2]]);
}
const trainX = tf.tensor3d(trainXArray);
const trainY = tf.tensor2d(trainYArray);
const model = tf.sequential();
model.add(tf.layers.lstm({ 
  units: 64,
  returnSequences: false,
  inputShape: [2, 26]
}));
model.add(tf.layers.dense({
  units: 64, 
  activation: "relu"
}));
model.add(tf.layers.dense({
  units: 26, 
  activation: "softmax"
}));
model.compile({
  optimizer: tf.train.adam(),
  loss: tf.losses.softmaxCrossEntropy
});
model.fit(trainX, trainY, {
  epochs: 160,
  verbose: 0,
  validationData: [trainX, trainY],
}).then(() => {
  const testX = tf.tensor3d(trainXArray);
  const predicts = model.predict(testX);
  const predictsArray = predicts.argMax().arraySync();
  for (let i = 0; i < predictsArray.length; i++) 
  {
    let predictAlphabet = String.fromCharCode('a'.charCodeAt(0) + predictsArray[i]);
    document.write(predictAlphabet);
  }
  alert("ok");
});

#成果物

以上。