mirror of
https://github.com/TaterTotterson/microWakeWord-Trainer-Nvidia-Docker.git
synced 2026-08-12 16:05:34 -06:00
113 lines
3.3 KiB
Bash
Executable File
113 lines
3.3 KiB
Bash
Executable File
#!/bin/bash
|
|
set -euo pipefail
|
|
|
|
PROGDIR="$(dirname "$(realpath "$0")")"
|
|
ROOTDIR="$(dirname "${PROGDIR}")"
|
|
|
|
KNOWN_ARGS=( data-dir force python )
|
|
source "${PROGDIR}/shell.functions"
|
|
|
|
if [ ${#UNKNOWN_ARGS[@]} -gt 0 ] ; then
|
|
echo "Unknown argument(s): ${UNKNOWN_ARGS[*]}" >&2
|
|
HELP=true
|
|
fi
|
|
|
|
if [ "${HELP}" == "true" ] ; then
|
|
cat <<EOF >&2
|
|
Usage: setup_blackwell_venv [ --data-dir=/data ] [ --force ] [ --python=python3.13 ]
|
|
|
|
Creates /data/.venv-blackwell for RTX 50 / Blackwell TensorFlow training.
|
|
Sample generation and augmentation continue to use /data/.venv.
|
|
|
|
Environment overrides:
|
|
MWW_BLACKWELL_TF_WHEEL_URL: TensorFlow Blackwell wheel URL.
|
|
|
|
EOF
|
|
exit 1
|
|
fi
|
|
|
|
[ -n "${DATA_DIR}" ] && DATA_DIR="$(realpath "${DATA_DIR}")"
|
|
[ -d "${DATA_DIR}" ] || {
|
|
echo "Data directory '${DATA_DIR}' doesn't exist." >&2
|
|
exit 1
|
|
}
|
|
|
|
PYTHON="${PYTHON:-python3.13}"
|
|
VENV="${DATA_DIR}/.venv-blackwell"
|
|
MARKER="${VENV}/.mww-blackwell-venv"
|
|
TF_WHEEL_URL="${MWW_BLACKWELL_TF_WHEEL_URL:-https://github.com/chivitiH/tensorflow-blackwell-python313/releases/download/v2.22.0-selfbuilt/tensorflow-2.22.0.dev0+selfbuilt-cp313-cp313-linux_x86_64.whl}"
|
|
|
|
if ! command -v "${PYTHON}" >/dev/null 2>&1 ; then
|
|
echo "Python 3.13 is required for the Blackwell TensorFlow wheel. Missing: ${PYTHON}" >&2
|
|
exit 1
|
|
fi
|
|
|
|
if [ "${FORCE:-false}" != "true" ] && [ -x "${VENV}/bin/python" ] && [ -f "${MARKER}" ] ; then
|
|
echo " Blackwell TensorFlow venv found (skipping setup_blackwell_venv)"
|
|
exit 0
|
|
fi
|
|
|
|
echo "===== Setting up Blackwell TensorFlow environment ${VENV} ====="
|
|
rm -rf "${VENV}" || :
|
|
"${PYTHON}" -m venv --upgrade-deps "${VENV}"
|
|
source "${VENV}/bin/activate"
|
|
|
|
export PIP_PROGRESS_BAR=off
|
|
export PIP_NO_COLOR=1
|
|
export PIP_QUIET=0
|
|
|
|
pip_install() {
|
|
if $VERBOSE ; then
|
|
pip install "$@" || return 1
|
|
else
|
|
{ pip install "$@" || return 1 ; } | stdbuf -i0 -o0 tr -d '[:print:]' | stdbuf -i0 -o0 tr '\n' '.'
|
|
fi
|
|
echo
|
|
}
|
|
|
|
echo " ===== Installing Blackwell TensorFlow wheel ====="
|
|
pip_install --upgrade pip setuptools wheel
|
|
pip_install "${TF_WHEEL_URL}"
|
|
|
|
echo " ===== Installing microWakeWord training dependencies ====="
|
|
pip_install \
|
|
audiomentations \
|
|
audio_metadata \
|
|
datasets \
|
|
mmap_ninja \
|
|
pymicro-features \
|
|
pyyaml \
|
|
webrtcvad-wheels \
|
|
ai-edge-litert \
|
|
numpy-minmax \
|
|
numpy-rms \
|
|
absl-py \
|
|
"numpy==2.3.5"
|
|
|
|
echo " ===== Checking microwakeword ====="
|
|
MWW="${DATA_DIR}/tools/microWakeWord"
|
|
if [ ! -d "${MWW}" ] || [ -n "$(git -C "${MWW}" status --porcelain 2>/dev/null || true)" ] ; then
|
|
rm -rf "${MWW}" || :
|
|
mkdir -p "${DATA_DIR}/tools"
|
|
echo " Cloning micro-wake-word to ${DATA_DIR}/tools"
|
|
git clone https://github.com/TaterTotterson/micro-wake-word "${MWW}" &>/dev/null
|
|
fi
|
|
echo " Installing microwakeword into Blackwell venv"
|
|
pip_install --no-deps -e "${MWW}"
|
|
|
|
echo " ===== Testing Blackwell TensorFlow environment ====="
|
|
"${VENV}/bin/python" - <<'PY'
|
|
import tensorflow as tf
|
|
from ai_edge_litert.interpreter import Interpreter
|
|
from microwakeword.data import FeatureHandler
|
|
from microwakeword.inference import Model
|
|
|
|
print("TensorFlow:", tf.__version__)
|
|
print("CUDA build:", tf.test.is_built_with_cuda())
|
|
print("GPU:", tf.config.list_physical_devices("GPU"))
|
|
print("microWakeWord Blackwell imports available")
|
|
PY
|
|
|
|
touch "${MARKER}"
|
|
echo "Blackwell TensorFlow environment ready: ${VENV}"
|