745a64b342
- VegaModel CNN-FCNN 34.5M params, 82 isotopes, val acc 99.89% - Generation 50k spectres synthetiques 1D (12-24h durees) - Entrainement 100 epochs sur RTX 5060 Ti (CUDA 12.8, Blackwell) - Detection continue avec soustraction du background - Capture background 24h avec gestion deconnexion - Docker Compose : conteneur train (GPU) + detect (CPU/USB) - Modele entraite inclus (vega_best.pt, 395 Mo) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
121 lines
3.0 KiB
Python
121 lines
3.0 KiB
Python
#!/usr/bin/env python
|
|
"""
|
|
Run Vega Training
|
|
|
|
Simple script to train the Vega model on synthetic gamma spectra.
|
|
Designed for both quick test runs and full-scale training.
|
|
"""
|
|
|
|
import sys
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
# Add project root to path
|
|
PROJECT_ROOT = Path(__file__).parent.parent.parent
|
|
sys.path.insert(0, str(PROJECT_ROOT))
|
|
|
|
from training.vega.train import train_vega, TrainingConfig
|
|
from training.vega.model import VegaConfig
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Train Vega model for isotope identification"
|
|
)
|
|
|
|
# Data paths
|
|
parser.add_argument(
|
|
"--data-dir", "-d",
|
|
type=str,
|
|
default="O:/master_data_collection/isotopev2",
|
|
help="Path to synthetic data directory"
|
|
)
|
|
parser.add_argument(
|
|
"--model-dir", "-m",
|
|
type=str,
|
|
default="models",
|
|
help="Directory to save trained models"
|
|
)
|
|
|
|
# Training parameters
|
|
parser.add_argument(
|
|
"--epochs", "-e",
|
|
type=int,
|
|
default=100,
|
|
help="Maximum number of training epochs"
|
|
)
|
|
parser.add_argument(
|
|
"--batch-size", "-b",
|
|
type=int,
|
|
default=64,
|
|
help="Batch size for training (default: 64 for better GPU utilization)"
|
|
)
|
|
parser.add_argument(
|
|
"--learning-rate", "-lr",
|
|
type=float,
|
|
default=1e-3,
|
|
help="Initial learning rate"
|
|
)
|
|
|
|
# Quick test mode
|
|
parser.add_argument(
|
|
"--test",
|
|
action="store_true",
|
|
help="Quick test mode with reduced epochs"
|
|
)
|
|
|
|
# Mixed precision
|
|
parser.add_argument(
|
|
"--no-amp",
|
|
action="store_true",
|
|
help="Disable automatic mixed precision training"
|
|
)
|
|
|
|
# Data loading parallelism
|
|
parser.add_argument(
|
|
"--workers", "-w",
|
|
type=int,
|
|
default=8,
|
|
help="Number of data loading workers (default: 8 for parallel I/O)"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Create training config
|
|
config = TrainingConfig(
|
|
data_dir=args.data_dir,
|
|
model_dir=args.model_dir,
|
|
batch_size=args.batch_size,
|
|
learning_rate=args.learning_rate,
|
|
num_epochs=args.epochs if not args.test else 5,
|
|
patience=10 if not args.test else 3,
|
|
use_amp=not args.no_amp,
|
|
num_workers=args.workers
|
|
)
|
|
|
|
# Create model config
|
|
model_config = VegaConfig()
|
|
|
|
print("\n" + "=" * 60)
|
|
print("VEGA TRAINING")
|
|
print("=" * 60)
|
|
print(f"Data directory: {args.data_dir}")
|
|
print(f"Model directory: {args.model_dir}")
|
|
print(f"Epochs: {config.num_epochs}")
|
|
print(f"Batch size: {config.batch_size}")
|
|
print(f"Learning rate: {config.learning_rate}")
|
|
print(f"Mixed precision: {config.use_amp}")
|
|
print(f"Data workers: {config.num_workers}")
|
|
if args.test:
|
|
print("MODE: Quick test run")
|
|
print("=" * 60 + "\n")
|
|
|
|
# Run training
|
|
model, results = train_vega(config=config, model_config=model_config)
|
|
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|