train_model.py - master - louintik/ML

train_model.py
295 строк10 KB
louisa
first_commit
14 май 2026, 22:04
14 май 2026, 22:0451ac654
"""
Обучение модели для классификации поверхностей.
Поддерживает все доступные модели через аргумент --model.
"""
import argparse
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler
import joblib
import os
import warnings

warnings.filterwarnings('ignore')

# Конфигурация моделей
MODEL_CONFIGS = {
    'random_forest': {
        'class': 'RandomForestClassifier',
        'module': 'sklearn.ensemble',
        'params': {
            'n_estimators': 100,
            'max_depth': 10,
            'random_state': 42,
            'class_weight': 'balanced'
        },
        'needs_scaling': False
    },
    'extra_trees': {
        'class': 'ExtraTreesClassifier',
        'module': 'sklearn.ensemble',
        'params': {
            'n_estimators': 100,
            'max_depth': 10,
            'random_state': 42,
            'class_weight': 'balanced'
        },
        'needs_scaling': False
    },
    'svm': {
        'class': 'SVC',
        'module': 'sklearn.svm',
        'params': {
            'kernel': 'rbf',
            'C': 1.0,
            'random_state': 42,
            'class_weight': 'balanced',
            'probability': True
        },
        'needs_scaling': True
    },
    'logistic_regression': {
        'class': 'LogisticRegression',
        'module': 'sklearn.linear_model',
        'params': {
            'random_state': 42,
            'max_iter': 1000,
            'class_weight': 'balanced'
        },
        'needs_scaling': True
    },
    'knn': {
        'class': 'KNeighborsClassifier',
        'module': 'sklearn.neighbors',
        'params': {
            'n_neighbors': 5,
            'weights': 'distance'
        },
        'needs_scaling': True
    },
    'mlp': {
        'class': 'MLPClassifier',
        'module': 'sklearn.neural_network',
        'params': {
            'hidden_layer_sizes': (100, 50),
            'max_iter': 500,
            'random_state': 42,
            'early_stopping': True
        },
        'needs_scaling': True
    },
    'xgboost': {
        'class': 'XGBClassifier',
        'module': 'xgboost',
        'params': {
            'n_estimators': 100,
            'max_depth': 6,
            'random_state': 42,
            'use_label_encoder': False,
            'eval_metric': 'logloss'
        },
        'needs_scaling': False,
        'optional': True
    },
    'lightgbm': {
        'class': 'LGBMClassifier',
        'module': 'lightgbm',
        'params': {
            'n_estimators': 100,
            'max_depth': 10,
            'random_state': 42,
            'class_weight': 'balanced',
            'verbose': -1
        },
        'needs_scaling': False,
        'optional': True
    }
}


def load_and_prepare_dataset(dataset_path='../data/processed/surface_dataset.csv'):
    """
    Загрузить датасет и подготовить данные для ML.
    """
    df = pd.read_csv(dataset_path)

    # Исключаем мета-колонки
    exclude_cols = ['file', 'press_type', 'surface_label', 'group', 'prefix_suffix_correlation']
    feature_cols = [col for col in df.columns if col not in exclude_cols]

    X = df[feature_cols]
    y = df['surface_label']

    print(f"Размер датасета: {X.shape}")
    print(f"Фичей: {len(feature_cols)}")
    print(f"Классы: {dict(zip(*np.unique(y, return_counts=True)))}")
    print(f"  0 (hard_surface): {(y == 0).sum()}")
    print(f"  1 (elastic_surface): {(y == 1).sum()}")

    return X, y, feature_cols


def create_model(model_name):
    """
    Создать экземпляр модели по имени.
    """
    if model_name not in MODEL_CONFIGS:
        print(f"❌ Модель '{model_name}' не найдена!")
        print(f"Доступные модели: {', '.join(MODEL_CONFIGS.keys())}")
        return None

    config = MODEL_CONFIGS[model_name]

    try:
        module = __import__(config['module'], fromlist=[config['class']])
        model_class = getattr(module, config['class'])
        model = model_class(**config['params'])
        return model
    except ImportError as e:
        print(f"❌ Не удалось импортировать {config['module']}!")
        print(f"Установите: pip install {config['module']}")
        return None
    except Exception as e:
        # Ловим ошибки загрузки нативных библиотек (например, libomp для LightGBM)
        if 'dlopen' in str(e).lower() or 'library not loaded' in str(e).lower():
            print(f"❌ Ошибка загрузки библиотеки {config['module']}!")
            print(f"Причина: {str(e).split(chr(10))[0]}")
            print()
            if 'libomp' in str(e).lower():
                print("Для LightGBM на macOS требуется OpenMP:")
                print("  brew install libomp")
            return None
        raise


def train_model(model_name, X, y, test_size=0.2, output_dir='../models'):
    """
    Обучить модель и сохранить её.
    """
    config = MODEL_CONFIGS[model_name]
    model = create_model(model_name)

    if model is None:
        return None, None, None

    # Заполняем NaN медианным значением
    imputer = SimpleImputer(strategy='median')
    X_imputed = imputer.fit_transform(X)

    # Разделяем на train/test
    X_train, X_test, y_train, y_test = train_test_split(
        X_imputed, y, test_size=test_size, random_state=42, stratify=y
    )

    # Масштабируем если нужно
    scaler = None
    if config['needs_scaling']:
        scaler = StandardScaler()
        X_train = scaler.fit_transform(X_train)
        X_test = scaler.transform(X_test)

    # Обучаем модель
    display_name = model_name.replace('_', ' ').title()
    print("\n" + "=" * 80)
    print(f"ОБУЧЕНИЕ МОДЕЛИ: {display_name}")
    print("=" * 80)

    model.fit(X_train, y_train)

    # Предсказания
    y_pred = model.predict(X_test)

    # Оценка
    acc = accuracy_score(y_test, y_pred)
    cm = confusion_matrix(y_test, y_pred)
    cr = classification_report(y_test, y_pred, target_names=['hard_surface', 'elastic_surface'])

    print(f"\nAccuracy: {acc:.3f}")
    print(f"\nConfusion Matrix:")
    print(cm)
    print(f"\nClassification Report:")
    print(cr)

    # Важность фич (для tree-based моделей)
    if hasattr(model, 'feature_importances_'):
        feature_importance = pd.DataFrame({
            'feature': X.columns,
            'importance': model.feature_importances_
        }).sort_values('importance', ascending=False)

        print(f"\nТоп-10 важных фич:")
        print(feature_importance.head(10).to_string(index=False))

    # Коэффициенты (для линейных моделей)
    elif hasattr(model, 'coef_'):
        coef_importance = pd.DataFrame({
            'feature': X.columns,
            'coefficient': model.coef_[0],
            'abs_coeff': np.abs(model.coef_[0])
        }).sort_values('abs_coeff', ascending=False)

        print(f"\nТоп-10 важных фич (по коэффициентам):")
        print(coef_importance.head(10).to_string(index=False))

    # Создаём директорию если не существует
    os.makedirs(output_dir, exist_ok=True)

    # Сохраняем модель
    joblib.dump(model, f'{output_dir}/{model_name}_model.pkl')
    joblib.dump(imputer, f'{output_dir}/{model_name}_imputer.pkl')
    joblib.dump(X.columns.tolist(), f'{output_dir}/{model_name}_features.pkl')

    if scaler is not None:
        joblib.dump(scaler, f'{output_dir}/{model_name}_scaler.pkl')

    print(f"\n{'=' * 80}")
    print(f"✓ Модель сохранена:")
    print(f"  - {output_dir}/{model_name}_model.pkl")
    print(f"  - {output_dir}/{model_name}_imputer.pkl")
    print(f"  - {output_dir}/{model_name}_features.pkl")
    if scaler is not None:
        print(f"  - {output_dir}/{model_name}_scaler.pkl")
    print(f"{'=' * 80}")

    return model, imputer, scaler


if __name__ == '__main__':
    # Определяем корень проекта (где лежит этот скрипт)
    PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))
    
    # Парсим аргументы командной строки
    parser = argparse.ArgumentParser(description='Обучение модели для классификации поверхностей')
    parser.add_argument('--model', type=str, required=True,
                        help=f"Модель для обучения: {', '.join(MODEL_CONFIGS.keys())}")
    parser.add_argument('--dataset', type=str,
                        default=os.path.join(PROJECT_ROOT, 'data', 'processed', 'surface_dataset.csv'),
                        help="Путь к датасету")
    parser.add_argument('--output-dir', type=str,
                        default=os.path.join(PROJECT_ROOT, 'models'),
                        help="Директория для сохранения модели")
    parser.add_argument('--test-size', type=float, default=0.2,
                        help="Размер тестовой выборки (по умолчанию 0.2)")

    args = parser.parse_args()

    print("=" * 80)
    print("ЗАГРУЗКА ДАТАСЕТА")
    print("=" * 80)

    # Загружаем датасет
    X, y, feature_cols = load_and_prepare_dataset(args.dataset)

    # Обучаем модель
    model, imputer, scaler = train_model(args.model, X, y, args.test_size, args.output_dir)

    if model is not None:
        print(f"\n{'=' * 80}")
        print(f"✓ Обучение завершено!")
        print(f"{'=' * 80}")

        print("\nДля предсказания на новых данных запустите:")
        print(f"  python predict_surface.py --model {args.model}")