36 lines
1.4 KiB
TypeScript
36 lines
1.4 KiB
TypeScript
import * as tf from "@tensorflow/tfjs-node";
|
|
import { readFile } from "node:fs/promises";
|
|
import { existsSync } from "node:fs";
|
|
import { join } from "node:path";
|
|
|
|
async function train() {
|
|
const savePath = process.argv[2] || 'sample/model';
|
|
const data = JSON.parse(await readFile(join(process.argv[3] || 'sample/dataset', 'dataset.json'), "utf-8"));
|
|
const X = tf.tensor2d(data.map((d: any) => d.x.map((v: any, i: number) => v / (i == 0 ? 20 : i == 1 ? 200 : 1))));
|
|
const y = tf.tensor2d(data.map((d: any) => [d.y]));
|
|
|
|
let model: tf.LayersModel;
|
|
const modelJsonPath = join(savePath, "model.json");
|
|
|
|
if (existsSync(modelJsonPath)) {
|
|
console.log("기존 모델을 불러와 추가 학습을 진행합니다.");
|
|
model = await tf.loadLayersModel(`file://${modelJsonPath}`);
|
|
} else {
|
|
console.log("새 모델을 생성합니다.");
|
|
model = tf.sequential({
|
|
layers: [
|
|
tf.layers.dense({ units: 32, activation: 'relu', inputShape: [5] }),
|
|
tf.layers.dense({ units: 16, activation: 'relu' }),
|
|
tf.layers.dense({ units: 1 })
|
|
]
|
|
});
|
|
}
|
|
|
|
model.compile({ optimizer: 'adam', loss: 'meanSquaredError' });
|
|
const epochs = parseInt(process.argv[4]) || 50;
|
|
await model.fit(X, y, { epochs: epochs, verbose: 0 });
|
|
await model.save(`file://${savePath}`);
|
|
console.log(`학습 완료: 모델이 ${savePath}에 저장되었습니다.`);
|
|
}
|
|
train().catch(console.error);
|