feat: Add optimized train script with fp16 precision

- Add new file for optimized training using fp16
- Changed bfp16 to fp16 in model initialization and training arguments
- Updated collate_fn to use fp16 instead of bfp16

article_base_train_fp16.py

@@ -0,0 +1,169 @@
+import os, time, math
+import pandas as pd
+from datasets import Dataset
+from transformers import PaliGemmaProcessor, PaliGemmaForConditionalGeneration, BitsAndBytesConfig, TrainingArguments, Trainer
+import torch
+from PIL import Image
+from peft import get_peft_model, LoraConfig
+import argparse
+
+
+def load_custom_dataset_from_csv(csv_file, image_folder):
+    data = pd.read_csv(csv_file)
+    
+    questions = data['question'].tolist()
+    images = [os.path.join(image_folder, img) for img in data['image'].tolist()]
+    answers = data['answer'].tolist()
+    
+    return Dataset.from_dict({
+        'question': questions,
+        'image': images,
+        'answer': answers
+    })
+
+
+def load_custom_dataset_from_parquet(parquet_file, image_folder):
+    data = pd.read_parquet(parquet_file)
+    
+    questions = data['question'].tolist()
+    images = [os.path.join(image_folder, img) for img in data['image'].tolist()]
+    answers = data['answer'].tolist()
+    
+    return Dataset.from_dict({
+        'question': questions,
+        'image': images,
+        'answer': answers
+    })
+
+
+def load_dataset_by_type(metadata_type, dataset_dir, image_folder):
+    if metadata_type == "csv":
+        return load_custom_dataset_from_csv(
+            os.path.join(dataset_dir, 'train_samples.csv'),
+            image_folder
+        )
+    elif metadata_type == "parquet":
+        return load_custom_dataset_from_parquet(
+            os.path.join(dataset_dir, 'train.parquet'),
+            image_folder
+        )
+    else:
+        raise ValueError("Unsupported metadata type. Use 'csv' or 'parquet'.")
+
+
+def load_model_and_args(use_qlora, model_id, device, output_dir):
+    if use_qlora:
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.float16  # Changed from bfloat16 to float16
+        )
+        lora_config = LoraConfig(
+            r=8,
+            target_modules=["q_proj", "o_proj", "k_proj", "v_proj", "gate_proj", "up_proj", "down_proj"],
+            task_type="CAUSAL_LM"
+        )
+        
+        model = PaliGemmaForConditionalGeneration.from_pretrained(model_id, quantization_config=bnb_config, device_map={"": 0})
+        model = get_peft_model(model, lora_config)
+        model.print_trainable_parameters()
+        
+        args = TrainingArguments(
+            output_dir=os.path.join(output_dir, f"{math.floor(time.time())}"),
+            num_train_epochs=2,
+            remove_unused_columns=False,
+            per_device_train_batch_size=1,
+            gradient_accumulation_steps=4,
+            warmup_steps=2,
+            learning_rate=2e-5,
+            weight_decay=1e-6,
+            logging_steps=100,
+            optim="adamw_hf",
+            save_strategy="steps",
+            save_steps=1000,
+            save_total_limit=1,
+            fp16=True,  # Changed from bf16 to fp16
+            report_to=["tensorboard"],
+            dataloader_pin_memory=False
+        )
+        
+        return model, args
+    else:
+        model = PaliGemmaForConditionalGeneration.from_pretrained(model_id, torch_dtype=torch.float16).to(device)  # Changed from bfloat16 to float16
+        for param in model.vision_tower.parameters():
+            param.requires_grad = False
+
+        for param in model.multi_modal_projector.parameters():
+            param.requires_grad = True
+        
+        args = TrainingArguments(
+            output_dir=os.path.join(output_dir, f"{math.floor(time.time())}"),
+            num_train_epochs=2,
+            remove_unused_columns=False,
+            per_device_train_batch_size=4,
+            gradient_accumulation_steps=4,
+            warmup_steps=2,
+            learning_rate=2e-5,
+            weight_decay=1e-6,
+            logging_steps=100,
+            optim="paged_adamw_8bit",
+            save_strategy="steps",
+            save_steps=1000,
+            save_total_limit=1,
+            fp16=True,  # Changed from bf16 to fp16
+            report_to=["tensorboard"],
+            dataloader_pin_memory=False
+        )
+
+        return model, args
+
+
+def main(args):
+    dataset_dir = args.dataset_dir
+    model_id = args.model_id
+    output_dir = args.output_dir
+    metadata_type = args.metadata_type
+    
+    dataset = load_dataset_by_type(metadata_type, dataset_dir, os.path.join(dataset_dir, 'images'))
+    train_val_split = dataset.train_test_split(test_size=0.1)
+    
+    train_ds = train_val_split['train']
+    val_ds = train_val_split['test']
+
+    processor = PaliGemmaProcessor.from_pretrained(model_id)
+    device = "cuda"
+
+    model, args = load_model_and_args(args.use_qlora, model_id, device, output_dir)
+
+    def collate_fn(examples):
+        texts = [example["question"] for example in examples]
+        labels = [example['answer'] for example in examples]
+        images = [Image.open(example['image']).convert("RGB") for example in examples]
+        tokens = processor(text=texts, images=images, suffix=labels, return_tensors="pt", padding="longest")
+        tokens = tokens.to(torch.float16).to(device)  # Changed from bfloat16 to float16
+        return tokens
+    
+    trainer = Trainer(
+        model=model,
+        train_dataset=train_ds,
+        eval_dataset=val_ds,
+        data_collator=collate_fn,
+        args=args
+    )
+    
+    trainer.train()
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Train a model with custom dataset")
+    parser.add_argument('--dataset_dir', type=str, default='./dataset', help='Path to the folder containing the images')
+    parser.add_argument('--model_id', type=str, default='google/paligemma-3b-pt-224', help='Model ID to use for training')
+    parser.add_argument('--output_dir', type=str, default='./output', help='Directory to save the output')
+    parser.add_argument('--use_qlora', type=bool, default=False, help='Use QLoRA for training')
+    parser.add_argument('--metadata_type', type=str, default='parquet', choices=['csv', 'parquet'], help='Metadata format (csv or parquet)')
+    return parser.parse_args()
+
+    
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)