99 lines
3.0 KiB
Python
99 lines
3.0 KiB
Python
# scripts/train_yolo.py
|
|
"""
|
|
Обучение YOLOv8 на детекцию.
|
|
Принимает параметры через командную строку.
|
|
"""
|
|
|
|
import sys
|
|
import time
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from ultralytics import YOLO
|
|
|
|
|
|
def check_environment():
|
|
print("=" * 60)
|
|
print("ПРОВЕРКА ОКРУЖЕНИЯ (DETECT)")
|
|
print("=" * 60)
|
|
print(f"Python: {sys.version.split()[0]}")
|
|
print(f"PyTorch: {torch.__version__}")
|
|
print(f"CUDA доступна: {torch.cuda.is_available()}")
|
|
if torch.cuda.is_available():
|
|
print(f"GPU: {torch.cuda.get_device_name(0)}")
|
|
print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
|
|
print("=" * 60)
|
|
|
|
|
|
def train(data_path: str, model: str, epochs: int, imgsz: int,
|
|
batch: int, device: str, project: str, name: str):
|
|
check_environment()
|
|
|
|
log_dir = Path(project) / name
|
|
log_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
print(f"\nЗагрузка модели: {model}")
|
|
yolo = YOLO(model)
|
|
|
|
args = {
|
|
'task': 'detect',
|
|
'data': data_path,
|
|
'epochs': epochs,
|
|
'imgsz': imgsz,
|
|
'batch': batch,
|
|
'device': device,
|
|
'project': project,
|
|
'name': name,
|
|
'exist_ok': True,
|
|
'pretrained': True,
|
|
'optimizer': 'auto',
|
|
'lr0': 0.01,
|
|
'lrf': 0.01,
|
|
'momentum': 0.937,
|
|
'weight_decay': 0.0005,
|
|
'warmup_epochs': 3.0,
|
|
'warmup_momentum': 0.8,
|
|
'warmup_bias_lr': 0.1,
|
|
'box': 7.5,
|
|
'cls': 0.5,
|
|
'dfl': 1.5,
|
|
'patience': 20,
|
|
'save': True,
|
|
'save_period': 10,
|
|
'plots': True,
|
|
'verbose': True,
|
|
}
|
|
|
|
print(f"\nЗапуск обучения DETECT...")
|
|
print(f"Датасет: {data_path}")
|
|
print(f"Модель: {model}")
|
|
print(f"Эпохи: {epochs}, Batch: {batch}, imgsz: {imgsz}")
|
|
print(f"Результаты: {log_dir}")
|
|
print("=" * 60)
|
|
|
|
t0 = time.time()
|
|
results = yolo.train(**args)
|
|
print(f"\nВремя обучения: {(time.time() - t0) / 60:.1f} мин")
|
|
print(f"Лучшая модель: {log_dir / 'weights' / 'best.pt'}")
|
|
return results
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='Обучение YOLOv8 detect')
|
|
parser.add_argument('--data', type=str, required=True, help='Путь к data.yaml')
|
|
parser.add_argument('--model', type=str, default='yolov8s.pt')
|
|
parser.add_argument('--epochs', type=int, default=100)
|
|
parser.add_argument('--imgsz', type=int, default=512)
|
|
parser.add_argument('--batch', type=int, default=16)
|
|
parser.add_argument('--device', type=str, default='0')
|
|
parser.add_argument('--project', type=str, default='runs/detect')
|
|
parser.add_argument('--name', type=str, default='experiment')
|
|
args = parser.parse_args()
|
|
|
|
train(args.data, args.model, args.epochs, args.imgsz,
|
|
args.batch, args.device, args.project, args.name)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main() |