diff options
Diffstat (limited to 'ml/model_trainer.py')
| -rw-r--r-- | ml/model_trainer.py | 393 |
1 files changed, 393 insertions, 0 deletions
diff --git a/ml/model_trainer.py b/ml/model_trainer.py new file mode 100644 index 0000000..ab3059f --- /dev/null +++ b/ml/model_trainer.py @@ -0,0 +1,393 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +from abc import ABC, abstractmethod +import datetime +import glob +import logging +import os +import pickle +import random +import sys +from types import SimpleNamespace +from typing import Any, List, NamedTuple, Optional, Set, Tuple + +import numpy as np +from sklearn.model_selection import train_test_split # type:ignore +from sklearn.preprocessing import MinMaxScaler # type: ignore + +from ansi import bold, reset +import argparse_utils +import config +from decorator_utils import timed +import parallelize as par + +logger = logging.getLogger(__file__) + +parser = config.add_commandline_args( + f"ML Model Trainer ({__file__})", + "Arguments related to training an ML model" +) +parser.add_argument( + "--ml_trainer_quiet", + action="store_true", + help="Don't prompt the user for anything." +) +parser.add_argument( + "--ml_trainer_delete", + action="store_true", + help="Delete invalid/incomplete features files in addition to warning." +) +group = parser.add_mutually_exclusive_group() +group.add_argument( + "--ml_trainer_dry_run", + action="store_true", + help="Do not write a new model, just report efficacy.", +) +group.add_argument( + "--ml_trainer_persist_threshold", + type=argparse_utils.valid_percentage, + metavar='0..100', + help="Persist the model if the test set score is >= this threshold.", +) + + +class InputSpec(SimpleNamespace): + file_glob: str + feature_count: int + features_to_skip: Set[str] + key_value_delimiter: str + training_parameters: List + label: str + basename: str + dry_run: Optional[bool] + quiet: Optional[bool] + persist_percentage_threshold: Optional[float] + delete_bad_inputs: Optional[bool] + + @staticmethod + def populate_from_config() -> InputSpec: + return InputSpec( + dry_run = config.config["ml_trainer_dry_run"], + quiet = config.config["ml_trainer_quiet"], + persist_percentage_threshold = config.config["ml_trainer_persist_threshold"], + delete_bad_inputs = config.config["ml_trainer_delete"], + ) + + +class OutputSpec(NamedTuple): + model_filename: Optional[str] + model_info_filename: Optional[str] + scaler_filename: Optional[str] + training_score: float + test_score: float + + +class TrainingBlueprint(ABC): + def __init__(self): + self.y_train = None + self.y_test = None + self.X_test_scaled = None + self.X_train_scaled = None + self.file_done_count = 0 + self.total_file_count = 0 + self.spec = None + + def train(self, spec: InputSpec) -> OutputSpec: + import smart_future + + random.seed() + self.spec = spec + + X_, y_ = self.read_input_files() + num_examples = len(y_) + + # Every example's features + X = np.array(X_) + + # Every example's label + y = np.array(y_) + + print("Doing random test/train split...") + X_train, X_test, self.y_train, self.y_test = self.test_train_split( + X, + y, + ) + + print("Scaling training data...") + scaler, self.X_train_scaled, self.X_test_scaled = self.scale_data( + X_train, + X_test, + ) + + print("Training model(s)...") + models = [] + modelid_to_params = {} + for params in self.spec.training_parameters: + model = self.train_model( + params, + self.X_train_scaled, + self.y_train + ) + models.append(model) + modelid_to_params[model.get_id()] = str(params) + + best_model = None + best_score = None + best_test_score = None + best_training_score = None + best_params = None + for model in smart_future.wait_any(models): + params = modelid_to_params[model.get_id()] + if isinstance(model, smart_future.SmartFuture): + model = model._resolve() + if model is not None: + training_score, test_score = self.evaluate_model( + model, + self.X_train_scaled, + self.y_train, + self.X_test_scaled, + self.y_test, + ) + score = (training_score + test_score * 20) / 21 + if not self.spec.quiet: + print( + f"{bold()}{params}{reset()}: " + f"Training set score={training_score:.2f}%, " + f"test set score={test_score:.2f}%", + file=sys.stderr, + ) + if best_score is None or score > best_score: + best_score = score + best_test_score = test_score + best_training_score = training_score + best_model = model + best_params = params + if not self.spec.quiet: + print( + f"New best score {best_score:.2f}% with params {params}" + ) + + if not self.spec.quiet: + msg = f"Done training; best test set score was: {best_test_score:.1f}%" + print(msg) + logger.info(msg) + scaler_filename, model_filename, model_info_filename = ( + self.maybe_persist_scaler_and_model( + best_training_score, + best_test_score, + best_params, + num_examples, + scaler, + best_model, + ) + ) + return OutputSpec( + model_filename = model_filename, + model_info_filename = model_info_filename, + scaler_filename = scaler_filename, + training_score = best_training_score, + test_score = best_test_score, + ) + + @par.parallelize(method=par.Method.THREAD) + def read_files_from_list( + self, + files: List[str], + n: int + ) -> Tuple[List, List]: + # All features + X = [] + + # The label + y = [] + + for filename in files: + wrote_label = False + with open(filename, "r") as f: + lines = f.readlines() + + # This example's features + x = [] + for line in lines: + + # We expect lines in features files to be of the form: + # + # key: value + line = line.strip() + try: + (key, value) = line.split(self.spec.key_value_delimiter) + except Exception as e: + logger.debug(f"WARNING: bad line in file {filename} '{line}', skipped") + continue + + key = key.strip() + value = value.strip() + if (self.spec.features_to_skip is not None + and key in self.spec.features_to_skip): + logger.debug(f"Skipping feature {key}") + continue + + value = self.normalize_feature(value) + + if key == self.spec.label: + y.append(value) + wrote_label = True + else: + x.append(value) + + # Make sure we saw a label and the requisite number of features. + if len(x) == self.spec.feature_count and wrote_label: + X.append(x) + self.file_done_count += 1 + else: + if wrote_label: + y.pop() + + if self.spec.delete_bad_inputs: + msg = f"WARNING: {filename}: missing features or label. DELETING." + print(msg, file=sys.stderr) + logger.warning(msg) + os.remove(filename) + else: + msg = f"WARNING: {filename}: missing features or label. Skipped." + print(msg, file=sys.stderr) + logger.warning(msg) + return (X, y) + + def make_progress_graph(self) -> None: + if not self.spec.quiet: + from text_utils import progress_graph + progress_graph( + self.file_done_count, + self.total_file_count + ) + + @timed + def read_input_files(self): + import list_utils + import smart_future + + # All features + X = [] + + # The label + y = [] + + results = [] + all_files = glob.glob(self.spec.file_glob) + self.total_file_count = len(all_files) + for n, files in enumerate(list_utils.shard(all_files, 500)): + file_list = list(files) + results.append(self.read_files_from_list(file_list, n)) + + for result in smart_future.wait_any(results, callback=self.make_progress_graph): + result = result._resolve() + for z in result[0]: + X.append(z) + for z in result[1]: + y.append(z) + if not self.spec.quiet: + print(" " * 80 + "\n") + return (X, y) + + def normalize_feature(self, value: str) -> Any: + if value in ("False", "None"): + ret = 0 + elif value == "True": + ret = 255 + elif isinstance(value, str) and "." in value: + ret = round(float(value) * 100.0) + else: + ret = int(value) + return ret + + def test_train_split(self, X, y) -> List: + logger.debug("Performing test/train split") + return train_test_split( + X, + y, + random_state=random.randrange(0, 1000), + ) + + def scale_data(self, + X_train: np.ndarray, + X_test: np.ndarray) -> Tuple[Any, np.ndarray, np.ndarray]: + logger.debug("Scaling data") + scaler = MinMaxScaler() + scaler.fit(X_train) + return (scaler, scaler.transform(X_train), scaler.transform(X_test)) + + # Note: children should implement. Consider using @parallelize. + @abstractmethod + def train_model(self, + parameters, + X_train_scaled: np.ndarray, + y_train: np.ndarray) -> Any: + pass + + def evaluate_model( + self, + model: Any, + X_train_scaled: np.ndarray, + y_train: np.ndarray, + X_test_scaled: np.ndarray, + y_test: np.ndarray) -> Tuple[np.float64, np.float64]: + logger.debug("Evaluating the model") + training_score = model.score(X_train_scaled, y_train) * 100.0 + test_score = model.score(X_test_scaled, y_test) * 100.0 + logger.info( + f"Model evaluation results: test_score={test_score:.5f}, " + f"train_score={training_score:.5f}" + ) + return (training_score, test_score) + + def maybe_persist_scaler_and_model( + self, + training_score: np.float64, + test_score: np.float64, + params: str, + num_examples: int, + scaler: Any, + model: Any) -> Tuple[Optional[str], Optional[str], Optional[str]]: + if not self.spec.dry_run: + import datetime_utils + import input_utils + import string_utils + + if ( + (self.spec.persist_percentage_threshold is not None and + test_score > self.spec.persist_percentage_threshold) + or + (not self.spec.quiet + and input_utils.yn_response("Write the model? [y,n]: ") == "y") + ): + scaler_filename = f"{self.spec.basename}_scaler.sav" + with open(scaler_filename, "wb") as f: + pickle.dump(scaler, f) + msg = f"Wrote {scaler_filename}" + print(msg) + logger.info(msg) + model_filename = f"{self.spec.basename}_model.sav" + with open(model_filename, "wb") as f: + pickle.dump(model, f) + msg = f"Wrote {model_filename}" + print(msg) + logger.info(msg) + model_info_filename = f"{self.spec.basename}_model_info.txt" + now: datetime.datetime = datetime_utils.now_pst() + info = f"""Timestamp: {datetime_utils.datetime_to_string(now)} +Model params: {params} +Training examples: {num_examples} +Training set score: {training_score:.2f}% +Testing set score: {test_score:.2f}%""" + with open(model_info_filename, "w") as f: + f.write(info) + msg = f"Wrote {model_info_filename}:" + print(msg) + logger.info(msg) + print(string_utils.indent(info, 2)) + logger.info(info) + return (scaler_filename, model_filename, model_info_filename) + return (None, None, None) |