File size: 6,265 Bytes
0601fa5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 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) |