add collate

Author: lukeingawesome

llm2clip/run.sh

@@ -5,18 +5,18 @@ python -m torch.distributed.launch --nproc_per_node=2 \
         --enable-deepspeed \
         --grad-checkpointing \
         --name="T_vitl336_mimic" \
-        --save-frequency 1  \
-        --zeroshot-frequency 1 \
+        --save-frequency 2  \
+        --zeroshot-frequency 2 \
         --report-to="tensorboard, wandb" \
         --wandb-project-name="LLM2CLIP" \
         --wandb-notes="EVA02-CLIP-L-14-336" \
         --train-data "/data/csv/llm2clip/mimic_clip.csv" \
         --pretrained=${PRETRAINED} \
         --precision "fp16" \
         --warmup 0 \
-        --batch-size=256 \
-        --eval-batch-size=256 \
-        --log-every-n-steps 100 \
+        --batch-size=150 \
+        --eval-batch-size=150 \
+        --log-every-n-steps 200 \
         --epochs=20 \
         --lr=1e-5 \
         --visual-lr=1e-5 \
@@ -41,4 +41,4 @@ python -m torch.distributed.launch --nproc_per_node=2 \
         --zero-stage=1 \
         --dataset-type "cxr" \
         --csv-img-key "img_path" \
-        --csv-caption-key "caption"
+        --csv-caption-key "caption2"

llm2clip/training/data.py

@@ -84,11 +84,51 @@ def __getitem__(self, idx):
         if self.transform:
             image = self.transform(image)
 
-        # Process caption if tokenizer is provided
-        if self.tokenizer:
-            caption = self.tokenizer([caption])[0]
+        # # Process caption if tokenizer is provided
+        # if self.tokenizer:
+        #     caption = self.tokenizer([caption])[0]
 
         return image, caption
+    def collate_fn(self, batch):
+        images, texts = zip(*batch)
+        images = torch.stack(images)
+        
+        # Split texts
+        texts_2 = []
+        original_texts = []
+        for text in texts:
+            t = text.split("!@#$%^&*()")
+            texts_2.append(t[1] if len(t) > 1 else "")
+            original_texts.append("".join(t))
+
+        # Tokenize original texts with padding
+        
+        original = self.tokenizer(
+            original_texts,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=512,
+        )
+
+        # Process secondary texts and create embed masks
+        embed_mask = torch.zeros_like(original["attention_mask"])
+        for i, t in enumerate(texts_2):
+            if t:  # Only process non-empty secondary texts
+                ids = self.tokenizer(
+                    [t],
+                    return_tensors="pt",
+                    padding=True,
+                    truncation=True,
+                    max_length=512,
+                    add_special_tokens=False,
+                )
+                if len(ids["input_ids"][0]) > 0:
+                    embed_mask[i, -len(ids["input_ids"][0]):] = 1
+
+        original["embed_mask"] = embed_mask
+        return images, original
+    
 
 # Example usage:
 
@@ -799,6 +839,7 @@ def get_cxr_dataset(args, preprocess_fn, is_train, epoch=0, tokenizer=None):
         pin_memory=True,
         sampler=sampler,
         drop_last=is_train,
+        collate_fn=dataset.collate_fn
     )
     dataloader.num_samples = num_samples
     dataloader.num_batches = len(dataloader)

llm2clip/training/params.py

@@ -168,7 +168,7 @@ def parse_args(args):
     parser.add_argument(
         "--logs",
         type=str,
-        default="logs",
+        default="/model/llm2clip/logs",
         help="Where to store tensorboard logs. Use None to avoid storing logs.",
     )
     parser.add_argument(