diff --git a/README.md b/README.md index ef65200..2b35294 100644 --- a/README.md +++ b/README.md @@ -133,12 +133,16 @@ Use the `dataset` command: This command collates all logs of dimension N+1 from the logs directory and saves the data in a compressed `.npz` file in the datasets directory. It will also print out some information about your dataset, in particular the total number of individual interactions. To have a useful dataset, it's good to start with more than 10,000 individual interactions but YMMV. -To train the model, use the `train` command---this can take a while on a normal computer, so be prepared to let your computer sit and think for a few hours! You'll have to decide what _size_ model to try to train: `xxs`, `xs`, `s`, `m`, `l`, `xl`. The size refers to the number of LSTM units in each layer of your model and roughly corresponds to "learning capacity" at a cost of slower training and predictions. +To train the model, use the `train` command. You don't need a GPU: an `s` model trained on a laptop CPU with around 10,000 interactions takes a few minutes. You'll have to decide what _size_ model to try to train: `xxs`, `xs`, `s`, `m`, `l`, `xl`. The size refers to the number of LSTM units in each layer of your model and roughly corresponds to "learning capacity" at a cost of slower training and predictions. It's a good idea to start with an `xs` or `s` model, and the larger models may work better for quite large datasets (e.g., >1M individual interactions). - poetry run ./start_impsy.py train --dimension (N+1) --modelsize s + poetry run ./start_impsy.py train datasets/training-dataset-(N+1)d.npz --modelsize s -Early stopping is enabled by default and will stop training after the model stops improving for 10 epochs (configurable with `--patience`). You can disable it with `--no-earlystopping`. +The model's dimension is read from the dataset. You can also leave out the dataset file and give `--dimension (N+1)` to use `datasets/training-dataset-(N+1)d.npz`. + +Early stopping is enabled by default and will stop training after the model stops improving for 10 epochs (configurable with `--patience`). You can disable it with `--no-earlystopping`. If you press `Ctrl-C` during training, IMPSY stops and saves the model as it is. + +You can also upload logs, make a dataset and train a model from the web UI (`impsy webui`, see the Train page) on whatever device is running IMPSY, e.g., a laptop running the Docker image or a Raspberry Pi. A Raspberry Pi 4 or 5 should be able to train small models, but we wouldn't recommend training on a Pi Zero 2 W. By default, your trained model will be saved in the `models` directory in `.keras` and `.tflite` format. diff --git a/docs/impsypi-workshop.md b/docs/impsypi-workshop.md index 7f6de25..225fd0e 100644 --- a/docs/impsypi-workshop.md +++ b/docs/impsypi-workshop.md @@ -123,7 +123,7 @@ Finally, the SD card is configured to allow ethernet over USB so you can plug yo Connecting to the Raspberry Pi should be straightforward: if your computer knows about it over USB-ethernet, you can just type into a web browser and see the web UI. The Pi gives itself the IP address 169.254.1.107 so you might try as well if the above doesn't work. -Now you can poke around the web UI and see what you can do. Most importantly, you can update the configuration file directly, download log files and generate or download datasets. You can't train new models from the web UI (yet!) and it's possibly a bad idea to train models on the slower Raspberry Pis anyway. +Now you can poke around the web UI and see what you can do. Most importantly, you can update the configuration file directly, upload or download log files and generate or download datasets. The web UI can also train models on the Train page. A Raspberry Pi 4 or 5 should be able to train a small model (try `xs` or `s`), but we wouldn't recommend training on a Pi Zero 2 W; instead, train on your laptop (e.g., with the Docker image) and upload the `.tflite` model to your Pi. > N.B.: I'm new to this whole ethernet over USB and SD card building business. I'm probably doing something wrong so if there are issues here please be patient and if you know a better way to set something up, let me know! @@ -308,7 +308,7 @@ docker run --rm \ --volume ./datasets:/impsy/datasets \ --volume ./models:/impsy/models \ charlepm/impsy \ - poetry run ./start_impsy.py train -D 4 -S datasets/training-dataset-4d.npz + poetry run ./start_impsy.py train datasets/training-dataset-4d.npz ``` to train a 4D model. This command makes sure that the model files that are saved by the train command get saved out of the docker container and in your regular computer filesystem. diff --git a/impsy/dataset.py b/impsy/dataset.py index aef8875..5418135 100755 --- a/impsy/dataset.py +++ b/impsy/dataset.py @@ -70,6 +70,7 @@ def generate_dataset( # Save Performance Data in a compressed numpy file. dataset_name = filename if filename else f"training-dataset-{dimension}d.npz" + Path(destination).mkdir(parents=True, exist_ok=True) dataset_file = Path(destination) / dataset_name # Input format is: diff --git a/impsy/mdrnn.py b/impsy/mdrnn.py index 5b64340..ce95c5d 100644 --- a/impsy/mdrnn.py +++ b/impsy/mdrnn.py @@ -227,8 +227,11 @@ def train( save_location="models", validation_split=0.1, patience=10, + callbacks=None, ): - """Train the network for a number of epochs with a specific dataset.""" + """Train the network for a number of epochs with a specific dataset. + + Extra Keras callbacks can be given in callbacks.""" save_location = Path(save_location) checkpoint_path = save_location / f"{self.model_name}-ckpt.keras" checkpoint_callback = tf.keras.callbacks.ModelCheckpoint( @@ -242,11 +245,13 @@ def train( early_stopping_callback = tf.keras.callbacks.EarlyStopping( monitor="val_loss", mode="min", verbose=1, patience=patience ) + extra_callbacks = list(callbacks or []) callbacks = [terminateOnNaN] if checkpointing: callbacks.append(checkpoint_callback) if early_stopping: callbacks.append(early_stopping_callback) + callbacks += extra_callbacks # Do the data scaling in here. X = np.array(X, dtype=np.float32) * SCALE_FACTOR diff --git a/impsy/templates/base.html b/impsy/templates/base.html index ba3a03c..4e3e8d9 100644 --- a/impsy/templates/base.html +++ b/impsy/templates/base.html @@ -34,6 +34,9 @@ Datasets {% if workflow.datasets %}✓{% endif %} +
  • + Train +
  • Models {% if workflow.models %}✓{% endif %} diff --git a/impsy/templates/datasets.html b/impsy/templates/datasets.html index 6e8c5d1..41d7508 100644 --- a/impsy/templates/datasets.html +++ b/impsy/templates/datasets.html @@ -31,6 +31,7 @@
    Dataset Files
    File Size Date + Dim Actions @@ -43,7 +44,9 @@
    Dataset Files
    {{ file.size }} {{ file.date }} + {% if file.dimension %}{{ file.dimension }}d{% else %}?{% endif %} +     Train Download
    diff --git a/impsy/templates/logs.html b/impsy/templates/logs.html index 952245e..1e91f71 100644 --- a/impsy/templates/logs.html +++ b/impsy/templates/logs.html @@ -2,7 +2,23 @@ {% block title %}IMPSY - Log Files{% endblock %} {% block content %}

    Log Files

    -

    Interaction logs recorded during performances. Run impsy run to record new data.

    +

    Interaction logs recorded during performances. Run impsy run to record new data, or upload logs from another IMPSY device.

    + +
    +
    +
    Upload Logs
    + +
    + + +
    +
    + +
    + +
    Log files from IMPSY, IMPSYpi, the AUv3 app and IMPSY Web all work. Names need to end in -{dimension}d-mdrnn.log.
    +
    +
    diff --git a/impsy/templates/models.html b/impsy/templates/models.html index 721a971..7ee9682 100644 --- a/impsy/templates/models.html +++ b/impsy/templates/models.html @@ -2,7 +2,7 @@ {% block title %}IMPSY - Models{% endblock %} {% block content %}

    Models

    -

    Upload trained models or train new ones with impsy train.

    +

    Upload trained models or train new ones on the Train page or with impsy train.

    diff --git a/impsy/templates/train.html b/impsy/templates/train.html new file mode 100644 index 0000000..0340689 --- /dev/null +++ b/impsy/templates/train.html @@ -0,0 +1,126 @@ +{% extends "base.html" %} +{% block title %}IMPSY - Train{% endblock %} +{% block content %} +

    Train

    +

    Train a model from a dataset. Training runs in the background and saves .keras and .tflite files to the models folder. A small model on a laptop CPU takes a few minutes.

    + +
    +
    +
    Train a Model
    + {% if dataset_files %} +
    +
    + + +
    +
    + + +
    +
    + + +
    +
    + + +
    +
    + +
    +
    +
    + Training stops early when the validation loss hasn't improved for patience epochs. + {% if config_dimension %}Your config uses dimension {{ config_dimension }}; the model's dimension comes from the dataset.{% endif %} +
    + {% else %} +

    No datasets yet. Generate a dataset from your log files first.

    + {% endif %} +
    +
    + +
    +
    +
    +
    Training Status
    +
    + +
    +
    +
    +
    Status
    {{ job.status }}
    +
    Dataset
    {{ job.dataset }}
    +
    Model size
    {{ job.model_size }}
    +
    Epoch
    {{ job.epoch }} / {{ job.max_epochs }}
    +
    Loss (val)
    –
    +
    Time
    {{ job.elapsed }} s
    +
    +
    {{ job.error }}
    +
    + {% for name in job.model_files %} + Download {{ name }} + {% endfor %} +
    +
    +
    +{% endblock %} + +{% block scripts %} + +{% endblock %} diff --git a/impsy/train.py b/impsy/train.py index ae9d4b2..f7d5729 100755 --- a/impsy/train.py +++ b/impsy/train.py @@ -46,8 +46,38 @@ def seq_to_singleton_format(examples): return (xs, ys) +def dataset_dimension(dataset_file) -> int: + """Returns the dimension (dt plus number of values) of the data in a .npz dataset.""" + with np.load(dataset_file, allow_pickle=True) as loaded: + corpus = loaded["perfs"] + for perf in corpus: + if len(perf) > 0: + return int(np.asarray(perf).shape[1]) + raise ValueError(f"Dataset {dataset_file} has no data.") + + +def find_dataset_file(dataset_location, dimension: int | None = None) -> Path: + """Resolves a dataset file from a .npz path or a directory of datasets. + + For a directory, uses training-dataset-{dimension}d.npz if dimension is given, + otherwise the only .npz file in the directory. + """ + dataset_location = Path(dataset_location) + if dataset_location.suffix == ".npz": + return dataset_location + if dimension is not None: + return dataset_location / f"training-dataset-{dimension}d.npz" + candidates = sorted(dataset_location.glob("*.npz")) + if len(candidates) == 1: + return candidates[0] + raise ValueError( + f"Found {len(candidates)} .npz files in {dataset_location}; " + "give a dataset file or a dimension to choose one." + ) + + def train_mdrnn( - dimension: int, + dimension: int | None, dataset_location: str, model_size: str, early_stopping: bool, @@ -58,8 +88,14 @@ def train_mdrnn( save_model: bool = True, save_weights: bool = False, save_tflite: bool = True, + callbacks: list | None = None, ): - """Loads a dataset, creates a model and runs the training procedure.""" + """Loads a dataset, creates a model and runs the training procedure. + + dimension can be None to read it from the dataset. If training is interrupted + (e.g., with Ctrl-C) the model is still saved with the latest weights. + Extra Keras callbacks can be given to follow or stop training. + """ from . import mdrnn model_config = mdrnn_config(model_size) @@ -78,13 +114,16 @@ def train_mdrnn( np.random.seed(SEED) # Load dataset - dataset_location = Path(dataset_location) - dataset_default_name = f"training-dataset-{str(dimension)}d.npz" - if dataset_location.suffix == "": - dataset_default_name = f"training-dataset-{str(dimension)}d.npz" - dataset_location = dataset_location / dataset_default_name - assert dataset_location.suffix == ".npz", "dataset file to load must end with .npz" + dataset_location = find_dataset_file(dataset_location, dimension) click.secho(f"Dataset: {dataset_location}") + data_dimension = dataset_dimension(dataset_location) + if dimension is None: + dimension = data_dimension + elif dimension != data_dimension: + raise ValueError( + f"Dataset has dimension {data_dimension} but dimension {dimension} was requested." + ) + click.secho(f"Dimension: {dimension}", fg="blue") with np.load(dataset_location, allow_pickle=True) as loaded: corpus = loaded["perfs"] print("Loaded performances:", len(corpus)) @@ -93,6 +132,11 @@ def train_mdrnn( # Restrict corpus to performances longer than the training sequence length. corpus = [l for l in corpus if len(l) > SEQ_LEN + 1] click.secho(f"Corpus Examples: {len(corpus)}", fg="blue") + if not corpus: + raise ValueError( + f"No performances in the dataset are longer than {SEQ_LEN + 1} events, " + "so there's nothing to train on. Record some longer logs." + ) # Prepare training data as X and Y. slices = [] @@ -121,17 +165,22 @@ def train_mdrnn( ) validation_split = 0.10 - history = training_mdrnn.train( - X, - y, - batch_size=batch_size, - epochs=num_epochs, - checkpointing=True, - early_stopping=early_stopping, - save_location=save_location, - validation_split=validation_split, - patience=patience, - ) + try: + history = training_mdrnn.train( + X, + y, + batch_size=batch_size, + epochs=num_epochs, + checkpointing=True, + early_stopping=early_stopping, + save_location=save_location, + validation_split=validation_split, + patience=patience, + callbacks=callbacks, + ) + except KeyboardInterrupt: + click.secho("Training interrupted, saving the latest weights.", fg="yellow") + history = getattr(training_mdrnn.model, "history", None) # Save final Model model_name = training_mdrnn.model_name @@ -162,19 +211,23 @@ def train_mdrnn( # Save .tflite file from .tflite_converter import model_to_tflite - tflite_file = model_to_tflite(inference_mdrnn.model, model_keras_file) + if not save_model: + inference_mdrnn.model.set_weights(training_mdrnn.model.get_weights()) + tflite_path = save_location / f"{inference_mdrnn.model_name}.tflite" + tflite_file = model_to_tflite(inference_mdrnn.model, tflite_path) output["tflite_file"] = tflite_file return output @click.command(name="train") +@click.argument("dataset", required=False, default=None) @click.option( "-D", "--dimension", type=int, - default=2, - help="The dimension of the data to model, must be >= 2.", + default=None, + help="The dimension of the data to model, must be >= 2 (read from the dataset if not given).", ) @click.option( "-S", @@ -218,7 +271,8 @@ def train_mdrnn( help="The destination directory to write trained model files to.", ) def train( - dimension: int, + dataset: str | None, + dimension: int | None, source: str, modelsize: str, earlystopping: bool, @@ -227,12 +281,20 @@ def train( batchsize: int, destination: str, ): - """Trains an IMPSY MDRNN model based on an existing dataset (run dataset command first!).""" + """Trains an IMPSY MDRNN model based on an existing dataset (run dataset command first!). + + DATASET is an optional .npz file to train on (overrides --source). The trained + model is saved in .keras and .tflite formats. Press Ctrl-C to stop training + early and save the model as it is. + """ + if dataset is not None: + source = dataset click.secho( - f"IMPSY: Going to train a {dimension}D, {modelsize} sized MDRNN model.", + f"IMPSY: Going to train a {modelsize} sized MDRNN model.", fg="green", ) - train_mdrnn( + Path(destination).mkdir(parents=True, exist_ok=True) + output = train_mdrnn( dimension, source, modelsize, @@ -243,3 +305,4 @@ def train( save_location=destination, ) click.secho("IMPSY: training completed.", fg="green") + click.secho(f"Model file: {output['tflite_file']}", fg="green") diff --git a/impsy/web_interface.py b/impsy/web_interface.py index 6b50bbf..945ec38 100644 --- a/impsy/web_interface.py +++ b/impsy/web_interface.py @@ -5,12 +5,14 @@ import shutil import platform import os +import re import time import tomllib from importlib.metadata import PackageNotFoundError, metadata as _pkg_metadata, version as _pkg_version from threading import Lock, Thread from impsy.data import default_config_template from impsy.dataset import generate_dataset +from impsy.utils import SIZE_TO_PARAMETERS from pathlib import Path from pythonosc import dispatcher, osc_server from datetime import datetime @@ -302,6 +304,161 @@ def allowed_dataset_file(filename): return "." in filename and filename.rsplit(".", 1)[1].lower() in {"npz"} +def log_file_dimension(filename): + """Returns the dimension in a log filename (*-{dimension}d-mdrnn.log), or None.""" + match = re.search(r"-(\d+)d-mdrnn\.log$", Path(filename).name) + return int(match.group(1)) if match else None + + +def get_dataset_file_info(filepath): + """Get metadata for a dataset file, including the dimension of its data. + + Uses the dimension in a standard training-dataset-{dimension}d.npz name, + otherwise loads the dataset to find it. + """ + info = get_file_info(filepath) + match = re.fullmatch(r"training-dataset-(\d+)d\.npz", filepath.name) + if match: + info["dimension"] = int(match.group(1)) + return info + try: + from impsy.train import dataset_dimension + + info["dimension"] = dataset_dimension(filepath) + except Exception: + info["dimension"] = None + return info + + +TRAIN_SIZES = list(SIZE_TO_PARAMETERS.keys()) +TRAIN_MAX_EPOCHS = 1000 + + +class TrainingJob: + """Runs one model training at a time in a background thread. + + Progress (epoch, losses, status) is kept here so the /train page can poll it. + Stopping ends training after the current batch and still saves the model. + """ + + def __init__(self): + self.status = "idle" # idle, running, stopping, finished, failed + self.dataset = None + self.model_size = None + self.max_epochs = 0 + self.epoch = 0 + self.loss = [] + self.val_loss = [] + self.model_files = [] + self.error = None + self.started_at = None + self.finished_at = None + self._stop_requested = False + self._thread = None + self._lock = Lock() + + @property + def running(self): + return self.status in ("running", "stopping") + + def start(self, dataset_file, model_size, max_epochs, patience, models_dir): + """Starts training in the background. Returns False if a job is already running.""" + with self._lock: + if self.running: + return False + self.status = "running" + self.dataset = Path(dataset_file).name + self.model_size = model_size + self.max_epochs = max_epochs + self.epoch = 0 + self.loss = [] + self.val_loss = [] + self.model_files = [] + self.error = None + self.started_at = time.time() + self.finished_at = None + self._stop_requested = False + self._thread = Thread( + target=self._run, + args=(dataset_file, model_size, max_epochs, patience, models_dir), + name="webui_training_thread", + daemon=True, + ) + self._thread.start() + return True + + def stop(self): + with self._lock: + if self.status == "running": + self._stop_requested = True + self.status = "stopping" + + def _run(self, dataset_file, model_size, max_epochs, patience, models_dir): + try: + import tensorflow as tf + from impsy.train import train_mdrnn + + job = self + + class ProgressCallback(tf.keras.callbacks.Callback): + def on_train_batch_end(self, batch, logs=None): + if job._stop_requested: + self.model.stop_training = True + + def on_epoch_end(self, epoch, logs=None): + logs = logs or {} + job.epoch = epoch + 1 + job.loss.append(logs.get("loss")) + job.val_loss.append(logs.get("val_loss")) + if job._stop_requested: + self.model.stop_training = True + + Path(models_dir).mkdir(parents=True, exist_ok=True) + output = train_mdrnn( + None, + dataset_file, + model_size, + early_stopping=True, + patience=patience, + num_epochs=max_epochs, + batch_size=64, + save_location=models_dir, + callbacks=[ProgressCallback()], + ) + model_files = [ + Path(output[k]).name + for k in ("tflite_file", "keras_file") + if k in output + ] + error, status = None, "finished" + except Exception as e: + model_files, error, status = [], str(e), "failed" + # Lock so a stop() request can't overwrite the final status. + with self._lock: + self.model_files = model_files + self.error = error + self.status = status + self.finished_at = time.time() + + def to_dict(self): + end = self.finished_at or time.time() + return { + "status": self.status, + "dataset": self.dataset, + "model_size": self.model_size, + "max_epochs": self.max_epochs, + "epoch": self.epoch, + "loss": self.loss, + "val_loss": self.val_loss, + "model_files": self.model_files, + "error": self.error, + "elapsed": int(end - self.started_at) if self.started_at else 0, + } + + +_training_job = TrainingJob() + + class MonitorListener: """Lazily-started OSC listener that mirrors IMPSY's /monitor/{in,out} stream. @@ -397,8 +554,27 @@ def index(): ) -@app.route("/logs") +@app.route("/logs", methods=["GET", "POST"]) def logs(): + if request.method == "POST": + uploaded = [] + for file in request.files.getlist("file"): + if not file or file.filename == "": + continue + filename = secure_filename(file.filename) + if not allowed_log_file(filename) or log_file_dimension(filename) is None: + flash( + f"Skipped {file.filename}: log files need names ending in -{{dimension}}d-mdrnn.log", + "error", + ) + continue + LOGS_DIR.mkdir(parents=True, exist_ok=True) + file.save(LOGS_DIR / filename) + uploaded.append(filename) + if uploaded: + flash(f"Uploaded {len(uploaded)} log file(s)", "success") + return redirect(url_for("logs")) + if LOGS_DIR.exists(): log_files = sorted( [get_log_file_info(f) for f in LOGS_DIR.iterdir() if f.suffix == ".log"], @@ -445,7 +621,11 @@ def datasets(): if DATASET_DIR.exists(): dataset_files = sorted( - [get_file_info(f) for f in DATASET_DIR.iterdir() if f.suffix == ".npz"], + [ + get_dataset_file_info(f) + for f in DATASET_DIR.iterdir() + if f.suffix == ".npz" + ], key=lambda x: x["name"], reverse=True, ) @@ -514,6 +694,77 @@ def delete_model(filename): return redirect(url_for("models")) +@app.route("/train", methods=["GET", "POST"]) +def train(): + """Train a model from a dataset in the background.""" + if request.method == "POST": + dataset_name = secure_filename(request.form.get("dataset", "")) + dataset_file = DATASET_DIR / dataset_name + model_size = request.form.get("model_size", "s") + max_epochs = request.form.get("max_epochs", type=int, default=100) + patience = request.form.get("patience", type=int, default=10) + if not dataset_name or not dataset_file.exists(): + flash("Choose a dataset to train on.", "error") + elif model_size not in TRAIN_SIZES: + flash(f"Unknown model size: {model_size}", "error") + elif not 1 <= max_epochs <= TRAIN_MAX_EPOCHS: + flash(f"Epochs must be between 1 and {TRAIN_MAX_EPOCHS}.", "error") + elif patience < 1: + flash("Patience must be at least 1.", "error") + elif not _training_job.start( + dataset_file, model_size, max_epochs, patience, MODEL_DIR + ): + flash( + "A model is already training. Wait for it to finish or stop it.", + "error", + ) + else: + flash(f"Started training on {dataset_name}.", "success") + return redirect(url_for("train")) + + if DATASET_DIR.exists(): + dataset_files = sorted( + [ + get_dataset_file_info(f) + for f in DATASET_DIR.iterdir() + if f.suffix == ".npz" + ], + key=lambda x: x["name"], + reverse=True, + ) + else: + dataset_files = [] + workflow = get_workflow_status() + default_size = "s" + try: + with open(CONFIG_FILE, "rb") as f: + default_size = tomllib.load(f).get("model", {}).get("size", "s") + except Exception: + pass + return render_template( + "train.html", + active_page="train", + dataset_files=dataset_files, + selected_dataset=request.args.get("dataset"), + config_dimension=workflow.get("dimension"), + sizes=TRAIN_SIZES, + default_size=default_size, + max_epochs=TRAIN_MAX_EPOCHS, + job=_training_job.to_dict(), + ) + + +@app.route("/train/status") +def train_status(): + return _training_job.to_dict() + + +@app.route("/train/stop", methods=["POST"]) +def train_stop(): + _training_job.stop() + return redirect(url_for("train")) + + @app.route("/download_log/") def download_log(filename): return send_file(os.path.join(LOGS_DIR, filename), as_attachment=True) diff --git a/tests/test_data.py b/tests/test_data.py index 8e54938..2a167c1 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -86,3 +86,16 @@ def test_log_with_rnn_source_filtered(tmp_path): log = dataset.transform_log_to_sequence_example(str(log_file), dimension) # All entries should be from interface source only assert log.shape[0] == 9 # 10 interface lines minus 1 (first diff is NaN) + + +def test_generate_dataset_creates_destination( + log_files, log_location, dimension, tmp_path +): + """A fresh workspace (e.g., the web UI's) has no datasets folder yet.""" + from impsy import dataset + + destination = tmp_path / "new" / "datasets" + dataset_file = dataset.generate_dataset( + dimension, source=log_location, destination=destination + ) + assert dataset_file.exists() diff --git a/tests/test_trained_model.py b/tests/test_trained_model.py index 9ec0a8b..c9b57bf 100644 --- a/tests/test_trained_model.py +++ b/tests/test_trained_model.py @@ -70,3 +70,123 @@ def test_train_function( # assert 'keras_file' in trained_model # assert 'tflite_file' in trained_model # assert 'weights_file' in trained_model + + +def test_dataset_dimension(dimension, dataset_file): + from impsy import train + + assert train.dataset_dimension(dataset_file) == dimension + + +def test_find_dataset_file(tmp_path): + import pytest + from impsy import train + + npz = tmp_path / "some-data.npz" + assert train.find_dataset_file(npz) == npz + assert train.find_dataset_file(tmp_path, 4) == tmp_path / "training-dataset-4d.npz" + with pytest.raises(ValueError): + train.find_dataset_file(tmp_path) # no .npz files + npz.write_bytes(b"") + assert train.find_dataset_file(tmp_path) == npz + (tmp_path / "other.npz").write_bytes(b"") + with pytest.raises(ValueError): + train.find_dataset_file(tmp_path) # two .npz files + + +def test_train_rejects_wrong_dimension(dimension, dataset_file, tmp_path): + import pytest + from impsy import train + + with pytest.raises(ValueError, match="dimension"): + train.train_mdrnn( + dimension=dimension + 1, + dataset_location=dataset_file, + model_size="xxs", + early_stopping=False, + patience=10, + num_epochs=1, + batch_size=1, + save_location=tmp_path, + ) + + +def test_train_stopped_early_still_saves_tflite(dimension, dataset_file, tmp_path): + """A callback that stops training (e.g., the web UI's stop button) still gets a model.""" + import tensorflow as tf + from impsy import train + + class StopAfterOneBatch(tf.keras.callbacks.Callback): + def on_train_batch_end(self, batch, logs=None): + self.model.stop_training = True + + output = train.train_mdrnn( + dimension=None, + dataset_location=dataset_file, + model_size="xxs", + early_stopping=False, + patience=10, + num_epochs=50, + batch_size=1, + save_location=tmp_path, + callbacks=[StopAfterOneBatch()], + ) + assert len(output["history"].history["loss"]) == 1 + assert output["tflite_file"].exists() + assert output["keras_file"].exists() + + +def test_train_interrupted_still_saves_tflite(dimension, dataset_file, tmp_path): + """Ctrl-C during training saves the model with the latest weights.""" + import tensorflow as tf + from impsy import train + + class Interrupt(tf.keras.callbacks.Callback): + def on_train_batch_end(self, batch, logs=None): + raise KeyboardInterrupt + + output = train.train_mdrnn( + dimension=None, + dataset_location=dataset_file, + model_size="xxs", + early_stopping=False, + patience=10, + num_epochs=50, + batch_size=1, + save_location=tmp_path, + callbacks=[Interrupt()], + ) + assert output["tflite_file"].exists() + + +def test_train_command_with_dataset_argument(dimension, dataset_file, tmp_path): + from click.testing import CliRunner + from impsy import train + + result = CliRunner().invoke( + train.train, + [str(dataset_file), "-M", "xxs", "-N", "1", "-B", "1", "-O", str(tmp_path)], + ) + assert result.exit_code == 0, result.output + assert list(tmp_path.glob(f"*dim{dimension}*.tflite")) + + +def test_train_rejects_short_performances(tmp_path): + import numpy as np + import pytest + from impsy import train + + short = np.array([np.zeros((10, 3), dtype=np.float32)], dtype=object) + dataset_file = tmp_path / "training-dataset-3d.npz" + np.savez_compressed(dataset_file, perfs=short) + with pytest.raises(ValueError, match="longer than"): + train.train_mdrnn( + dimension=None, + dataset_location=dataset_file, + model_size="xxs", + early_stopping=False, + patience=10, + num_epochs=1, + batch_size=1, + save_location=tmp_path, + ) diff --git a/tests/test_webui.py b/tests/test_webui.py index f1e1650..77cfbab 100644 --- a/tests/test_webui.py +++ b/tests/test_webui.py @@ -439,3 +439,108 @@ def test_realtime_route_renders(client): finally: web_interface._monitor_listener.stop() web_interface._monitor_listener = None + + +# Log upload and training + + +def test_log_file_dimension(): + assert webui_mod.log_file_dimension("2024-06-01T12-00-00-8d-mdrnn.log") == 8 + assert webui_mod.log_file_dimension("x-12d-mdrnn.log") == 12 + assert webui_mod.log_file_dimension("notes.log") is None + assert webui_mod.log_file_dimension("x-d-mdrnn.log") is None + assert webui_mod.log_file_dimension("2d-mdrnn.log") is None # no separator + + +def test_dataset_file_info_uses_standard_name(tmp_path): + # Not a real .npz: the dimension must come from the name without loading it. + fake = tmp_path / "training-dataset-5d.npz" + fake.write_bytes(b"not a dataset") + assert webui_mod.get_dataset_file_info(fake)["dimension"] == 5 + other = tmp_path / "mine.npz" + other.write_bytes(b"not a dataset") + assert webui_mod.get_dataset_file_info(other)["dimension"] is None + + +def test_training_job_stop_does_not_overwrite_finished(): + job = webui_mod.TrainingJob() + job.status = "finished" + job.stop() + assert job.status == "finished" + assert not job._stop_requested + + +def test_logs_upload(client, tmp_path, restored_workspace): + set_workspace(tmp_path) + data = { + "file": [ + (io.BytesIO(b"2024-06-01T12:00:00,interface,0.5\n"), "a-2d-mdrnn.log"), + (io.BytesIO(b"stuff"), "badname.log"), + (io.BytesIO(b"stuff"), "a-2d-mdrnn.exe"), + ] + } + response = client.post( + "/logs", data=data, content_type="multipart/form-data", follow_redirects=True + ) + assert response.status_code == 200 + assert b"Uploaded 1 log file" in response.data + assert b"Skipped badname.log" in response.data + assert [f.name for f in (tmp_path / "logs").iterdir()] == ["a-2d-mdrnn.log"] + + +def test_train_get_without_datasets(client, tmp_path, restored_workspace): + set_workspace(tmp_path) + response = client.get("/train") + assert response.status_code == 200 + assert b"No datasets yet" in response.data + + +def test_train_status_idle(client): + response = client.get("/train/status") + assert response.status_code == 200 + assert "status" in response.get_json() + + +def test_train_post_rejects_bad_input(client, tmp_path, restored_workspace): + set_workspace(tmp_path) + (tmp_path / "datasets").mkdir() + (tmp_path / "datasets" / "d.npz").write_bytes(b"") + response = client.post( + "/train", data={"dataset": "missing.npz"}, follow_redirects=True + ) + assert b"Choose a dataset" in response.data + response = client.post( + "/train", data={"dataset": "d.npz", "model_size": "huge"}, follow_redirects=True + ) + assert b"Unknown model size" in response.data + response = client.post( + "/train", data={"dataset": "d.npz", "max_epochs": "0"}, follow_redirects=True + ) + assert b"Epochs must be" in response.data + response = client.post( + "/train", data={"dataset": "d.npz", "patience": "0"}, follow_redirects=True + ) + assert b"Patience must be" in response.data + + +def test_training_job_trains_and_stops(dimension, dataset_file, tmp_path): + """Run a real (tiny) training job in the background, then stop one early.""" + job = webui_mod.TrainingJob() + assert job.start(dataset_file, "xxs", 1, 10, tmp_path) + job._thread.join(timeout=300) + state = job.to_dict() + assert state["status"] == "finished", state["error"] + assert state["epoch"] == 1 + assert len(state["val_loss"]) == 1 + assert any(name.endswith(".tflite") for name in state["model_files"]) + for name in state["model_files"]: + assert (tmp_path / name).exists() + + # Stop a long job straight away; it should still finish with a model. + assert job.start(dataset_file, "xxs", 500, 500, tmp_path / "stopped") + assert not job.start(dataset_file, "xxs", 1, 10, tmp_path) # one at a time + job.stop() + job._thread.join(timeout=300) + assert job.status == "finished", job.error + assert job.epoch < 500 + assert (tmp_path / "stopped").exists()