How to replace the classes in a pretrained network's output classification layer?

Question

Dominik Mattioli 2021-6-10

0
链接

此问题的直接链接

https://ww2.mathworks.cn/matlabcentral/answers/853415-how-to-replace-the-classes-in-a-pretrained-network-s-output-classification-layer

回答： Anshika Chaurasia 2021-6-14

I've followed the MATLAB tutorial for training an image classification model (link below) and it works well for my 2 classes, but when I export it and then import it later on, it's basing predictions off of the original pretrained network's 1000 classes.

exportONNXNetwork( net, 'my_class_model.onnx' ), 'NetworkName', 'myModel' );
%... saved, then import it later on when needed:
fn	= 'my_class_model.onnx';
classes	= ["Class1" "Class2"];
mynet	= importONNXNetwork( fn, 'OutputLayerType', 'classification', 'Classes', classes );
% Error results:
% Error using nnet.internal.cnn.onnx.importONNXNetwork>iCheckOutputLayerTypeAndClasses (line 116)
% Expected the number of classes specified to match the number of classes in the classification output
% layer: 1000. Instead it was 2.

If you remove the 'Classes' input, there is no error. However, the predictions are then given for the original 1,000 classes, not my two.

Do I need to replace the classification layer before training? Or is there an exporting parameter for this?

Link to tutorial:

Image Category Classification Using Deep Learning - MATLAB & Simulink (mathworks.com)

0 个评论
显示 -2更早的评论隐藏 -2更早的评论

请先登录，再进行评论。

请先登录，再回答此问题。

Answer 1

Anshika Chaurasia 2021-6-14

0
链接

此回答的直接链接

https://ww2.mathworks.cn/matlabcentral/answers/853415-how-to-replace-the-classes-in-a-pretrained-network-s-output-classification-layer#answer_724570

Hi,

As correctly mentioned you need to replace the classification layer and fully connected layer in net before training, so that network classify image for two classes.

You can refer to the following documentation and video for transfer learning in pretrained network:

Deep Learning with MATLAB: Transfer Learning in 10 Lines of MATLAB Code Video - MATLAB (mathworks.com)

Transfer Learning Using AlexNet - MATLAB & Simulink (mathworks.com)

Hope it helps!