js实现lstm
LSTM 简介
LSTM(Long Short-Term Memory)是一种特殊的循环神经网络(RNN),用于处理时间序列数据或序列数据的长期依赖问题。在 JavaScript 中,可以通过 TensorFlow.js 实现 LSTM 模型。
安装 TensorFlow.js
在项目中引入 TensorFlow.js:
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@latest"></script>
创建 LSTM 模型
使用 TensorFlow.js 的 tf.layers.lstm 构建 LSTM 层:
const model = tf.sequential();
model.add(tf.layers.lstm({
units: 50, // LSTM 单元数量
inputShape: [10, 1], // 输入形状:[时间步长, 特征维度]
returnSequences: false // 是否返回序列
}));
model.add(tf.layers.dense({units: 1})); // 输出层
编译模型
配置优化器和损失函数:
model.compile({
optimizer: tf.train.adam(0.01),
loss: 'meanSquaredError'
});
准备数据
将输入数据转换为张量:
const xs = tf.tensor2d([/* 输入数据 */], [样本数, 时间步长]);
const ys = tf.tensor2d([/* 标签数据 */], [样本数, 1]);
训练模型
调用 fit 方法进行训练:
await model.fit(xs, ys, {
epochs: 100,
batchSize: 32,
callbacks: {
onEpochEnd: (epoch, logs) => console.log(`Epoch ${epoch}: loss = ${logs.loss}`)
}
});
预测
使用训练好的模型进行预测:
const input = tf.tensor2d([/* 新数据 */], [1, 10]);
const prediction = model.predict(input);
prediction.print();
注意事项
- LSTM 对输入数据的形状敏感,需确保输入张量的形状与模型定义一致。
- 训练数据需进行归一化处理以提高模型性能。
- 调整
units和epochs等超参数以优化结果。
完整示例
以下是一个完整的 LSTM 时间序列预测示例:

// 生成示例数据
function generateData(numPoints) {
const data = [];
for (let i = 0; i < numPoints; i++) {
data.push(Math.sin(i * 0.1));
}
return data;
}
// 准备训练数据
const data = generateData(100);
const xs = [];
const ys = [];
for (let i = 0; i < data.length - 10; i++) {
xs.push(data.slice(i, i + 10));
ys.push(data[i + 10]);
}
// 转换为张量
const xsTensor = tf.tensor2d(xs, [xs.length, 10]);
const ysTensor = tf.tensor2d(ys, [ys.length, 1]);
// 创建模型
const model = tf.sequential();
model.add(tf.layers.lstm({units: 50, inputShape: [10, 1]}));
model.add(tf.layers.dense({units: 1}));
// 编译与训练
model.compile({optimizer: 'adam', loss: 'meanSquaredError'});
model.fit(xsTensor, ysTensor, {epochs: 20}).then(() => {
// 预测
const testInput = tf.tensor2d([data.slice(90, 100)], [1, 10]);
model.predict(testInput).print();
});
通过以上步骤,可以在 JavaScript 中实现 LSTM 模型并用于时间序列预测任务。






