TensorFlow JS大数据集无法内存加载的异步训练解决方案及示例求助

我在TensorFlow JS中遇到了数据集过大无法一次性载入内存的问题，需要对全部数据条目进行训练。我的数据来自MongoDB实例，需异步加载。

我尝试使用生成器函数，但尚未成功实现异步生成器，同时考虑是否可以分批拟合模型。希望能获得通过分批或数据库游标异步加载数据进行训练的最小示例。

遇到的TypeScript错误示例 #

1. 返回Promise的生成器报错：

const generate = function* () {
    yield new Promise(() => {});
};

tf.data.generator(generate);

错误信息：

Argument of type '() => Generator<Promise<unknown>, void, unknown>' is not assignable to parameter of type '() => Iterator<TensorContainer, any, undefined> | Promise<Iterator<TensorContainer, any, undefined>>'.

2. 异步生成器报错：

tf.data.generator(async function* () {})

错误信息：

Argument of type '() => AsyncGenerator<any, void, unknown>' is not assignable to parameter of type '() => Iterator<TensorContainer, any, undefined> | Promise<Iterator<TensorContainer, any, undefined>>'.

核心原因 #

tf.data.generator仅支持同步迭代器或返回同步迭代器的Promise，不直接支持异步生成器或yield Promise的生成器。要处理异步数据源，推荐用tf.data.Dataset.fromAsync或手动分批加载逻辑。

方案1：用tf.data.fromAsync结合MongoDB游标 #

tf.data.fromAsync可以接收一个返回异步迭代器的函数，完美适配MongoDB的游标（MongoDB游标本身就是异步迭代器）。

最小示例代码

import * as tf from '@tensorflow/tfjs';
import { MongoClient } from 'mongodb';

// 1. 连接MongoDB并获取数据游标
async function getMongoCursor() {
  const client = await MongoClient.connect('mongodb://localhost:27017');
  const db = client.db('your-db-name');
  // 假设集合中每个文档有features（输入特征）和label（标签）字段
  return db.collection('your-collection').find({}, { projection: { features: 1, label: 1 } });
}

// 2. 用tf.data.fromAsync创建异步数据集
async function createDataset() {
  const cursor = await getMongoCursor();
  // 将MongoDB文档转换为TensorFlow可处理的格式
  return tf.data.fromAsync(async function* () {
    for await (const doc of cursor) {
      // 将文档中的features和label转为张量
      const features = tf.tensor(doc.features);
      const label = tf.tensor(doc.label);
      yield { xs: features, ys: label };
    }
  }).batch(32); // 设置每批大小
}

// 3. 定义模型并训练
async function trainModel() {
  const dataset = await createDataset();
  
  // 简单的示例模型
  const model = tf.sequential({
    layers: [
      tf.layers.dense({ inputShape: [10], units: 32, activation: 'relu' }),
      tf.layers.dense({ units: 1, activation: 'sigmoid' })
    ]
  });
  
  model.compile({ optimizer: 'adam', loss: 'binaryCrossentropy', metrics: ['accuracy'] });
  
  // 用数据集训练模型
  await model.fitDataset(dataset, {
    epochs: 5,
    callbacks: {
      onEpochEnd: (epoch, logs) => {
        console.log(`Epoch ${epoch + 1}: loss = ${logs.loss}, accuracy = ${logs.acc}`);
      }
    }
  });
  
  // 关闭MongoDB连接
  const client = await MongoClient.connect('mongodb://localhost:27017');
  await client.close();
}

trainModel();

方案2：手动分批加载并训练 #

如果不想用tf.data API，可以手动从MongoDB分批拉取数据，逐批训练模型。

最小示例代码

import * as tf from '@tensorflow/tfjs';
import { MongoClient } from 'mongodb';

async function trainBatchByBatch() {
  const client = await MongoClient.connect('mongodb://localhost:27017');
  const db = client.db('your-db-name');
  const collection = db.collection('your-collection');
  
  const batchSize = 32;
  const totalDocs = await collection.countDocuments();
  const totalBatches = Math.ceil(totalDocs / batchSize);
  
  // 定义模型
  const model = tf.sequential({
    layers: [
      tf.layers.dense({ inputShape: [10], units: 32, activation: 'relu' }),
      tf.layers.dense({ units: 1, activation: 'sigmoid' })
    ]
  });
  
  model.compile({ optimizer: 'adam', loss: 'binaryCrossentropy', metrics: ['accuracy'] });
  
  // 逐批加载并训练
  for (let batchIdx = 0; batchIdx < totalBatches; batchIdx++) {
    // 从MongoDB拉取当前批次的数据
    const docs = await collection.find({}, { projection: { features: 1, label: 1 } })
      .skip(batchIdx * batchSize)
      .limit(batchSize)
      .toArray();
    
    // 转换为张量
    const features = tf.tensor2d(docs.map(d => d.features));
    const labels = tf.tensor2d(docs.map(d => d.label), [docs.length, 1]);
    
    // 训练当前批次
    const logs = await model.fit(features, labels, {
      batchSize: batchSize,
      epochs: 1,
      verbose: 0
    });
    
    console.log(`Batch ${batchIdx + 1}/${totalBatches}: loss = ${logs.history.loss[0]}, accuracy = ${logs.history.acc[0]}`);
    
    // 手动清理张量，避免内存泄漏
    features.dispose();
    labels.dispose();
  }
  
  await client.close();
}

trainBatchByBatch();

注意事项 #

• 确保MongoDB查询只返回需要的字段（用projection），减少数据传输量
• 训练时注意清理不再使用的张量，避免内存泄漏
• 如果数据集极大，可以考虑添加shuffle操作（方案1中可以在.batch()前加.shuffle(bufferSize)）
• 方案1中的tf.data.fromAsync会自动处理异步迭代的完成逻辑，但训练结束后最好手动关闭MongoDB连接

内容的提问来源于stack exchange，提问作者Sir hennihau