Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import torch
- from ultralytics import YOLO
- import os
- import glob
- import datetime
- def get_latest_model():
- """Mencari model terbaru di folder hasil training."""
- model_paths = glob.glob("runs/detect/train*/weights/best.pt")
- if not model_paths:
- print("❌ Tidak ada model ditemukan! Pastikan training sudah dilakukan.")
- return None
- latest_model = max(model_paths, key=os.path.getctime)
- print(f"✅ Model terbaru ditemukan: {latest_model}")
- return latest_model
- def train_models_few_shot():
- """Training YOLO dengan dataset kecil (Few-Shot Learning)."""
- dataset_path = "assets/datasets/dataset.yaml"
- model_path = "yolov8n.pt" # Bisa diganti dengan yolov8s.pt untuk akurasi lebih baik
- # Cek apakah model YOLO tersedia
- if not os.path.exists(model_path):
- print(f"❌ Model {model_path} tidak ditemukan! Harap unduh model terlebih dahulu.")
- return
- # Buat direktori training unik berdasarkan timestamp
- timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
- save_dir = f"runs/detect/train_{timestamp}"
- os.makedirs(save_dir, exist_ok=True)
- print("📂 Memuat model YOLO untuk Few-Shot Learning...")
- model = YOLO(model_path)
- print("🚀 Memulai training YOLO Few-Shot Learning...")
- # Training dengan strategi Few-Shot Learning
- model.train(
- data=dataset_path,
- epochs=50, # Kurangi epoch untuk menghindari overfitting
- imgsz=512, # Resolusi gambar
- batch=16, # Ukuran batch kecil untuk dataset kecil
- device="cuda" if torch.cuda.is_available() else "cpu", # Gunakan GPU jika tersedia
- workers=0, # Worker kecil karena dataset kecil
- project="runs/detect",
- name=f"train_{timestamp}",
- amp=True, # Mixed Precision Training (menghemat memori dan mempercepat training)
- exist_ok=True,
- save_period=5, # Simpan checkpoint lebih sering (setiap 5 epoch)
- freeze=10, # Freeze 10 layer pertama agar hanya head layer yang belajar
- patience=10, # Early stopping jika tidak ada peningkatan dalam 10 epoch
- augment=True, # Aktifkan augmentasi agar model lebih kuat
- )
- print(f"✅ Training Few-Shot Learning selesai! Model disimpan di {save_dir}/weights/best.pt")
- def evaluate_latest_model():
- """Evaluasi model terbaru secara otomatis."""
- latest_model = get_latest_model()
- if not latest_model:
- return
- print("📊 Evaluasi model Few-Shot Learning...")
- model = YOLO(latest_model)
- model.val(save_json=True) # Simpan hasil evaluasi dalam format JSON
- print("✅ Evaluasi selesai!")
- # Eksekusi training Few-Shot Learning
- train_models_few_shot()
- evaluate_latest_model()
Add Comment
Please, Sign In to add comment
