This medical Contrastive Language-Image Pre-Training will be used on MIMIC dataset to predict classes
-
Define our custom image and text encoders. In this example, we will use ResNet50 for image encoding and a simple LSTM for text encoding. These encoders will be passed to the CLIP model later.
-
Load our dataset. In this example, we will use the CIFAR10 dataset as an example.
-
Initialize our model with the custom encoders. Here we will initialize our CLIP model with our custom image and text encoders.
-
Define our loss function and optimizer. In this example, we will use the contrastive loss function and AdamW optimizer.
-
Train our model. In this example, we will train the model for 10 epochs.
-
In the training code , we first evaluate the model on the validation set and calculate the F1 score. To do this, we set the model to evaluation mode using
model.eval()
, then we loop over the validation set and calculate the cosine similarity between the image and text embeddings. We then threshold the similarities at 0.5 to obtain binary predictions, and compute the F1 score using these predictions and the ground truth labels. Finally, we set the model back to training mode usingmodel.train()
. We repeat this evaluation and F1 score calculation after each epoch of training.
To use the Roberta model for text encoding and add a hyperparameter search using Optuna, we proceeded as follow:
-
Define our custom image and text encoders. In this example, we will use ResNet50 for image encoding and Roberta for text encoding. These encoders will be passed to the CLIP model later
-
Load our dataset. In this example, we will use the CIFAR10 dataset as an example.
-
Initialize your model with the custom encoders. Here we will initialize our CLIP model with our custom image and text encoders.
-
Define your loss function and optimizer. In this example, we will use the contrastive loss function and AdamW optimizer.
-
Define a function to train the model with Optuna hyperparameter search. We will define the number of epochs, learning rate, and batch size as hyperparameters to search over.
-
Inside the training loop, we will train the model using the current hyperparameters.
-
At the end of each epoch, we will compute the validation loss and return it as the objective value for Optuna to minimize.
-
Run the Optuna hyperparameter search.
-
Get the best hyperparameters found by Optuna and retrain the model using those hyperparameters.
-
Using a function, we calculated the
F1 score
on the test set by comparing the binary predictions to the ground truth labels using thecompute_f1_score
function. Note that we need to convert the predictions and labels back to CPU before computing theF1 score
.
In stack_model.py
, we define a new ImageEncoder
class that loads ResNet152, ResNet18, and DenseNet from the torchvision.models module, removes their last fully connected layer, and concatenates their features. We freeze all parameters of these models to prevent them from being trained during training.
We can use the same dataset and the Roberta model for text encoding as before, and we can use the same training and evaluation code with the new ImageEncoder
class.
To do so, import stack_model.ImageEncoder
and replace the image_encoder
in the roberta_clip_optuna.py
(row 63) to use the staking model instead of ResNet50.