使用TensorFlow.js无法提取输出概率数组的技术问题求助

问题原因 #

编译报错Property 'data' does not exist on type 'Tensor<Rank> | Tensor<Rank>[]'，是因为TypeScript中net.predict()的返回类型被定义为Tensor<Rank> | Tensor<Rank>[]（兼容多输出模型），即使你的模型是单输出的顺序模型，TS也无法自动推断返回单个张量，导致无法直接调用data()方法。

解决方案 #

方案1：类型断言修复基础调用 #

直接将predict的返回值断言为单个tf.Tensor，即可正常调用data()：

// 替换原错误代码段
const outputTensor = net.predict(input) as tf.Tensor;
const output = await outputTensor.data();  
const predictedPortfolio = Object.keys(portfolios)[output.indexOf(Math.max(...output))];
return predictedPortfolio;

方案2：优化内存管理与计算效率 #

TensorFlow.js的张量需要手动管理内存，同时用TFJS内置方法替代原生JS数组操作更高效：

// 替换原预测逻辑部分
const predictedIndex = tf.tidy(() => {
  // 断言为单张量
  const outputTensor = net.predict(input) as tf.Tensor;
  // 直接在张量上计算最大值索引（axis=1对应每行的最大值位置）
  return outputTensor.argMax(1).dataSync()[0];
});

const predictedPortfolio = Object.keys(portfolios)[predictedIndex];
return predictedPortfolio;

修改后的完整代码 #

import * as tf from '@tensorflow/tfjs';

interface FinancialInformation {
  age: number;
  riskTolerance: number;
  currentNetWorth: number;
  annualIncome: number;
  debt: number;
}

interface Portfolios {
  [key: string]: string[];
}

// 定义投资组合
const portfolios: Portfolios = {
  conservative: ['bonds', 'real estate'],
  balanced: ['stocks', 'bonds', 'real estate', 'commodities'],
  aggressive: ['stocks', 'commodities', 'crypto'],
};

export async function generateSuggestion(financialInfo: FinancialInformation): Promise<string> {
  // 定义神经网络
  const net = tf.sequential({
    layers: [
      tf.layers.dense({ inputShape: [5], units: 10, activation: 'sigmoid' }),
      tf.layers.dense({ units: 10, activation: 'sigmoid' }),
      tf.layers.dense({ units: 3, activation: 'softmax' }),
    ],
  });

  // 训练神经网络
  const trainingData = tf.tensor2d([
    [25, 2, 100000, 60000, 0],
    [30, 4, 150000, 80000, 20000],
    [40, 6, 200000, 100000, 50000],
    [50, 8, 300000, 120000, 100000],
    [60, 10, 400000, 150000, 150000],
  ]);

  const outputData = tf.tensor2d([
    [1, 0, 0],
    [0, 1, 0],
    [0, 0, 1],
    [0, 0, 1],
    [0, 0, 1],
  ]);

  const options = {
    epochs: 500,
    learningRate: 0.3,
  };

  net.compile({ optimizer: tf.train.adam(options.learningRate), loss: 'categoricalCrossentropy' });

  await net.fit(trainingData, outputData, options);

  // 构造输入张量
  const input = tf.tensor2d([
    [financialInfo.age, financialInfo.riskTolerance, financialInfo.currentNetWorth, financialInfo.annualIncome, financialInfo.debt],
  ]);

  // 优化后的预测逻辑
  const predictedIndex = tf.tidy(() => {
    const outputTensor = net.predict(input) as tf.Tensor;
    return outputTensor.argMax(1).dataSync()[0];
  });

  const predictedPortfolio = Object.keys(portfolios)[predictedIndex];
  // 清理训练和输入张量，避免内存泄漏
  tf.dispose([trainingData, outputData, input]);
  return predictedPortfolio;
}

补充说明：

• tf.tidy()会自动清理其回调函数内创建的所有临时张量，仅返回指定结果，有效避免内存泄漏
• argMax(1)比原生Math.max(...output)更适合张量操作，批量数据场景下性能优势明显
• 训练完成后手动清理训练数据和输入张量，避免长期运行内存堆积

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