构建和训练模型的脚本如下:
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。但这在模型对象上不可用。那么我该如何进行?
答案 0 :(得分: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