| 20 | const sz_model = require('./strike_zone'); |
| 21 | |
| 22 | async function run(epochCount, savePath) { |
| 23 | sz_model.model.summary(); |
| 24 | await sz_model.model.fitDataset(sz_model.trainingData, { |
| 25 | epochs: epochCount, |
| 26 | callbacks: { |
| 27 | onEpochEnd: async (epoch, logs) => { |
| 28 | console.log(`Epoch: ${epoch} - loss: ${logs.loss.toFixed(3)}`); |
| 29 | } |
| 30 | } |
| 31 | }); |
| 32 | |
| 33 | // Eval against test data: |
| 34 | await sz_model.testValidationData.forEachAsync(data => { |
| 35 | const evalOutput = |
| 36 | sz_model.model.evaluate(data.xs, data.ys, sz_model.TEST_DATA_LENGTH); |
| 37 | |
| 38 | console.log( |
| 39 | `\nEvaluation result:\n` + |
| 40 | ` Loss = ${evalOutput[0].dataSync()[0].toFixed(3)}; ` + |
| 41 | `Accuracy = ${evalOutput[1].dataSync()[0].toFixed(3)}`); |
| 42 | }); |
| 43 | |
| 44 | if (savePath !== null) { |
| 45 | await sz_model.model.save(`file://${savePath}`); |
| 46 | console.log(`Saved model to path: ${savePath}`); |
| 47 | } |
| 48 | } |
| 49 | |
| 50 | const parser = new argparse.ArgumentParser( |
| 51 | {description: 'TensorFlow.js Strike Zone Training Example', addHelp: true}); |