如何在 Tensorflow js 中将 Tensor 格式更改为通道优先?

How to change Tensor format to channels first in Tensorflow js?

我是计算机视觉模型结构的新手,我正在使用适用于 Node JS 的 Tensorflow @tensorflow/tfjs-node 使一些模型检测到一些对象。对于 Mobilenet 和 Resnet SSD,模型使用 Channels Last 格式,因此当我使用 tf.node.decodeImage 创建 Tensor 时,格式默认为 Channels Last,例如 3 通道的 shape: [1, 1200, 1200, 3] ,而且预测数据很好,能够识别物体。

但是 Pytorch 的模型,转换为 ONNX,然后转换为 Protobuf PB 格式,saved_model.pb 具有 Channels First 格式,如 shape: [1, 3, 1200, 1200]

现在我需要使用 Channels First 格式从图像创建 Tensor。我发现许多创建 conv1d, conv2d 指定格式 dataFormat='channelsFirst' 的例子。但我不知道如何将它应用于图像数据。这是 API https://js.tensorflow.org/api/latest/#layers.conv2d .

这是张量代码:

const tf = require('@tensorflow/tfjs-node');
let imgTensor = tf.node.decodeImage(new Uint8Array(subBuffer), 3);
imgTensor = imgTensor.cast('float32').div(255);
imgTensor = imgTensor.expandDims(0); // to add the most left axis of size 1
console.log('tensor', imgTensor);

这给了我一个最后一个通道的形状,它与第一个通道的模型形状不兼容:

tensor Tensor {
  kept: false,
  isDisposedInternal: false,
  shape: [ 1, 1200, 1200, 3 ],
  dtype: 'float32',
  size: 4320000,
  strides: [ 4320000, 3600, 3 ],
  dataId: {},
  id: 7,
  rankType: '4',
  scopeId: 4
}

我知道 tf.shape,但它在没有先转换为通道的情况下进行了重塑,结果在预测结果中似乎没有用。不知道我错过了什么。

你可以使用这样的东西:

const nchw = tf.transpose(nhwc, [0, 3, 1, 2]);