Spaces:
Paused
Paused
File size: 2,791 Bytes
3ac15b8 bc6660c 3ac15b8 bc6660c ba93638 bc6660c 3ac15b8 | 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 | from transformers import TrainingArguments, Trainer, DetrForObjectDetection, DetrImageProcessor
import torch
import torchvision
from datasets import load_dataset, DatasetDict, Image
import os
# Set TRANSFORMERS_CACHE to a writable directory
os.environ["TRANSFORMERS_CACHE"] = "/app/cache"
# Ensure the directory exists
os.makedirs("/app/cache", exist_ok=True)
# Check if GPU is available
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")
if torch.cuda.is_available():
print(f"CUDA Device: {torch.cuda.get_device_name(0)}")
# Configuration
MODEL_NAME = "facebook/detr-resnet-50" # DETR (DEtection TRansformer) model
BATCH_SIZE = 4
EPOCHS = 10
LEARNING_RATE = 5e-5
OUTPUT_DIR = "./results"
# Load Pre-trained Model and Processor
processor = DetrImageProcessor.from_pretrained(MODEL_NAME)
model = DetrForObjectDetection.from_pretrained(MODEL_NAME, num_labels=91).to(device) # 91 is the number of COCO classes
# Load COCO Dataset
# Replace with paths to your custom dataset
coco_dataset = DatasetDict({
"train": load_dataset("coco", data_files={"train": "dataset.coco.json"}, split="train"),
"validation": load_dataset("coco", data_files={"validation": "dataset.coco.json"}, split="validation"),
})
# Preprocess Function for Images
def preprocess_coco(example):
image = Image.open(example["file_name"]).convert("RGB") # Ensure RGB format
target = {
"boxes": torch.tensor(example["bbox"]),
"labels": torch.tensor(example["category_id"]),
}
encoding = processor(images=image, annotations=target, return_tensors="pt")
return encoding
# Apply Preprocessing
coco_dataset = coco_dataset.map(preprocess_coco, batched=True)
# Data Collator
def collate_fn(batch):
images = [item["pixel_values"].squeeze(0) for item in batch]
annotations = [item["labels"] for item in batch]
return {"pixel_values": torch.stack(images), "labels": annotations}
# Training Arguments
training_args = TrainingArguments(
output_dir=OUTPUT_DIR,
evaluation_strategy="epoch",
learning_rate=LEARNING_RATE,
per_device_train_batch_size=BATCH_SIZE,
num_train_epochs=EPOCHS,
save_strategy="epoch",
logging_dir="./logs",
logging_steps=10,
load_best_model_at_end=True,
push_to_hub=False,
)
# Trainer for Object Detection
trainer = Trainer(
model=model,
args=training_args,
train_dataset=coco_dataset["train"],
eval_dataset=coco_dataset["validation"],
tokenizer=processor, # Not used for object detection, but required by Trainer
data_collator=collate_fn,
)
# Train the Model
print("Starting Training...")
trainer.train()
print("Training Complete!")
# Save Model
print("Saving Model...")
trainer.save_model(OUTPUT_DIR)
print(f"Model saved to {OUTPUT_DIR}")
|