File size: 6,735 Bytes
1758c0c |
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 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 |
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
# Function to load custom dataset from CSV
def load_custom_dataset_from_csv(csv_file, image_folder):
# Load CSV data using pandas
data = pd.read_csv(csv_file)
# Prepare dataset format for Hugging Face
questions = data['question'].tolist()
images = [os.path.join(image_folder, img) for img in data['image'].tolist()]
answers = data['answer'].tolist()
# Create a Hugging Face dataset from the loaded CSV
return Dataset.from_dict({
'question': questions,
'image': images,
'answer': answers
})
# Function to load custom dataset from Parquet
def load_custom_dataset_from_parquet(parquet_file, image_folder):
# Load Parquet data using pandas
data = pd.read_parquet(parquet_file)
# Prepare dataset format for Hugging Face
questions = data['question'].tolist()
images = [os.path.join(image_folder, img) for img in data['image'].tolist()]
answers = data['answer'].tolist()
# Create a Hugging Face dataset from the loaded Parquet
return Dataset.from_dict({
'question': questions,
'image': images,
'answer': answers
})
# Choose the appropriate loader based on metadata_type argument
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.bfloat16
)
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()
# TODO: Customize training setting
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,
bf16=True,
report_to=["tensorboard"],
dataloader_pin_memory=False
)
return model, args
else:
model = PaliGemmaForConditionalGeneration.from_pretrained(model_id, torch_dtype=torch.bfloat16).to(device)
for param in model.vision_tower.parameters():
param.requires_grad = False
for param in model.multi_modal_projector.parameters():
param.requires_grad = True
# TODO: Customize training setting
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,
bf16=True,
report_to=["tensorboard"],
dataloader_pin_memory=False
)
return model, args
# Main training function
def main(args):
dataset_dir = args.dataset_dir
model_id = args.model_id
output_dir = args.output_dir
metadata_type = args.metadata_type
# Load custom datasetsㄴ
# dataset = load_custom_dataset_from_csv(
# os.path.join(dataset_dir, 'train_samples.csv'),
# os.path.join(dataset_dir, 'images/train')) # TODO: change to appropriate path
dataset = load_dataset_by_type(metadata_type, dataset_dir, os.path.join(dataset_dir, 'images/train'))
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)
# Custom collate function
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.bfloat16).to(device)
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) |