After lots of method still cannot handle Multi-class classification with class imbalance dataset

Keahi: <p>Hi, everyone! I’m currently working on a medical image detection task based on chest X-ray images as my undergraduate final project, in which the dataset is composed of Pneumonia with 4273 images, Normal class with 1583 images, COVID-19 with 576 images, and I used ViT(replicate by myself from scratch) as the model architecture in this project. Above is the basic information of this project. But as the title said, I found that my model cannot learn anything from the dataset through these metrics (confusion matrix, precision, and so on), and then I tried my method from our forum to deal with this problem, some results are shown as the following:</p>
<p><strong>Experiment 1</strong>: custom ViT, batch size: 64, learning rate:3e-4, epochs:30.</p>
<pre><code class="lang-auto"> 32 Learning Rate: 0.000300
 33 Epoch: 3 | Train Loss: 0.5388, Train Acc: 0.8243, Test Loss: 1.5342, Test Acc: 0.6638, Learning Rate: 0.000300
 34 Confusion Matrix:
 35             COVID  Normal  Pneumonia
 36 COVID          0       0        116
 37 Normal         0       0        317
 38 Pneumonia      0       0        855
 39 Classification Report:
 40                precision    recall  f1-score     support
 41 0              0.000000  0.000000  0.000000   116.00000
 42 1              0.000000  0.000000  0.000000   317.00000
 43 2              0.663820  1.000000  0.797947   855.00000
 44 accuracy       0.663820  0.663820  0.663820     0.66382
 45 macro avg      0.221273  0.333333  0.265982  1288.00000
 46 weighted avg   0.440657  0.663820  0.529693  1288.00000
 47 Accuracy Score: 0.6638
 48 Learning Rate: 0.000300
 49 Epoch: 4 | Train Loss: 0.5498, Train Acc: 0.7905, Test Loss: 1.4384, Test Acc: 0.6638, Learning Rate: 0.000300
 50 Confusion Matrix:
 51             COVID  Normal  Pneumonia
 52 COVID          0       0        116
 53 Normal         0       0        317
 54 Pneumonia      0       0        855
 55 Classification Report:
 56                precision    recall  f1-score     support
 57 0              0.000000  0.000000  0.000000   116.00000
 58 1              0.000000  0.000000  0.000000   317.00000
 59 2              0.663820  1.000000  0.797947   855.00000
 60 accuracy       0.663820  0.663820  0.663820     0.66382
 61 macro avg      0.221273  0.333333  0.265982  1288.00000
 62 weighted avg   0.440657  0.663820  0.529693  1288.00000
 63 Accuracy Score: 0.6638
 64 Learning Rate: 0.000300
</code></pre>
<p>this experiment is an initial experiment without any techniques to deal with the class imbalance issue, my model seems that learned nothing from the dataset. and then after a search, some solutions are: <strong>Learning rate</strong>,  “Model complexity”<br>
“Class imbalance”: before expt 1, I didn’t note my dataset was imbalanced<br>
“Optimizer and loss function choice”: I used the Adam and CrossEntropy as the same as the original paper, so I didn’t pay attention to this part.<br>
and then I continued to other experiments, firstly, I used the <code>lr_scheduler.ReduceLROnPlateau</code>  to decrease my lr util 3e-8, the confusion matrix was the same as experiment 1.<br>
<strong>Experiment 2</strong>: lr_rate:3e-3, others is the same as last experiment, but in this expt, two different kinds of stuff are:</p>
<ol>
<li>I used <code>WeightedRandomSampler</code> in train_dataloader, after this I checked the distribution of the dataloader, the proportion of COVID(minority class)was almost equal to others.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/0/a009866e413ed82677bdf069e3e65454d1e1a613.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/a009866e413ed82677bdf069e3e65454d1e1a613" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/0/a009866e413ed82677bdf069e3e65454d1e1a613_2_388x500.jpeg" alt="image" data-base62-sha1="mPKNLL5HOA6FrFZ8KYYP8PGb1DR" width="388" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/0/a009866e413ed82677bdf069e3e65454d1e1a613_2_388x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/a/0/a009866e413ed82677bdf069e3e65454d1e1a613_2_582x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/a/0/a009866e413ed82677bdf069e3e65454d1e1a613.jpeg 2x" data-dominant-color="F0F0F0"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">714×920 64.1 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></li>
<li>In the stage of splitting the dataset, I used stratified sampling, which resulted in a good distribution. (no code in this after, I manually calculated the proportion of each class in train_dataset), in this case, i think this part may be okay.</li>
<li>but the result is not good as the last experiment.</li>
</ol>
<pre><code class="lang-auto"> 95 Accuracy Score: 0.2436
 96 Learning Rate: 0.000300
 97 Epoch: 7 | Train Loss: 1.1096, Train Acc: 0.3183, Test Loss: 1.1112, Test Acc: 0.0881, Learning Rate: 0.000300
 98 Confusion Matrix:
 99             COVID  Normal  Pneumonia
100 COVID        170       0          0
101 Normal       470       0          0
102 Pneumonia   1289       0          0
103 Classification Report:
104                precision    recall  f1-score      support
105 0              0.088129  1.000000  0.161982   170.000000
106 1              0.000000  0.000000  0.000000   470.000000
107 2              0.000000  0.000000  0.000000  1289.000000
108 accuracy       0.088129  0.088129  0.088129     0.088129
109 macro avg      0.029376  0.333333  0.053994  1929.000000
110 weighted avg   0.007767  0.088129  0.014275  1929.000000
111 Accuracy Score: 0.0881
112 Learning Rate: 0.000300
113 Epoch 00008: reducing learning rate of group 0 to 3.0000e-05.
114 Epoch: 8 | Train Loss: 1.1089, Train Acc: 0.3267, Test Loss: 1.1991, Test Acc: 0.0881, Learning Rate: 0.000030
115 Confusion Matrix:
116             COVID  Normal  Pneumonia
117 COVID        170       0          0
118 Normal       470       0          0
119 Pneumonia   1289       0          0
120 Classification Report:
121                precision    recall  f1-score      support
122 0              0.088129  1.000000  0.161982   170.000000
123 1              0.000000  0.000000  0.000000   470.000000
124 2              0.000000  0.000000  0.000000  1289.000000
125 accuracy       0.088129  0.088129  0.088129     0.088129
126 macro avg      0.029376  0.333333  0.053994  1929.000000
127 weighted avg   0.007767  0.088129  0.014275  1929.000000
128 Accuracy Score: 0.0881
129 Learning Rate: 0.000030
130 Epoch: 9 | Train Loss: 1.1058, Train Acc: 0.3155, Test Loss: 1.0895, Test Acc: 0.6682, Learning Rate: 0.000030
131 Confusion Matrix:
132             COVID  Normal  Pneumonia
133 COVID          0       0        170
134 Normal         0       0        470
135 Pneumonia      0       0       1289
</code></pre>
<p>Code:</p>
<pre><code class="lang-auto">dataset = torchvision.datasets.ImageFolder(data_path, transform=common_transform)
train_dataset, test_dataset = torch.utils.data.random_split(dataset,[0.7, 0.3])

# Deal with Imbalanced Dataset
class_weights={}
for root, sudir, files in os.walk(data_path):
    if files:
        class_weights[dataset.class_to_idx[os.path.basename(root)]] = len(files)
print(class_weights) # output:{1: 1583, 2: 4273, 0: 576}
sample_weights = [0] * len(train_dataset)
for idx, (data, label) in enumerate(train_dataset):
    class_weight = class_weights[label]
    sample_weights[idx] = 1/class_weight

sampler = WeightedRandomSampler(sample_weights, num_samples=len(sample_weights), replacement=True) 

train_dataloader = DataLoader(train_dataset,
                             batch_size=batch_size,
                             sampler=sampler)
test_dataloader = DataLoader(test_dataset,
                            batch_size=batch_size)
</code></pre>
<p><strong>Note</strong>:</p>
<ol>
<li><code>class_to_idx = (COVID:0, NORMAL:1, PNEUMONIA:2}</code></li>
<li>After this expt, I thought the reason of the output was still is the same as the expt without weighted sampling may due to the order of sample_weights being incorrect, but eventually no error on that.</li>
<li>there are differences in this case, as you can see, the model always overrepresented a specific class in each epoch. I have no idea why this situation appeared after the two techniques I mentioned above were used in this experiment, and I checked the experiment simultaneously. And then I tried changing the lr using lr_scheduler as well, but nothing new. someone mentioned that combining the weightedSampling and specifying the weights for the loss function is a recommended method, so I continued to the next experiment.</li>
</ol>
<p>experiment 4:<br>
Parameters:</p>
<pre><code class="lang-auto">config:{batch_size:64
Dataset:"dataset with weighted Sampling and class weight in loss function"
epochs:50
image_size:224
initial_learning_rate:0.0003
mlp_dropout:0.1
model_name:"ViT"
transformer_layers:12}
</code></pre>
<p>Result:</p>
<pre><code class="lang-auto">  1 Epoch: 1 | Train Loss: 0.9572, Train Acc: 0.3302, Test Loss: 1.4187, Test Acc: 0.0876, Learning Rate: 0.000300
  2 Confusion Matrix:
  3             COVID  Normal  Pneumonia
  4 COVID        169       1          0
  5 Normal       470       0          0
  6 Pneumonia   1286       3          0
  7 Classification Report:
  8                precision    recall  f1-score     support
  9 0              0.087792  0.994118  0.161337   170.00000
 10 1              0.000000  0.000000  0.000000   470.00000
 11 2              0.000000  0.000000  0.000000  1289.00000
 12 accuracy       0.087610  0.087610  0.087610     0.08761
 13 macro avg      0.029264  0.331373  0.053779  1929.00000
 14 weighted avg   0.007737  0.087610  0.014218  1929.00000
 15 Accuracy Score: 0.0876
 16 Learning Rate: 0.000300
 17 Epoch: 2 | Train Loss: 0.8471, Train Acc: 0.3275, Test Loss: 1.3836, Test Acc: 0.0881, Learning Rate: 0.000300
 18 Confusion Matrix:
 19             COVID  Normal  Pneumonia
 20 COVID        170       0          0
 21 Normal       470       0          0
 22 Pneumonia   1289       0          0
 23 Classification Report:
 24                precision    recall  f1-score      support
 25 0              0.088129  1.000000  0.161982   170.000000
 26 1              0.000000  0.000000  0.000000   470.000000
 27 2              0.000000  0.000000  0.000000  1289.000000
 28 accuracy       0.088129  0.088129  0.088129     0.088129
 29 macro avg      0.029376  0.333333  0.053994  1929.000000
 30 weighted avg   0.007767  0.088129  0.014275  1929.000000
 31 Accuracy Score: 0.0881
 32 Learning Rate: 0.000300
 33 Epoch: 3 | Train Loss: 0.8325, Train Acc: 0.3444, Test Loss: 1.4703, Test Acc: 0.0881, Learning Rate: 0.000300
 34 Confusion Matrix:
 35             COVID  Normal  Pneumonia
 36 COVID        170       0          0
 37 Normal       470       0          0
 38 Pneumonia   1289       0          0
 39 Classification Report:
 40                precision    recall  f1-score      support
 41 0              0.088129  1.000000  0.161982   170.000000
 42 1              0.000000  0.000000  0.000000   470.000000
..............(same as above)
</code></pre>
<p>Code:</p>
<pre><code class="lang-auto"># Calculating weights: inverse of class frequencies
total_samples = 576 + 1583 + 4273
weights = torch.tensor([total_samples / 576, total_samples / 1583, total_samples / 4273], dtype=torch.float32)

loss_function = torch.nn.CrossEntropyLoss(weights).to(device)
optimizer = torch.optim.Adam(params=vit_model_own.parameters(),
                            lr=learning_rate,
                            betas=(0.9, 0.999),
                            weight_decay=weight_decay_optimizer)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,patience=5,threshold=0.1,verbose="True")
</code></pre>
<p>Model Summary:</p>
<pre><code class="lang-auto">==========================================================================================
Layer (type:depth-idx)                                            Param #
==========================================================================================
ViT                                                               152,064
├─Dropout: 1-1                                                    --
├─PatchEmbedding: 1-2                                             --
│    └─Conv2d: 2-1                                                590,592
│    └─Flatten: 2-2                                               --
├─Sequential: 1-3                                                 --
│    └─TransformerEncoderBlock: 2-3                               --
│    │    └─MultiheadSelfAttentionBlock: 3-1                      2,363,904
│    │    └─MLPBlock: 3-2                                         4,723,968
│    └─TransformerEncoderBlock: 2-4                               --
│    │    └─MultiheadSelfAttentionBlock: 3-3                      2,363,904
│    │    └─MLPBlock: 3-4                                         4,723,968
│    └─TransformerEncoderBlock: 2-5                               --
│    │    └─MultiheadSelfAttentionBlock: 3-5                      2,363,904
│    │    └─MLPBlock: 3-6                                         4,723,968
│    └─TransformerEncoderBlock: 2-6                               --
│    │    └─MultiheadSelfAttentionBlock: 3-7                      2,363,904
│    │    └─MLPBlock: 3-8                                         4,723,968
│    └─TransformerEncoderBlock: 2-7                               --
│    │    └─MultiheadSelfAttentionBlock: 3-9                      2,363,904
│    │    └─MLPBlock: 3-10                                        4,723,968
│    └─TransformerEncoderBlock: 2-8                               --
│    │    └─MultiheadSelfAttentionBlock: 3-11                     2,363,904
│    │    └─MLPBlock: 3-12                                        4,723,968
│    └─TransformerEncoderBlock: 2-9                               --
│    │    └─MultiheadSelfAttentionBlock: 3-13                     2,363,904
│    │    └─MLPBlock: 3-14                                        4,723,968
│    └─TransformerEncoderBlock: 2-10                              --
│    │    └─MultiheadSelfAttentionBlock: 3-15                     2,363,904
│    │    └─MLPBlock: 3-16                                        4,723,968
│    └─TransformerEncoderBlock: 2-11                              --
│    │    └─MultiheadSelfAttentionBlock: 3-17                     2,363,904
│    │    └─MLPBlock: 3-18                                        4,723,968
│    └─TransformerEncoderBlock: 2-12                              --
│    │    └─MultiheadSelfAttentionBlock: 3-19                     2,363,904
│    │    └─MLPBlock: 3-20                                        4,723,968
│    └─TransformerEncoderBlock: 2-13                              --
│    │    └─MultiheadSelfAttentionBlock: 3-21                     2,363,904
│    │    └─MLPBlock: 3-22                                        4,723,968
│    └─TransformerEncoderBlock: 2-14                              --
│    │    └─MultiheadSelfAttentionBlock: 3-23                     2,363,904
│    │    └─MLPBlock: 3-24                                        4,723,968
├─Sequential: 1-4                                                 --
│    └─LayerNorm: 2-15                                            1,536
│    └─Linear: 2-16                                               2,307
==========================================================================================
Total params: 85,800,963
Trainable params: 85,800,963
Non-trainable params: 0
==========================================================================================
</code></pre>
<p>After experiment 3, I tried to use TinyVGG to classify it as well, the result is the same. I felt I was cooked, and too many factors will lead to failure because I’m so new to this area, I thought I may not understand what was going on even from the beginning. I’m not sure whether I provided all the necessary information, leave a comment please if not.</p>
<p>Looking forward to your help!!!</p>

KFrank: <p>Hi Keahi!</p>
<p>Check for bugs, of course, but it could simply be that you have to train for much<br>
longer than thirty epochs.</p>
<aside class="quote no-group" data-username="Keahi" data-post="1" data-topic="202582" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/keahi/48/69792_2.png" class="avatar"> Keahi:</div>
<blockquote>
<p>Pneumonia with 4273 images, Normal class with 1583 images, COVID-19 with 576 images</p>
</blockquote>
</aside>
<p>Your dataset is imbalanced, but not ridiculously so.  Your dataset is not particularly<br>
large, but with 576 images in your smallest class I would think you would have enough<br>
data to make progress.  (<code>WeightedRandomSampler</code> is probably the best way to<br>
address your class imbalance.)</p>
<aside class="quote no-group">
<blockquote>
<p>batch size: 64</p>
</blockquote>
</aside>
<p>I would suggest seeing whether you can overfit your model with a single (balanced)<br>
batch.  Randomly select 21 (or 22) samples from each of your three classes and<br>
repeatedly train on that one single batch.  Can you successfully overfit your model<br>
so that it “predicts” with essentially 100% accuracy?  (Don’t expect such an overfit<br>
model to have good performance on an independent validation dataset.)</p>
<aside class="quote no-group">
<blockquote>
<p>epochs:30</p>
</blockquote>
</aside>
<p>Thirty epochs is not a lot (especially as it appears that you are not starting with a<br>
pre-trained model).  Try training <em>a lot</em> longer.</p>
<aside class="quote no-group">
<blockquote>
<p>I tried to use TinyVGG</p>
</blockquote>
</aside>
<p>I would also try to overfit TinyVGG.  If you can’t overfit either of your models you<br>
almost certainly have a bug somewhere.</p>
<p>It’s probably not the cause of your problem, but when things aren’t working, I always<br>
suggest training with plain-vanilla <code>SGD</code> and then adding momentum to it.  <code>Adam</code> is a<br>
good optimizer, but can sometimes be unstable or be more likely to get stuck.</p>
<p>Best.</p>
<p>K. Frank</p>

Keahi: <p>Hi Frank!</p>
<p>Thank you VERY MUCH, I’ve found out and fixed the problem after I used the methods from you, the bug is that I use the parameters from the original ViT paper for the optimizer(Adam), in which the weight_decay is 0.3, I didn’t notice it’s for prevent overfitting, that’s why my model learned nothing. And the reason TinyVGG has the same performance(cannot learn anything) on this task, is I used the optimizer for ViT as the optimizer of TinyVGG.</p>
<p>Thank you very much again, you said that trying to overfit a model first when you’re training a model is a very USEFUL method as well, you let me learn a lot, thank you!</p>
<p>Best regards,</p>
<p>Keahi</p>