Spaces:
Sleeping
Sleeping
File size: 2,419 Bytes
f3f6f5d | 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 | """Train RF-DETR model on car_data for aerial car detection."""
import argparse
from pathlib import Path
import rfdetr
MODEL_CLASSES: dict[str, type] = {
"nano": rfdetr.RFDETRNano,
"small": rfdetr.RFDETRSmall,
"base": rfdetr.RFDETRBase,
"medium": rfdetr.RFDETRMedium,
"large": rfdetr.RFDETRLarge,
}
def run_training(
dataset_dir: str | Path,
epochs: int = 50,
batch_size: int = 4,
lr: float = 1e-4,
resolution: int = 640,
output_dir: str = "output",
model_size: str = "base",
grad_accum_steps: int = 1,
num_classes: int = 1,
) -> None:
"""Run RF-DETR training.
Args:
dataset_dir: Path to dataset (YOLO or COCO format, auto-detected).
epochs: Number of training epochs.
batch_size: Batch size.
lr: Learning rate.
resolution: Input resolution.
output_dir: Checkpoint output directory.
model_size: Model variant (nano/small/base/medium/large).
grad_accum_steps: Gradient accumulation steps.
num_classes: Number of object classes.
"""
model_cls = MODEL_CLASSES.get(model_size)
if model_cls is None:
raise ValueError(
f"Unknown model_size {model_size!r}, "
f"choose from: {', '.join(MODEL_CLASSES)}"
)
model = model_cls()
model.train(
dataset_dir=str(dataset_dir),
epochs=epochs,
batch_size=batch_size,
lr=lr,
resolution=resolution,
output_dir=output_dir,
num_classes=num_classes,
grad_accum_steps=grad_accum_steps,
run_test=False,
)
def main() -> None:
parser = argparse.ArgumentParser(description="Train RF-DETR on car_data")
parser.add_argument("--epochs", type=int, default=50)
parser.add_argument("--batch-size", type=int, default=4)
parser.add_argument("--lr", type=float, default=1e-4)
parser.add_argument("--resolution", type=int, default=640)
parser.add_argument("--output-dir", type=str, default="output")
args = parser.parse_args()
training_dir = Path(__file__).resolve().parent
dataset_dir = training_dir / "car_data" / "mydata" / "mydata"
run_training(
dataset_dir=dataset_dir,
epochs=args.epochs,
batch_size=args.batch_size,
lr=args.lr,
resolution=args.resolution,
output_dir=args.output_dir,
)
if __name__ == "__main__":
main()
|