You can prune the model now and write some shell scripts to call the vai_p_tensorflow
functions.
WORKSPACE=./models
BASELINE_GRAPH=${WORKSPACE}/mnist.pbtxt
BASELINE_CKPT=${WORKSPACE}/train/model.ckpt
PRUNED_GRAPH=${WORKSPACE}/pruned/graph.pbtxt
PRUNED_CKPT=${WORKSPACE}/pruned/sparse.ckpt
INPUT_NODES="image"
OUTPUT_NODES="logits/add"
action=prune
mkdir -p $(dirname "${PRUNED_GRAPH}")
vai_p_tensorflow \
--action=${action} \
--input_graph=${BASELINE_GRAPH} \
--input_ckpt=${BASELINE_CKPT} \
--output_graph=${PRUNED_GRAPH} \
--output_ckpt=${PRUNED_CKPT} \
--workspace=${WORKSPACE} \
--input_nodes="${INPUT_NODES}" \
--input_node_shapes="1,28,28,1" \
--output_nodes="${OUTPUT_NODES}" \
--sparsity=0.5 \
--gpu="0,1,2,3" \
2>&1 | tee prune.log