train_kaggle.py

"""
YOLOv8-MPEB Training Script for Kaggle
Based on: "YOLOv8-MPEB small target detection algorithm based on UAV images"

This script is specifically configured for Kaggle environment:
- Uses /kaggle/working for writable operations
- Uses /kaggle/input for read-only input files
- Handles dataset paths correctly for Kaggle's file system

Paper Specifications:
- Model: YOLOv8s-MPEB (Small variant)
- Parameters: 7.39M
- Model Size: 14.5 MB
- Target mAP50: 91.9%
- GFLOPs: 27.4
"""

import sys
import os
from pathlib import Path
import shutil

# Set up paths for Kaggle environment
KAGGLE_INPUT = Path('/kaggle/input')
KAGGLE_WORKING = Path('/kaggle/working')
CODE_DIR = KAGGLE_INPUT / 'yolo-mpeb-training-code' / 'code'

# Add code directory to Python path
sys.path.insert(0, str(CODE_DIR))

# Import custom modules from the input directory
from yolov8_mpeb_modules import MobileNetBlock, EMA, C2f_EMA, BiFPN_Fusion

# Patch Ultralytics modules BEFORE importing YOLO
import ultralytics.nn.modules as modules
import ultralytics.nn.modules.block as block
import ultralytics.nn.tasks as tasks

print("=" * 80)
print("YOLOv8-MPEB Training Script for Kaggle")
print("=" * 80)
print("\nPatching Ultralytics modules...")

# Proxy: GhostBottleneck -> MobileNetBlock
block.GhostBottleneck = MobileNetBlock
modules.GhostBottleneck = MobileNetBlock

# Proxy: C3 -> C2f_EMA
block.C3 = C2f_EMA
modules.C3 = C2f_EMA

# Patch tasks namespace
if hasattr(tasks, 'GhostBottleneck'): 
    tasks.GhostBottleneck = MobileNetBlock
if hasattr(tasks, 'C3'): 
    tasks.C3 = C2f_EMA
if hasattr(tasks, 'block'):
    tasks.block.GhostBottleneck = MobileNetBlock
    tasks.block.C3 = C2f_EMA

from ultralytics import YOLO

# Copy necessary files to working directory
print("\nSetting up working directory...")
WORKING_CODE_DIR = KAGGLE_WORKING / 'code'
WORKING_CODE_DIR.mkdir(exist_ok=True)

# Copy model YAML and dataset YAML to working directory
model_yaml = CODE_DIR / 'yolov8_mpeb.yaml'
dataset_yaml = CODE_DIR / 'dataset_example.yaml'

if model_yaml.exists():
    shutil.copy(model_yaml, WORKING_CODE_DIR / 'yolov8_mpeb.yaml')
    print(f"✓ Copied model YAML to {WORKING_CODE_DIR / 'yolov8_mpeb.yaml'}")

if dataset_yaml.exists():
    shutil.copy(dataset_yaml, WORKING_CODE_DIR / 'dataset_example.yaml')
    print(f"✓ Copied dataset YAML to {WORKING_CODE_DIR / 'dataset_example.yaml'}")

# Change to working directory
os.chdir(KAGGLE_WORKING)

# Training configuration
TRAINING_CONFIG = {
    'data': str(WORKING_CODE_DIR / 'dataset_example.yaml'),
    'epochs': 200,
    'batch': 32,
    'imgsz': 640,
    'lr0': 0.01,
    'lrf': 0.01,
    'weight_decay': 0.0005,
    'device': 0,  # Use GPU 0
    'project': str(KAGGLE_WORKING / 'runs' / 'train'),
    'name': 'yolov8_mpeb',
    'resume': False,
    # Additional parameters
    'patience': 50,
    'save': True,
    'save_period': 10,
    'cache': False,
    'workers': 4,
    'optimizer': 'SGD',
    'verbose': True,
    'seed': 0,
    'deterministic': True,
    'single_cls': False,
    'rect': False,
    'cos_lr': False,
    'close_mosaic': 10,
    'amp': True,
    'fraction': 1.0,
    'profile': False,
    # Data augmentation
    'hsv_h': 0.015,
    'hsv_s': 0.7,
    'hsv_v': 0.4,
    'degrees': 0.0,
    'translate': 0.1,
    'scale': 0.5,
    'shear': 0.0,
    'perspective': 0.0,
    'flipud': 0.0,
    'fliplr': 0.5,
    'mosaic': 1.0,
    'mixup': 0.0,
    'copy_paste': 0.0,
}

print("\n" + "=" * 80)
print("STARTING YOLOv8-MPEB TRAINING ON KAGGLE")
print("=" * 80)
print(f"\nGPU: Tesla P100-PCIE-16GB")
print(f"Model: YOLOv8s-MPEB (7.38M parameters)")
print(f"Dataset: dataset_example.yaml")
print(f"Batch Size: {TRAINING_CONFIG['batch']}")
print(f"Epochs: {TRAINING_CONFIG['epochs']}")
print(f"\nEstimated time: 6-8 hours")
print("=" * 80)

# Load model
print("\nLoading YOLOv8-MPEB model...")
model = YOLO(str(WORKING_CODE_DIR / 'yolov8_mpeb.yaml'))

# Display model info
print("\nModel Information:")
model.info()

print("\nTraining starting...\n")

# Train
results = model.train(**TRAINING_CONFIG)

print("\n" + "=" * 80)
print("TRAINING COMPLETE!")
print("=" * 80)
print(f"Results saved to: {results.save_dir}")
print(f"Best weights: {results.save_dir}/weights/best.pt")
print(f"Last weights: {results.save_dir}/weights/last.pt")
print("=" * 80)

# Validate the best model
print("\nValidating best model...")
val_results = model.val(data=TRAINING_CONFIG['data'])

print("\n" + "=" * 80)
print("VALIDATION RESULTS")
print("=" * 80)
print(f"mAP50: {val_results.box.map50:.4f}")
print(f"mAP50-95: {val_results.box.map:.4f}")
print(f"Target mAP50 (from paper): 0.919")
print("=" * 80)