当前位置：首页 > JavaScript

js实现lstm

2026-02-01 00:38:12JavaScript

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 层：

js实现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'
});

准备数据

将输入数据转换为张量：

js实现lstm

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 模型并用于时间序列预测任务。