2.1

Dataset

The dataset is retrieved from the Kaggle repository named kermany2018. It has three folders (train, test, validation), each with subfolders for four categories of images (NORMAL, CNV, DME, DRUSEN). The dataset consists of 83,484 training images, 41,741 validation images, and 968 test images.

2.2

Vision Transformer

The vision transformer model was proposed by Dosovitskiy in 2020 [12]. The model contains three part: token generation, transformer encoder and classification layer. Fig. 2 shows the structure of vision transformer. The input image is separated into several patches called tokens, firstly. Then the linear embedding is performed on each token to consider the position information. A new token called the class token is added to the sequence of the tokens which plays the role of feature expression for each scene. Next the token sequence is sent into the encoder to consider the interaction of the tokens using multi-head self-attention mechanism. Finally, the output class token is classified into different scene categories by a MLP layer. The vision transformer model directly considers the contextual information and the spatial distribution of the objects contained in the image.

Fig. 2

The structure of vision transformer.

2.3

Symmetric cross-entropy loss function

The symmetrical cross-entropy loss function can minimize the effect of noise on the training set and prevent overfitting [13]. Therefore, we used the symmetric cross-entropy loss function as the loss function of the model to reduce the impact of noise during model training. The definition of the symmetric cross-entropy loss function is:

I_ce here is the cross-entropy loss function and l_rce is the reverse cross-entropy function. Their definitions are shown as follows:

Where, q(k|x) is the ground truth class distribution conditioned on sample x, whilst p(k|x) is the predicted distribution over labels by the classifier.

3.1

Experimental environment

The vision transformer model is simulated using PyCharm software installed on 64-bit windows operating system, with Intel (R) Xeon (R) W-2245 CPU, NVIDIA GeForce RTX 3090 and 32GB memory. Further, we conducted various experiments to evaluate its performance in predicting the candidate retinal conditions (healthy, CNV, DME, and Drusen) using OCT scans.

3.2

Evaluation standard

We adopt the evaluation metrics of classification accuracy, sensitivity and specificity to evaluate the performance of the vision transformer model. Accuracy is defined as the ratio of correctly recognized positive instances and negative instances to total instances. Sensitivity is the detection rate of positive instances and specificity is the detection rate of negative instances. Their definitions are as follows:

where, TP is the true positive instance, TN is the true negative instance, FP is the false positive instance, and FN is the false negative instance. In addition, we also used the confusion matrix, the comparison matrix between the predicted labels and the real labels, to analyze the performance of each category of OCT image.

3.3

Results

We trained the model 30 times to obtain the best results and saved the parameters of the model with the highest accuracy of the model on the validation set. To verify the reliability of the model, accuracy and loss of iterations are shown in Fig. 3. The best classification accuracy rate is 96.57% and the loss is 0.0929.

Fig. 3

The loss changes (a) and the accuracy changes (b) of the validation set.

From the above loss changes and accuracy changes, it can be determined that the transformer model is normally fitted without overfitting. The results were summarized in the form of a confusion matrix according to real labels. The columns of the matrix represent true labels, and the rows represent predicted labels. Confusion matrix for diagnosing retinal diseases is shown in Fig. 4.

Fig. 4

The confusion matrix of the test set.

As shown in Fig. 5, the proposed method can achieve classification accuracy of 95.76%, sensitivity of 95.77% and specificity of 98.59% on the test set. In this study, the model presented in this paper was compared with the traditional CNN to verify the effectiveness of the method proposed. We choose EfficentNet, Resnet50, and VGG16 as the comparison models. The experimental steps and hyperparameters of these models are consistent with the vision transformer. From the graphical comparison of the validation set presented in Fig. 5, it is shown that the validation loss of vision transformer for classifying retinal diseases is 0.0929, which is less than that of Vgg16 by 0.0272 and ResNet50 by 0.0913. Compared with Vgg16 and ResNet50, the accuracy is increased by 3.72% and 8.45%, respectively.

Fig. 5

The accuracy and loss comparison of vision transformer with the other CNN models.

The accuracy comparison of validation set and test set between vision transformer and five CNN models are shown in Table 1.

Table 1.

Model comparison of accuracy

As we can see from Table 1, in the CNN image classification models, VGG16 has the highest classification accuracy of 91.42% in test set which is still less than the classification accuracy of vision transformer. Through the attention mechanism, vision transformer can concentrate on the image region related to the semantics of the classification target, thereby obtaining higher accuracy. Considering the recognition accuracy, vision transformer is superior to CNN models in diagnosis of retinal disease.