File size: 2,947 Bytes
87b3d4b |
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 |
from datasets import load_dataset
import torch
from transformers import PaliGemmaProcessor, PaliGemmaForConditionalGeneration, Trainer, TrainingArguments, BitsAndBytesConfig
from peft import get_peft_model, LoraConfig
import os
USE_LORA = False
USE_QLORA = False
FREEZE_VISION = False
ds = load_dataset('merve/vqav2-small', split="validation")
ds = ds.train_test_split(test_size=0.5)["train"]
model_id = "google/paligemma2-3b-pt-448"
processor = PaliGemmaProcessor.from_pretrained(model_id)
device = "cuda" if torch.cuda.is_available() else "cpu"
image_token = processor.tokenizer.convert_tokens_to_ids("<image>")
def collate_fn(examples):
texts = ["<image>answer en " + example["question"] for example in examples]
labels= [example['multiple_choice_answer'] for example in examples]
images = [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
if USE_LORA or USE_QLORA:
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",
)
if USE_QLORA:
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_type=torch.bfloat16
)
model = PaliGemmaForConditionalGeneration.from_pretrained(model_id, device_map="auto",
quantization_config=bnb_config if USE_QLORA else None,
torch_dtype=torch.bfloat16)
model = get_peft_model(model, lora_config)
model = model.to(device)
model.print_trainable_parameters()
else:
model = PaliGemmaForConditionalGeneration.from_pretrained(model_id, device_map="auto").to(device)
model = model.to(device)
if FREEZE_VISION:
for param in model.vision_tower.parameters():
param.requires_grad = False
for param in model.multi_modal_projector.parameters():
param.requires_grad = False
args=TrainingArguments(
num_train_epochs=3,
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,
adam_beta2=0.999,
logging_steps=100,
optim="adamw_hf",
save_strategy="steps",
save_steps=1000,
save_total_limit=1,
push_to_hub=True
output_dir="paligemma_vqav2",
bf16=True,
report_to=["tensorboard"],
dataloader_pin_memory=False
)
trainer = Trainer(
model=model,
train_dataset=ds ,
data_collator=collate_fn,
args=args
)
trainer.train()
|