Okay so to state first:
I have no coding / AI background, i have been learning this over the past few months, and i want to learn more!
After some time i created an image classifier with the help of pytorch lightning with the medmnist dataset, when i was confident that everything worked, i changed to another dataset.
But now it only guesses one class. I have been researching this problem and it often says data imbalance, which is an issue for me. But i thought with correct weights it should fix the problem of only guessing 1 class. But this did not work for me. I have tried to have more balanced data but still same issue.
So i have been trying a lot of solutions and checking if the optimizers work correctly, which i think they do..
Also
def main():
import torch
import os
from torch import optim, utils, Tensor
import torch.nn as nn
import torch.utils.data as data
from torch.utils.tensorboard import SummaryWriter
from torchvision.transforms import ToTensor
from tqdm import tqdm
import torchmetrics
from torchmetrics.classification import MulticlassROC, MulticlassAUROC
from lightning.pytorch.callbacks import ModelCheckpoint
import lightning as L
import numpy as np
from sklearn.utils.class_weight import compute_class_weight
import datetime
current_datetime = datetime.datetime.now()
month = current_datetime.month
day = current_datetime.day
hour = current_datetime.hour
minute = current_datetime.minute
todaysdate = "{:02d}-{:02d}_{:02d}:{:02d}".format(month, day, hour, minute)
"""# Model - Data Prep"""
# Importing the datasets
print(f"{todaysdate} nImporting Perfusion dataset... ")
from MyFile.Tryouts.NewLoader import perfusionmaps_dataset
## Dataset:
dataset_train = perfusionmaps_dataset(mode='train', method='sv', relative=True)
dataset_val = perfusionmaps_dataset(mode='val', method='sv', relative=True)
dataset_test = perfusionmaps_dataset(mode='test', method='sv', relative=True)
print(f' Trainset: {len(dataset_train)}, '
f'valset: {len(dataset_val)}, testset: {len(dataset_test)}')
lr = 0.001
## Dataloader
batch_size_final = 1
train_loader = data.DataLoader(dataset_train, batch_size=batch_size_final, shuffle=True)
val_loader = data.DataLoader(dataset_val, batch_size=batch_size_final, shuffle=False)
test_loader = data.DataLoader(dataset_test, batch_size=batch_size_final, shuffle=False)
# ----------------------------------------------------------------------------------------------------------------------------
# Label names printing
classes_labels = {
0: "Other (m3 or anterior)",
1: "ICA",
2: "ICA-T",
3: "M1",
4: "M2",
5: "No occlusion"
}
num_classes = len(classes_labels)
classes = list(classes_labels.keys())
label_names = list(classes_labels.values())
print(f" Classes with Labels: {num_classes}n"
f" Classes numerical: {classes} n"
f" Label names: {label_names[:3]}n"
f" ...{label_names[3:]}n"
f" Learning Rate: {lr}n"
f" Batch Size: {batch_size_final}n"
print("... Done!n")
# CCN MODEL DEFINITION
print(f"Model setup:")
# Ques_print = input(" Model size diagnosis? (Yes/No) ").lower()
Ques_print = "No".lower()
Q_New_or_save = input(" New training (New) or from earlier save (Save)?: ").lower()
# Q_New_or_save = "new".lower()
if Q_New_or_save == "new":
# NUM_EPOCHS = int(input(" How many epochs?: "))
NUM_EPOCHS = 5
print(f"... Setup done")
writer = SummaryWriter("")
class CNNModel(L.LightningModule):
def __init__(self):
super(CNNModel, self).__init__()
# input (1, 256,256,32)
## Layer 1
self.conv1 = nn.Conv3d(1, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=1)
self.relu1 = nn.ReLU()
self.pooling1 = nn.MaxPool3d(kernel_size=2, stride=2)
# output (32, 128,128,16)
## Layer 2
self.conv2 = nn.Conv3d(32, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=1)
self.relu2 = nn.ReLU()
self.pooling2 = nn.MaxPool3d(kernel_size=2, stride=2)
# output (64, 64, 64, 8)
## Layer 3
self.conv3 = nn.Conv3d(64, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=1)
self.relu3 = nn.ReLU()
self.pooling3 = nn.MaxPool3d(kernel_size=2, stride=2)
# output (128, 32,32,4)
## Layer 4 - Extra layered
self.conv4 = nn.Conv3d(128, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=1)
self.relu4 = nn.ReLU()
self.pooling4 = nn.MaxPool3d(kernel_size=2, stride=2)
# output (256, 16,16,2)
self.flatten = nn.Flatten()
# output (524288) - (131072)
self.fc1 = nn.Linear(256 *16*16*2, 1028)
self.relu5 = nn.ReLU()
# self.fc1 = nn.Linear(128 * 32 * 32 * 4, 512)
# self.relu4 = nn.ReLU()
# output (512) - (1028)
self.fc2 = nn.Linear(1028, 256)
self.relu5 = nn.ReLU()
# output (256)
# self.fc2 = nn.Linear(512, num_classes)
self.fc3 = nn.Linear(256, num_classes)
self.softmax = nn.Softmax(dim=1)
# output (6) Classes
# Extra
self.on_validation_step_outputs = []
self.on_test_step_outputs = []
## Torch Lightning metrics
# Accuracy
self.train_acc = torchmetrics.classification.Accuracy(task='multiclass',
num_classes=num_classes,
average='weighted')
self.val_acc = torchmetrics.classification.Accuracy(task='multiclass',
num_classes=num_classes,
average='none')
self.val_W_acc = torchmetrics.classification.Accuracy(task='multiclass',
num_classes=num_classes,
average='weighted',
multidim_average='global')
self.test_acc = torchmetrics.classification.Accuracy(task='multiclass',
num_classes=num_classes,
# Loss functions
self.criterion = nn.CrossEntropyLoss()
def forward(self, x):
batch_size, channels, width, height, depth = x.size()
# print(f"batch_size:{batch_size}")
if Ques_print == "yes":
print(f"Before anything: {x.size()}")
x = self.relu1(self.conv1(x))
x = self.pooling1(x)
if Ques_print == "yes":
print(f"after conv1: {x.size()}")
x = self.relu2(self.conv2(x))
x = self.pooling2(x)
if Ques_print == "yes":
print(f"after conv2: {x.size()}")
x = self.relu3(self.conv3(x))
x = self.pooling3(x)
if Ques_print == "yes":
print(f"after conv3: {x.size()}")
x = self.relu4(self.conv4(x))
x = self.pooling4(x)
if Ques_print == "yes":
print(f"after conv4: {x.size()}")
# Flattening
x = x.view(-1, 131072)
if Ques_print == "yes":
print(f"after flattening: {x.size()}")
# x = self.relu4(self.fc1(x))
x = self.relu5(self.fc1(x))
if Ques_print == "yes":
print(f"after fc1: {x.size()}")
# x = self.fc2(x)
x = self.relu5(self.fc2(x))
if Ques_print == "yes":
print(f"after fc2: {x.size()}")
x = self.fc3(x)
if Ques_print == "yes":
print(f"after fc3: {x.size()}")
x = self.softmax(x)
if Ques_print == "yes":
print(f"after softmax: {x.size()}")
print(f'end x: {x}')
return x
def training_step(self, batch, batch_idx):
# Prep
x, labels, Labels_int = batch
labels_int = [tensor[0].item() for tensor in Labels_int]
comp_class_weights = compute_class_weight('balanced',
classes=np.unique(labels_int),
y=labels_int)
class_weights = torch.tensor(comp_class_weights,
dtype=torch.float32).to(self.device)
print(f'Class Weights: {class_weights}')
# # Training
x = x.float()
x = x.view(x.size(0), 1, 256, 256, 32)
logits = self.forward(x)
# print(f"nLogits: {logits}nlabels: {labels}")
# Loss
criterion_weighted = nn.CrossEntropyLoss(weight=class_weights)
loss = criterion_weighted(logits, labels)
self.log("train/loss", loss)
tensorboard_logs = {'train_loss': loss}
# Track accuracy
self.train_acc(logits, labels)
self.log('train/acc', self.train_acc,
on_step=False, on_epoch=True)
# return loss
return {'loss': loss, 'log': tensorboard_logs}
def validation_step(self, batch, batch_idx):
x_val, labels_val = batch
x_val = x_val.float()
logits = self.forward(x_val)
loss_val = self.criterion(logits, labels_val)
self.on_validation_step_outputs.append(loss_val)
# Acc per Class
labels_val = torch.argmax(labels_val, dim=1)
# print(f"Problem with val_acc?: {logits.shape},{labels_val.shape}")
self.val_acc(logits, labels_val)
accuracy_value = self.val_acc.compute()
# print(f"accuracy_value: {accuracy_value}")
# print(f"accuracy_value.shape: {accuracy_value.shape}")
for i in range(accuracy_value.shape[0]):
self.log(f'val/acc/class_{i}', accuracy_value[i],
on_step=True, on_epoch=True)
# Acc weighted mean
# print(f"Problem with val_W_acc?: {logits.shape},{labels_val.shape}")
self.val_W_acc(logits, labels_val)
accuracy_W_value = self.val_W_acc.compute()
accuracy_W_value = accuracy_W_value.item()
self.log('val/W_acc', accuracy_W_value,
on_step=False, on_epoch=True)
return loss_val, accuracy_W_value
def on_validation_epoch_end(self):
average_val_loss = torch.stack(self.on_validation_step_outputs).mean()
self.log("val/loss", average_val_loss, prog_bar=True)
tensorboard_logs = {'val/loss': average_val_loss}
self.on_validation_step_outputs.clear()
return {'avg_val_loss': average_val_loss, 'log': tensorboard_logs}
def test_step(self, batch, batch_idx):
x_test, labels_test = batch
x_test = x_test.float()
y_test = self.forward(x_test)
print(f'ny_test: {y_test}')
print(f'labels_test: {labels_test}')
# labels_test = labels_test.view(-1).long()
# print("Before loss function")
loss_test = self.criterion(y_test, labels_test)
self.on_test_step_outputs.append(loss_test)
labels_test = torch.argmax(labels_test, dim=1)
print(f"y_test: {y_test}"
f"labels_test: {labels_test}")
# print("Before accuracy function")
accuracy_value_test = self.test_acc(y_test, labels_test)
# accuracy_value_test = test_acc_.compute()
for i in range(accuracy_value_test.shape[0]):
self.log(f'test/acc/class_{i}', accuracy_value_test[i],
on_step=True, on_epoch=False)
return loss_test
def on_test_epoch_end(self):
average_test_loss = torch.stack(self.on_test_step_outputs).mean()
self.log("test/loss", average_test_loss, prog_bar=True)
tensorboard_logs = {'test/loss': average_test_loss}
self.on_test_step_outputs.clear()
return {'avg_test_loss': average_test_loss, 'log': tensorboard_logs}
def configure_optimizers(self):
optimizer = optim.Adam(self.parameters(), lr=lr)
return optimizer
# Initialize the CNN model
cnn_model = CNNModel()
# Callbacks
ckpt_callback = ModelCheckpoint(monitor='val/W_acc', save_top_k=1, mode='max',
dirpath=save_callback_path,
filename=
'callback_Ep{epoch}_{val/W_acc:.2f}'+f'_{todaysdate}',
auto_insert_metric_name=False)
## Train the model
trainer = L.Trainer(max_epochs=NUM_EPOCHS,
callbacks=ckpt_callback
)
trainer.fit(model=cnn_model,
train_dataloaders=train_loader,
val_dataloaders=val_loader
)
if __name__ == '__main__':
main()
And this is the summary of the model:
| Name | Type | Params
--------------------------------------------------
0 | conv1 | Conv3d | 896
1 | relu1 | ReLU | 0
2 | pooling1 | MaxPool3d | 0
3 | conv2 | Conv3d | 55.4 K
4 | relu2 | ReLU | 0
5 | pooling2 | MaxPool3d | 0
6 | conv3 | Conv3d | 221 K
7 | relu3 | ReLU | 0
8 | pooling3 | MaxPool3d | 0
9 | conv4 | Conv3d | 884 K
10 | relu4 | ReLU | 0
11 | pooling4 | MaxPool3d | 0
12 | flatten | Flatten | 0
13 | fc1 | Linear | 134 M
14 | relu5 | ReLU | 0
15 | fc2 | Linear | 263 K
16 | fc3 | Linear | 1.5 K
17 | softmax | Softmax | 0
18 | train_acc | MulticlassAccuracy | 0
19 | val_acc | MulticlassAccuracy | 0
20 | val_W_acc | MulticlassAccuracy | 0
21 | test_acc | MulticlassAccuracy | 0
22 | criterion | CrossEntropyLoss | 0
--------------------------------------------------
136 M Trainable params
0 Non-trainable params
136 M Total params
544.682 Total estimated model params size (MB)
Further as i said before, im learning and still want to learn so is there someone who can see the issue ‘:)
I have tried:
balancing the datasets
tinkering with the optimizers
etc.
New contributor
Nick Staalstra is a new contributor to this site. Take care in asking for clarification, commenting, and answering.
Check out our Code of Conduct.