bangoc123 / mlp-mixer

Implementation for paper MLP-Mixer: An all-MLP Architecture for Vision

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Which class is predicted?

abderhasan opened this issue · comments

Thanks so much for the nice implementation.

I just have a question about the predicted class.

For instance, I have two classes in two directories: cat and dog. After training the model, I tried to predict the image you had in the test directory, and got the following:

Output Softmax: [[0. 1.]]
This image belongs to class: 1

Which class does 1 refer to in this case?

Thanks.

Hi @abderhasan.

You can use print('--->', train_ds.class_names) in train.py to list an array of the classes.

The index of a label in that array is corresponding to its class name.