import os
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import (AdamW, AutoModelForCausalLM, AutoProcessor,
get_scheduler)
from data import ObjectDetectionDataset
# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load the model and processor
# model = AutoModelForCausalLM.from_pretrained("model/Florence-2-base-ft", trust_remote_code=True).to(device)
# processor = AutoProcessor.from_pretrained("model/Florence-2-base-ft", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("microsoft/Florence-2-large-ft", revision="refs/pr/10", trust_remote_code=True, device_map="cuda") # load the model on GPU
processor = AutoProcessor.from_pretrained("microsoft/Florence-2-large-ft", revision="refs/pr/10", trust_remote_code=True)
IGNORE_ID = -100 # Pytorch ignore index when computing loss
MAX_LENGTH = 512
def collate_fn(examples):
prompt_texts = [example[0] for example in examples]
label_texts = [example[1] for example in examples]
images = [example[2] for example in examples]
inputs = processor(
images=images,
text=prompt_texts,
return_tensors="pt",
padding="longest",
max_length=MAX_LENGTH,
).to(device)
return inputs, label_texts
# Create datasets
train_dataset = ObjectDetectionDataset("train", processor=processor)
val_dataset = ObjectDetectionDataset("test", processor=processor)
# Create DataLoader
batch_size = 4
num_workers = 0
train_loader = DataLoader(
train_dataset,
batch_size=batch_size,
collate_fn=collate_fn,
num_workers=num_workers,
shuffle=True,
)
val_loader = DataLoader(
val_dataset, batch_size=batch_size, collate_fn=collate_fn, num_workers=num_workers
)
def train_model(train_loader, val_loader, model, processor, epochs=10, lr=1e-6):
optimizer = AdamW(model.parameters(), lr=lr)
num_training_steps = epochs * len(train_loader)
lr_scheduler = get_scheduler(
name="cosine",
optimizer=optimizer,
num_warmup_steps=100,
num_training_steps=num_training_steps,
)
for epoch in range(epochs):
# Training phase
model.train()
train_loss = 0
i = -1
for batch in tqdm(train_loader, desc=f"Training Epoch {epoch + 1}/{epochs}"):
i += 1
inputs, label_texts = batch
labels = processor.tokenizer(
label_texts,
return_tensors="pt",
padding="longest",
max_length=MAX_LENGTH,
return_token_type_ids=False, # no need to set this to True since BART does not use token type ids
)["input_ids"].to(device)
labels[labels == processor.tokenizer.pad_token_id] = IGNORE_ID # do not learn to predict pad tokens during training
input_ids = inputs["input_ids"]
pixel_values = inputs["pixel_values"]
outputs = model(
input_ids=input_ids, pixel_values=pixel_values, labels=labels
)
loss = outputs.loss
if i % 25 == 0:
print(loss)
generated_ids = model.generate(
input_ids=inputs["input_ids"],
pixel_values=inputs["pixel_values"],
max_new_tokens=128,
early_stopping=False,
do_sample=False,
num_beams=3,
)
generated_texts = processor.batch_decode(
generated_ids, skip_special_tokens=False
)
for generated_text, answer in zip(generated_texts, label_texts):
parsed_answer = processor.post_process_generation(
generated_text,
task="<OD>",
image_size=(
inputs["pixel_values"].shape[-2],
inputs["pixel_values"].shape[-1],
),
)
print("GT:", answer)
print("Generated Text:", generated_text)
print("Pred:", parsed_answer["<OD>"])
loss.backward()
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
train_loss += loss.item()
avg_train_loss = train_loss / len(train_loader)
print(f"Average Training Loss: {avg_train_loss}")
# Validation phase
model.eval()
val_loss = 0
with torch.no_grad():
for batch in tqdm(
val_loader, desc=f"Validation Epoch {epoch + 1}/{epochs}"
):
inputs, labels = batch
input_ids = inputs["input_ids"]
pixel_values = inputs["pixel_values"]
labels = processor.tokenizer(
text=labels,
return_tensors="pt",
padding=True,
return_token_type_ids=False,
).input_ids.to(device)
outputs = model(
input_ids=input_ids, pixel_values=pixel_values, labels=labels
)
loss = outputs.loss
val_loss += loss.item()
avg_val_loss = val_loss / len(val_loader)
print(f"Average Validation Loss: {avg_val_loss}")
# Save model checkpoint
output_dir = f"./model_checkpoints/epoch_{epoch+1}"
os.makedirs(output_dir, exist_ok=True)
model.save_pretrained(output_dir)
processor.save_pretrained(output_dir)
for param in model.vision_tower.parameters():
param.requires_grad = False
model_total_params = sum(p.numel() for p in model.parameters())
model_train_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Number of trainable parameters {model_train_params} out of {model_total_params}, rate: {model_train_params/model_total_params:0.3f}")
train_model(train_loader, val_loader, model, processor, epochs=3, lr=1e-6)
model.push_to_hub("danelcsb/Florence-2-FT-cppe-5")
processor.push_to_hub("danelcsb/Florence-2-FT-cppe-5")
