PRUNED_GRAPH=${WORKSPACE}/pruned/graph.pbtxt
PRUNED_CKPT=${WORKSPACE}/pruned/sparse.ckpt
action=prune
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.2 \
--gpu="0,1,2,3"