如何在 tensorflow.js 中使用自定义模型对图像进行分类？

Question

构建和训练模型的脚本如下所示：

const model = tf.sequential();

model.add(tf.layers.conv2d({
  inputShape: [160, 200, 3],
  filters: 32,
  kernelSize: 3,
  activation: 'relu',
}));

model.add(tf.layers.flatten());

model.add(tf.layers.dense({units: labels.length, activation: 'softmax'}));  

model.compile({
  optimizer: 'sgd',
  loss: 'categoricalCrossentropy',
  metrics: ['accuracy']
});    

const info = await model.fitDataset(ds, {epochs: 5});
console.log('Accuracy:', info.history.acc);

console.log('Saving...');
await model.save('file://' + MODEL_PATH);
console.log('Saved model');

ds 由图像和标签组成。对于 100 张图片，我得到以下结果：

4478ms 135692us/step - acc=0.109 loss=14.37

它生成了一个 20 MB 的 weights.bin 文件...

坦率地说，我不知道这是否好用，因为我不知道如何使用它对新图像进行分类。

我知道如何加载模型：

const model = await tf.loadLayersModel('file://' + MODEL_PATH + '/model.json');

但就是这样。

mobilenet 有一个 .classify 方法，我可以将图像传递给该方法，它会输出预测的 laabel。但这在模型对象上不可用。那么我该如何继续？

Answer 1

训练您的模型后，为了对新图像进行分类，将使用predict 方法。它将返回给定输入图像的每个标签的概率。

output = model.predict(image) // It can be a tensor of one image or a batch of many 
// output is a tensor that contain the probability for each label
images.argMax([-1]) // contain the label of high probability for each input