add baseline

src/run_medclip.py → run_medclip.py

@@ -28,6 +28,7 @@ import logging
 import os
 import sys
 import time
+import getpass
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Callable, Optional
@@ -47,9 +48,9 @@ from flax import jax_utils
 from flax.jax_utils import unreplicate
 from flax.training import train_state
 from flax.training.common_utils import get_metrics, shard, shard_prng_key
-from modeling_hybrid_clip import FlaxHybridCLIP
+from src.modeling_medclip import FlaxHybridCLIP
 from transformers import AutoTokenizer, HfArgumentParser, TrainingArguments, is_tensorboard_available, set_seed
-
+import wandb
 
 logger = logging.getLogger(__name__)
 
@@ -210,7 +211,6 @@ class ImageTextDataset(VisionDataset):
         self,
         root: str,
         file_path: str,
-        captions_per_image=2,
         transform: Optional[Callable] = None,
         target_transform: Optional[Callable] = None,
         transforms: Optional[Callable] = None,
@@ -224,15 +224,21 @@ class ImageTextDataset(VisionDataset):
         self.image_paths = []
 
         for example in examples:
-            self.captions.extend(example["captions"][:captions_per_image])
-            self.image_paths.extend([example["image_path"]] * captions_per_image)
+            self.captions.append(example["caption"])
+            self.image_paths.append(f'{root}/{example["image_path"]}')
 
     def _load_image(self, idx: int):
         path = self.image_paths[idx]
         return read_image(path, mode=ImageReadMode.RGB)
 
     def _load_target(self, idx):
-        return self.captions[idx]
+        sections = self.captions[idx]
+        longest_section = max(
+            filter(lambda x: isinstance(x, str), sections.values()), 
+            key=len
+        )
+
+        return longest_section
 
     def __getitem__(self, index: int):
         image = self._load_image(index)
@@ -290,6 +296,17 @@ def main():
     else:
         model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
+    if jax.process_index() == 0:
+        wandb.init(
+            entity=getpass.getuser(),
+            project='medclip',
+            sync_tensorboard=True
+        )
+
+        wandb.config.update(model_args)
+        wandb.config.update(data_args)
+        wandb.config.update(training_args)
+
     if (
         os.path.exists(training_args.output_dir)
         and os.listdir(training_args.output_dir)
@@ -351,14 +368,12 @@ def main():
     train_dataset = ImageTextDataset(
         data_args.data_dir,
         data_args.train_file,
-        captions_per_image=2,
         transform=preprocess,
     )
 
     eval_dataset = ImageTextDataset(
         data_args.data_dir,
         data_args.validation_file,
-        captions_per_image=1,
         transform=preprocess,
     )
 

src/modeling_medclip.py

@@ -18,7 +18,7 @@ from typing import Optional, Tuple
 import flax.linen as nn
 import jax
 import jax.numpy as jnp
-from configuration_hybrid_clip import HybridCLIPConfig
+from src.configuration_medclip import HybridCLIPConfig
 from flax.core.frozen_dict import FrozenDict
 from transformers import FLAX_MODEL_MAPPING, FlaxCLIPVisionModel
 from transformers.modeling_flax_utils import FlaxPreTrainedModel

train_model.sh

@@ -0,0 +1,15 @@
+python run_medclip.py \
+    --output_dir model \
+    --text_model_name_or_path="allenai/scibert_scivocab_uncased" \
+    --vision_model_name_or_path="openai/clip-vit-base-patch32" \
+    --tokenizer_name="allenai/scibert_scivocab_uncased" \
+    --data_dir="/home/shared/data/mimic-cxr" \
+    --train_file="/home/shared/data/mimic-cxr/train_dataset.json" \
+    --validation_file="/home/shared/data/mimic-cxr/validate_dataset.json" \
+    --do_train --do_eval \
+    --num_train_epochs="40" --max_seq_length 512 \
+    --per_device_train_batch_size="64" \
+    --per_device_eval_batch_size="64" \
+    --learning_rate="5e-5" --warmup_steps="0" --weight_decay 0.1 \
+    --overwrite_output_dir \
+    --preprocessing_num_workers 32 \