Add a script to download model from NGC

This commit is contained in:
Adrian Lancucki 2021-09-23 18:37:59 +02:00 committed by Krzysztof Kudrynski
parent f4d6a00cb5
commit f24917b3ee
5 changed files with 34 additions and 7 deletions

View file

@ -0,0 +1,2 @@
pretrained_models/
results/

View file

@ -264,7 +264,7 @@ To train your model using mixed or TF32 precision with Tensor Cores or using FP3
7. Start inference/predictions.
Inside the container, use the following script to run inference.
Make sure the downloaded and preprocessed dataset is located at `$DATA_DIR/LibriSpeech` on the host, which is mounted as `/datasets/LibriSpeech` inside the container.
A pretrained model checkpoint can be downloaded from [NGC model repository](https://ngc.nvidia.com/catalog/models).
A pretrained model checkpoint can be downloaded from [NGC model repository](https://ngc.nvidia.com/catalog/models), manually or automatically using `scripts/download_quartznet.sh`.
```bash
[OPTION1=value1 OPTION2=value2 ...] bash scripts/inference.sh

View file

@ -4,13 +4,11 @@ SCRIPT_DIR=$(cd $(dirname $0); pwd)
QN_REPO=${QN_REPO:-"${SCRIPT_DIR}/../.."}
DATA_DIR=${1:-${DATA_DIR-${QN_REPO}"/datasets"}}
CHECKPOINT_DIR=${2:-${CHECKPOINT_DIR:-${QN_REPO}"/checkpoints"}}
RESULT_DIR=${3:-${RESULT_DIR:-${QN_REPO}"/results"}}
PROGRAM_PATH=${PROGRAM_PATH}
RESULT_DIR=${2:-${RESULT_DIR:-${QN_REPO}"/results"}}
SCRIPT=${3:-${SCRIPT:-""}}
MOUNTS=""
MOUNTS+=" -v $DATA_DIR:/datasets"
MOUNTS+=" -v $CHECKPOINT_DIR:/checkpoints"
MOUNTS+=" -v $RESULT_DIR:/results"
MOUNTS+=" -v ${QN_REPO}:/quartznet"
@ -21,4 +19,4 @@ docker run -it --rm --gpus all\
--ulimit stack=67108864 \
$MOUNTS \
-w /quartznet \
quartznet:latest bash $PROGRAM_PATH
quartznet:latest bash $SCRIPT

View file

@ -0,0 +1,27 @@
#!/usr/bin/env bash
set -e
: ${MODEL_DIR:="pretrained_models/quartznet"}
MODEL_ZIP="quartznet_pyt_ckpt_amp_21.03.0.zip"
MODEL="nvidia_quartznet_210504.pt"
MODEL_URL="https://api.ngc.nvidia.com/v2/models/nvidia/quartznet_pyt_ckpt_amp/versions/21.03.0/zip"
mkdir -p "$MODEL_DIR"
if [ ! -f "${MODEL_DIR}/${MODEL_ZIP}" ]; then
echo "Downloading ${MODEL_ZIP} ..."
wget -qO ${MODEL_DIR}/${MODEL_ZIP} ${MODEL_URL} \
|| { echo "ERROR: Failed to download ${MODEL_ZIP} from NGC"; exit 1; }
fi
if [ ! -f "${MODEL_DIR}/${MODEL}" ]; then
echo "Extracting ${MODEL} ..."
unzip -qo ${MODEL_DIR}/${MODEL_ZIP} -d ${MODEL_DIR} \
|| { echo "ERROR: Failed to extract ${MODEL_ZIP}"; exit 1; }
echo "OK"
else
echo "${MODEL} already downloaded."
fi

View file

@ -17,7 +17,7 @@
: ${DATA_DIR:=${1:-"/datasets/LibriSpeech"}}
: ${MODEL_CONFIG:=${2:-"configs/quartznet15x5_speedp-online-1.15_speca.yaml"}}
: ${OUTPUT_DIR:=${3:-"/results"}}
: ${CHECKPOINT:=${4:-"/checkpoints/quartznet_fp16.pt"}}
: ${CHECKPOINT:=${4:-"pretrained_models/quartznet/nvidia_quartznet_210504.pt"}}
: ${DATASET:="test-other"}
: ${LOG_FILE:=""}
: ${CUDNN_BENCHMARK:=false}