frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Move helper scripts to new repo (#6159)

Browse Source
This commit is contained in:
bddppq
2018-04-02 11:06:29 -07:00
committed by Edward Z. Yang
parent 63af898d46
commit 28e66705ff
5 changed files with 783 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

91
scripts/fbcode-dev-setup/ccache_setup.sh Executable file
View File

@ -0,0 +1,91 @@
#!/bin/bash
# This script installs CCache with CUDA support.
# Example usage:
# ./ccache_setup.sh --path /installed/folder
set -e
shopt -s expand_aliases
# Setup the proxy
alias with_proxy="HTTPS_PROXY=http://fwdproxy.any:8080 HTTP_PROXY=http://fwdproxy.any:8080 FTP_PROXY=http://fwdproxy.any:8080 https_proxy=http://fwdproxy.any:8080 http_proxy=http://fwdproxy.any:8080 ftp_proxy=http://fwdproxy.any:8080 http_no_proxy='*.facebook.com|*.tfbnw.net|*.fb.com'"
# Parse options
path="$HOME/ccache"
force=false
while [[ $# -gt 0 ]]; do
case "$1" in
--path)
shift
path="$1"
path=$(realpath "$path")
;;
--force) # Force install
force=true
;;
--help)
echo 'usage: ./ccache_setup.py --path /installed/folder [--force]'
exit 0
;;
*)
echo "Invalid option: $1"
exit 1
;;
esac
shift
done
# Check whether you put nvcc in PATH
set +e
nvcc_path=$(which nvcc)
if [[ -z "$nvcc_path" ]]; then
nvcc_path="/usr/local/cuda/bin/nvcc"
fi
set -e
if [ ! -f "$nvcc_path" ] && ! $force; then
# shellcheck disable=SC2016
echo 'nvcc is not detected in $PATH'
exit 1
fi
echo "nvcc is detected at $nvcc_path"
if [ -f "$CUDA_NVCC_EXECUTABLE" ] && [[ "$CUDA_NVCC_EXECUTABLE" == *"ccache"* ]]; then # Heuristic rule
if $CUDA_NVCC_EXECUTABLE --version; then
if ! $force; then
echo "CCache with nvcc support is already installed at $CUDA_NVCC_EXECUTABLE, please add --force"
exit 0
fi
fi
fi
# Installing CCache
echo "CCache will be installed at $path"
if [ -e "$path" ]; then
mv --backup=t -T "$path" "${path}.old"
fi
with_proxy git clone https://github.com/colesbury/ccache.git "$path" -b ccbin
cd "$path"
./autogen.sh
./configure
make install prefix="$path"
mkdir -p "$path/lib"
mkdir -p "$path/cuda"
ln -sf "$path/bin/ccache" "$path/lib/cc"
ln -sf "$path/bin/ccache" "$path/lib/c++"
ln -sf "$path/bin/ccache" "$path/lib/gcc"
ln -sf "$path/bin/ccache" "$path/lib/g++"
ln -sf "$path/bin/ccache" "$path/cuda/nvcc"
"$path/bin/ccache" -M 25Gi
# Make sure the nvcc wrapped in CCache is runnable
"$path/cuda/nvcc" --version
echo 'Congrats! The CCache with nvcc support is installed!'
echo -e "Please add the following lines to your bash init script:\\n"
echo "################ Env Var for CCache with CUDA support ################"
# shellcheck disable=SC2016
echo 'export PATH="'"$path"'/lib:$PATH"'
echo 'export CUDA_NVCC_EXECUTABLE="'"$path"'/cuda/nvcc"'
echo '######################################################################'

10
scripts/fbcode-dev-setup/onnx_c2_sanity_check.sh Executable file
View File

@ -0,0 +1,10 @@
#!/bin/bash
set -e
python -c 'from caffe2.python import build; from pprint import pprint; pprint(build.build_options)'
python -c 'from caffe2.python import core, workspace; print("GPUs found: " + str(workspace.NumCudaDevices()))'
python -c "import onnx"
python -c "import torch"
echo "Caffe2, PyTorch and ONNX installed successfully!!"

180
scripts/fbcode-dev-setup/onnx_c2_setup.sh Executable file
View File

@ -0,0 +1,180 @@
#!/bin/bash
# This script helps developers set up the ONNX and Caffe2 develop environment on devgpu.
# It creates an virtualenv instance, and installs all the dependencies in this environment.
# The script will creates a folder called onnx-dev folder under the $HOME directory.
# onnx, pytorch and caffe2 are installed as submodules in $HOME/onnx-dev/onnx-fb-universe/repos.
# Please source $HOME/onnx-dev/.onnx_env_init to initialize the development before starting developing.
# TODO: support python 3.
# Set script configuration
set -e
shopt -s expand_aliases
# Proxy setup
alias with_proxy="HTTPS_PROXY=http://fwdproxy.any:8080 HTTP_PROXY=http://fwdproxy.any:8080 FTP_PROXY=http://fwdproxy.any:8080 https_proxy=http://fwdproxy.any:8080 http_proxy=http://fwdproxy.any:8080 ftp_proxy=http://fwdproxy.any:8080 http_no_proxy='*.facebook.com|*.tfbnw.net|*.fb.com'"
# Set the variables
RED='\033[0;31m'
CYAN='\033[0;36m'
NC='\033[0m'
onnx_root="$HOME/onnx-dev" # I think hardcoding the onnx root dir is fine, just like fbsource
venv="$onnx_root/onnxvenv"
onnx_init_file="$onnx_root/.onnx_env_init"
ccache_root="$onnx_root/ccache"
ccache_script="$(pwd)/ccache_install.sh"
sanity_script="$onnx_root/sanity.sh"
# Check whether default CUDA exists
# TODO check the required header and lib files
default_cuda="/usr/local/cuda"
if [[ ! -e "$default_cuda" ]]; then
echo "Default CUDA is not found at $default_cuda"
fi
# Checking to see if CuDNN is present, and install it if not exists
if [ -f /usr/local/cuda/include/cudnn.h ]; then
echo "CuDNN header already exists!!"
else
sudo cp -R /home/engshare/third-party2/cudnn/6.0.21/src/cuda/include/* /usr/local/cuda/include/
sudo cp -R /home/engshare/third-party2/cudnn/6.0.21/src/cuda/lib64/* /usr/local/cuda/lib64/
fi
# TODO set the specific version for each package
# Install the dependencies for Caffe2
sudo yum install python-virtualenv freetype-devel libpng-devel glog gflags protobuf protobuf-devel protobuf-compiler -y
rpm -q protobuf # check the version and if necessary update the value below
protoc --version # check protoc
protoc_path=$(which protoc)
if [[ "$protoc_path" != "/bin/protoc" ]]; then
echo "Warning: Non-default protoc is detected, the script may not work with non-default protobuf!!!"
echo "Please try to remove the protoc at $protoc_path and rerun this script."
exit 1
fi
# Upgrade Cmake to the right version (>3.0)
sudo yum remove cmake3 -y
sudo yum install cmake -y
# Install the dependencies for CCache
sudo yum install autoconf asciidoc -y
# Create the root folder
if [ -e "$onnx_root" ]; then
timestamp=$(date "+%Y.%m.%d-%H.%M.%S")
mv --backup=t -T "$onnx_root" "${onnx_root}.old.$timestamp"
fi
mkdir -p "$onnx_root"
# Set the name of virtualenv instance
with_proxy virtualenv "$venv"
# Creating a script that can be sourced in the future for the environmental variable
touch "$onnx_init_file"
{
# shellcheck disable=SC2016
echo 'if [ -z "$LD_LIBRARY_PATH" ]; then';
echo ' export LD_LIBRARY_PATH=/usr/local/cuda/lib64';
echo 'else'
# shellcheck disable=SC2016
echo ' export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH';
echo "fi"
# shellcheck disable=SC2016
echo 'export PATH='"$ccache_root"'/lib:/usr/local/cuda/bin:$PATH';
echo "source $venv/bin/activate";
echo 'alias with_proxy="HTTPS_PROXY=http://fwdproxy.any:8080 HTTP_PROXY=http://fwdproxy.any:8080 FTP_PROXY=http://fwdproxy.any:8080 https_proxy=http://fwdproxy.any:8080 http_proxy=http://fwdproxy.any:8080 ftp_proxy=http://fwdproxy.any:8080 http_no_proxy='"'"'*.facebook.com|*.tfbnw.net|*.fb.com'"'"'"'
} >> "$onnx_init_file"
chmod u+x "$onnx_init_file"
# Installing CCache
cd "$onnx_root"
if [ ! -f "$ccache_script" ]; then
ccache_script="$onnx_root/ccache_install.sh"
with_proxy wget https://raw.githubusercontent.com/onnxbot/onnx-fb-universe/master/scripts/ccache_setup.sh -O "$ccache_script"
fi
chmod u+x "$ccache_script"
"$ccache_script" --path "$ccache_root"
# Test nvcc with CCache
own_ccache=true
if [ -f "$CUDA_NVCC_EXECUTABLE" ] && [[ "$ccache_root/cuda/nvcc" != "$CUDA_NVCC_EXECUTABLE" ]] && \
[[ "$CUDA_NVCC_EXECUTABLE" == *"ccache"* ]]; then # Heuristic rule
if $CUDA_NVCC_EXECUTABLE --version; then
own_ccache=false
fi
fi
if $own_ccache; then
echo "export CUDA_NVCC_EXECUTABLE=$ccache_root/cuda/nvcc" >> "$onnx_init_file"
fi
# Loading env vars
# shellcheck disable=SC1090
source "$onnx_init_file"
"$CUDA_NVCC_EXECUTABLE" --version
# Create a virtualenv, activate it, upgrade pip
if [ -f "$HOME/.pip/pip.conf" ]; then
echo "${RED}Warning: $HOME/.pip/pip.conf is detected, pip install may fail!${NC}"
fi
with_proxy python -m pip install -U pip setuptools
with_proxy python -m pip install future numpy "protobuf>3.2" pytest-runner pyyaml typing ipython
# Cloning repos
cd "$onnx_root"
with_proxy git clone https://github.com/onnxbot/onnx-fb-universe --recursive
# Build ONNX
cd "$onnx_root/onnx-fb-universe/repos/onnx"
with_proxy python setup.py develop
# Build PyTorch
cd "$onnx_root/onnx-fb-universe/repos/pytorch"
with_proxy pip install -r "requirements.txt"
with_proxy python setup.py build develop
# Build Caffe2
set +e
cd "$onnx_root"
with_proxy wget https://raw.githubusercontent.com/onnxbot/onnx-fb-universe/master/scripts/onnx_c2_sanity_check.sh -O "$sanity_script"
chmod u+x "$sanity_script"
cd "$onnx_root/onnx-fb-universe/repos/caffe2"
with_proxy python setup.py develop
caffe2_exit_code=$?
caffe2_ok=true
if [ $caffe2_exit_code != 0 ]; then
caffe2_ok=false
fi
if ! $caffe2_ok; then
# Possible failure reasons when building Caffe2
ninja_path=$(which ninja)
if [[ ! -z "$ninja_path" ]]; then
echo "${RED}Warning: ninja is installed at $ninja_path, which may cause Caffe2 building issue!!!${NC}"
echo "${RED}Please try to remove the ninja at ${ninja_path}.${NC}"
fi
echo "${RED}We are almost there, only building Caffe2 fails. We can fix this problem seperately.${NC}"
echo "###### Please run the following command before development/fixing the problem: ######"
echo "${CYAN}source $onnx_init_file${NC}"
echo "#####################################################################################"
echo "########## Please run the following command to install Caffe2 seperately: ##########"
echo "${CYAN}cd $onnx_root/onnx-fb-universe/repos/caffe2; python setup.py develop${NC}"
echo "#####################################################################################"
echo "########### Please run the following command to check your installation: ###########"
echo "${CYAN}$sanity_script${NC}"
echo "#####################################################################################"
exit 1
fi
set -e
# Sanity checks and useful info
$sanity_script
echo "Congrats, you are ready to rock!!"
echo "################ Please run the following command before development ################"
echo -e "${CYAN}source $onnx_init_file${NC}"
echo "#####################################################################################"

165
scripts/model_zoo/update-caffe2-models.py Executable file
View File

@ -0,0 +1,165 @@
#! /usr/bin/env python
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
import os
import subprocess
import sys
import tarfile
import tempfile
from six.moves.urllib.request import urlretrieve
from caffe2.python.models.download import downloadFromURLToFile, getURLFromName, deleteDirectory
class SomeClass:
# largely copied from
# https://github.com/onnx/onnx-caffe2/blob/master/tests/caffe2_ref_test.py
def _download(self, model):
model_dir = self._caffe2_model_dir(model)
assert not os.path.exists(model_dir)
os.makedirs(model_dir)
for f in ['predict_net.pb', 'init_net.pb', 'value_info.json']:
url = getURLFromName(model, f)
dest = os.path.join(model_dir, f)
try:
try:
downloadFromURLToFile(url, dest,
show_progress=False)
except TypeError:
# show_progress not supported prior to
# Caffe2 78c014e752a374d905ecfb465d44fa16e02a28f1
# (Sep 17, 2017)
downloadFromURLToFile(url, dest)
except Exception as e:
print("Abort: {reason}".format(reason=e))
print("Cleaning up...")
deleteDirectory(model_dir)
exit(1)
def _caffe2_model_dir(self, model):
caffe2_home = os.path.expanduser('~/.caffe2')
models_dir = os.path.join(caffe2_home, 'models')
return os.path.join(models_dir, model)
def _onnx_model_dir(self, model):
onnx_home = os.path.expanduser('~/.onnx')
models_dir = os.path.join(onnx_home, 'models')
model_dir = os.path.join(models_dir, model)
return model_dir, os.path.dirname(model_dir)
# largely copied from
# https://github.com/onnx/onnx/blob/master/onnx/backend/test/runner/__init__.py
def _prepare_model_data(self, model):
model_dir, models_dir = self._onnx_model_dir(model)
if os.path.exists(model_dir):
return
os.makedirs(model_dir)
url = 'https://s3.amazonaws.com/download.onnx/models/{}.tar.gz'.format(model)
# On Windows, NamedTemporaryFile cannot be opened for a
# second time
download_file = tempfile.NamedTemporaryFile(delete=False)
try:
download_file.close()
print('Start downloading model {} from {}'.format(model, url))
urlretrieve(url, download_file.name)
print('Done')
with tarfile.open(download_file.name) as t:
t.extractall(models_dir)
except Exception as e:
print('Failed to prepare data for model {}: {}'.format(model, e))
raise
finally:
os.remove(download_file.name)
models = [
'bvlc_alexnet',
'densenet121',
'inception_v1',
'inception_v2',
'resnet50',
'shufflenet',
# TODO currently onnx can't translate squeezenet :(
# 'squeezenet',
'vgg16',
# TODO currently vgg19 doesn't work in the CI environment,
# possibly due to OOM
# 'vgg19'
]
def download_models():
sc = SomeClass()
for model in models:
print('update-caffe2-models.py: downloading', model)
caffe2_model_dir = sc._caffe2_model_dir(model)
onnx_model_dir, onnx_models_dir = sc._onnx_model_dir(model)
if not os.path.exists(caffe2_model_dir):
sc._download(model)
if not os.path.exists(onnx_model_dir):
sc._prepare_model_data(model)
def generate_models():
sc = SomeClass()
for model in models:
print('update-caffe2-models.py: generating', model)
caffe2_model_dir = sc._caffe2_model_dir(model)
onnx_model_dir, onnx_models_dir = sc._onnx_model_dir(model)
subprocess.check_call(['echo', model])
with open(os.path.join(caffe2_model_dir, 'value_info.json'), 'r') as f:
value_info = f.read()
subprocess.check_call([
'convert-caffe2-to-onnx',
'--caffe2-net-name', model,
'--caffe2-init-net', os.path.join(caffe2_model_dir, 'init_net.pb'),
'--value-info', value_info,
'-o', os.path.join(onnx_model_dir, 'model.pb'),
os.path.join(caffe2_model_dir, 'predict_net.pb')
])
subprocess.check_call([
'tar',
'-czf',
model + '.tar.gz',
model
], cwd=onnx_models_dir)
def upload_models():
sc = SomeClass()
for model in models:
print('update-caffe2-models.py: uploading', model)
onnx_model_dir, onnx_models_dir = sc._onnx_model_dir(model)
subprocess.check_call([
'aws',
's3',
'cp',
model + '.tar.gz',
"s3://download.onnx/models/{}.tar.gz".format(model),
'--acl', 'public-read'
], cwd=onnx_models_dir)
def cleanup():
sc = SomeClass()
for model in models:
onnx_model_dir, onnx_models_dir = sc._onnx_model_dir(model)
os.remove(os.path.join(os.path.dirname(onnx_model_dir), model + '.tar.gz'))
if __name__ == '__main__':
try:
subprocess.check_call(['aws', 'sts', 'get-caller-identity'])
except:
print('update-caffe2-models.py: please run `aws configure` manually to set up credentials')
sys.exit(1)
if sys.argv[1] == 'download':
download_models()
if sys.argv[1] == 'generate':
generate_models()
elif sys.argv[1] == 'upload':
upload_models()
elif sys.argv[1] == 'cleanup':
cleanup()

337
scripts/model_zoo/update-models-from-caffe2.py Normal file
View File

@ -0,0 +1,337 @@
#! /usr/bin/env python
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
import argparse
import caffe2.python.workspace as c2_workspace
import glob
import json
import math
import numpy as np
import onnx
import caffe2.python.onnx.frontend
import caffe2.python.onnx.backend
import os
import shutil
import subprocess
import sys
import tarfile
import tempfile
import boto3
from six.moves.urllib.request import urlretrieve
from caffe2.python.models.download import downloadFromURLToFile, getURLFromName, deleteDirectory
from caffe2.proto import caffe2_pb2
from onnx import numpy_helper
from filechunkio import FileChunkIO
"""A script converting Caffe2 models to ONNX, and updating ONNX model zoos.
Arguments:
-v, verbose
--local-dir, where we store the ONNX and Caffe2 models
--no-cache, ignore existing models in local-dir
--clean-test-data, delete all the existing test data when updating ONNX model zoo
--add-test-data, add add-test-data sets of test data for each ONNX model
--only-local, run locally (for testing purpose)
Examples:
# store the data in /home/username/zoo-dir, delete existing test data, ignore local cache,
# and generate 3 sets of new test data
python update-caffe2-models.py --local-dir /home/username/zoo-dir --clean-test-data --no-cache --add-test-data 3
"""
# TODO: Add GPU support
def upload_onnx_model(model_name, zoo_dir, backup=False, only_local=False):
if only_local:
print('No uploading in local only mode.')
return
model_dir = os.path.join(zoo_dir, model_name)
suffix = '-backup' if backup else ''
if backup:
print('Backing up the previous version of ONNX model {}...'.format(model_name))
rel_file_name = '{}{}.tar.gz'.format(model_name, suffix)
abs_file_name = os.path.join(zoo_dir, rel_file_name)
print('Compressing {} model to {}'.format(model_name, abs_file_name))
with tarfile.open(abs_file_name, 'w:gz') as f:
f.add(model_dir, arcname=model_name)
file_size = os.stat(abs_file_name).st_size
print('Uploading {} ({} MB) to s3 cloud...'.format(abs_file_name, float(file_size) / 1024 / 1024))
client = boto3.client('s3', 'us-east-1')
transfer = boto3.s3.transfer.S3Transfer(client)
transfer.upload_file(abs_file_name, 'download.onnx', 'models/latest/{}'.format(rel_file_name),
extra_args={'ACL': 'public-read'})
print('Successfully uploaded {} to s3!'.format(rel_file_name))
def download_onnx_model(model_name, zoo_dir, use_cache=True, only_local=False):
model_dir = os.path.join(zoo_dir, model_name)
if os.path.exists(model_dir):
if use_cache:
upload_onnx_model(model_name, zoo_dir, backup=True, only_local=only_local)
return
else:
shutil.rmtree(model_dir)
url = 'https://s3.amazonaws.com/download.onnx/models/latest/{}.tar.gz'.format(model_name)
download_file = tempfile.NamedTemporaryFile(delete=False)
try:
download_file.close()
print('Downloading ONNX model {} from {} and save in {} ...\n'.format(
model_name, url, download_file.name))
urlretrieve(url, download_file.name)
with tarfile.open(download_file.name) as t:
print('Extracting ONNX model {} to {} ...\n'.format(model_name, zoo_dir))
t.extractall(zoo_dir)
except Exception as e:
print('Failed to download/backup data for ONNX model {}: {}'.format(model_name, e))
if not os.path.exists(model_dir):
os.makedirs(model_dir)
finally:
os.remove(download_file.name)
if not only_local:
upload_onnx_model(model_name, zoo_dir, backup=True, only_local=only_local)
def download_caffe2_model(model_name, zoo_dir, use_cache=True):
model_dir = os.path.join(zoo_dir, model_name)
if os.path.exists(model_dir):
if use_cache:
return
else:
shutil.rmtree(model_dir)
os.makedirs(model_dir)
for f in ['predict_net.pb', 'init_net.pb', 'value_info.json']:
url = getURLFromName(model_name, f)
dest = os.path.join(model_dir, f)
try:
try:
downloadFromURLToFile(url, dest,
show_progress=False)
except TypeError:
# show_progress not supported prior to
# Caffe2 78c014e752a374d905ecfb465d44fa16e02a28f1
# (Sep 17, 2017)
downloadFromURLToFile(url, dest)
except Exception as e:
print("Abort: {reason}".format(reason=e))
print("Cleaning up...")
deleteDirectory(model_dir)
raise
def caffe2_to_onnx(caffe2_model_name, caffe2_model_dir):
caffe2_init_proto = caffe2_pb2.NetDef()
caffe2_predict_proto = caffe2_pb2.NetDef()
with open(os.path.join(caffe2_model_dir, 'init_net.pb'), 'rb') as f:
caffe2_init_proto.ParseFromString(f.read())
caffe2_init_proto.name = '{}_init'.format(caffe2_model_name)
with open(os.path.join(caffe2_model_dir, 'predict_net.pb'), 'rb') as f:
caffe2_predict_proto.ParseFromString(f.read())
caffe2_predict_proto.name = caffe2_model_name
with open(os.path.join(caffe2_model_dir, 'value_info.json'), 'rb') as f:
value_info = json.loads(f.read())
print('Converting Caffe2 model {} in {} to ONNX format'.format(caffe2_model_name, caffe2_model_dir))
onnx_model = caffe2.python.onnx.frontend.caffe2_net_to_onnx_model(
init_net=caffe2_init_proto,
predict_net=caffe2_predict_proto,
value_info=value_info
)
return onnx_model, caffe2_init_proto, caffe2_predict_proto
def tensortype_to_ndarray(tensor_type):
shape = []
for dim in tensor_type.shape.dim:
shape.append(dim.dim_value)
if tensor_type.elem_type == onnx.TensorProto.FLOAT:
type = np.float32
elif tensor_type.elem_type == onnx.TensorProto.INT:
type = np.int32
else:
raise
array = np.random.rand(*shape).astype(type)
return array
def generate_test_input_data(onnx_model, scale):
real_inputs_names = list(set([input.name for input in onnx_model.graph.input]) - set([init.name for init in onnx_model.graph.initializer]))
real_inputs = []
for name in real_inputs_names:
for input in onnx_model.graph.input:
if name == input.name:
real_inputs.append(input)
test_inputs = []
for input in real_inputs:
ndarray = tensortype_to_ndarray(input.type.tensor_type)
test_inputs.append((input.name, ndarray * scale))
return test_inputs
def generate_test_output_data(caffe2_init_net, caffe2_predict_net, inputs):
p = c2_workspace.Predictor(caffe2_init_net, caffe2_predict_net)
inputs_map = {input[0]:input[1] for input in inputs}
output = p.run(inputs_map)
c2_workspace.ResetWorkspace()
return output
def onnx_verify(onnx_model, inputs, ref_outputs):
prepared = caffe2.python.onnx.backend.prepare(onnx_model)
onnx_inputs = []
for input in inputs:
if isinstance(input, tuple):
onnx_inputs.append(input[1])
else:
onnx_inputs.append(input)
onnx_outputs = prepared.run(inputs=onnx_inputs)
np.testing.assert_almost_equal(onnx_outputs, ref_outputs, decimal=3)
model_mapping = {
'bvlc_alexnet': 'bvlc_alexnet',
'bvlc_googlenet': 'bvlc_googlenet',
'bvlc_reference_caffenet': 'bvlc_reference_caffenet',
'bvlc_reference_rcnn_ilsvrc13': 'bvlc_reference_rcnn_ilsvrc13',
'densenet121': 'densenet121',
#'finetune_flickr_style': 'finetune_flickr_style',
'inception_v1': 'inception_v1',
'inception_v2': 'inception_v2',
'resnet50': 'resnet50',
'shufflenet': 'shufflenet',
'squeezenet': 'squeezenet_old',
'vgg16': 'vgg16',
'vgg19': 'vgg19',
}
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Update the ONNX models.')
parser.add_argument('-v', action="store_true", default=False, help="verbose")
parser.add_argument("--local-dir", type=str, default=os.path.expanduser('~'),
help="local dir to store Caffe2 and ONNX models")
parser.add_argument("--no-cache", action="store_true", default=False,
help="whether use local ONNX models")
parser.add_argument('--clean-test-data', action="store_true", default=False,
help="remove the old test data")
parser.add_argument('--add-test-data', type=int, default=0,
help="add new test data")
parser.add_argument('--only-local', action="store_true", default=False,
help="no upload including backup")
args = parser.parse_args()
delete_test_data = args.clean_test_data
add_test_data = args.add_test_data
use_cache = not args.no_cache
only_local = args.only_local
root_dir = args.local_dir
caffe2_zoo_dir = os.path.join(root_dir, ".caffe2", "models")
onnx_zoo_dir = os.path.join(root_dir, ".onnx", "models")
for onnx_model_name in model_mapping:
c2_model_name = model_mapping[onnx_model_name]
print('####### Processing ONNX model {} ({} in Caffe2) #######'.format(onnx_model_name, c2_model_name))
download_caffe2_model(c2_model_name, caffe2_zoo_dir, use_cache=use_cache)
download_onnx_model(onnx_model_name, onnx_zoo_dir, use_cache=use_cache, only_local=only_local)
onnx_model_dir = os.path.join(onnx_zoo_dir, onnx_model_name)
if delete_test_data:
print('Deleting all the existing test data...')
# NB: For now, we don't delete the npz files.
#for f in glob.glob(os.path.join(onnx_model_dir, '*.npz')):
# os.remove(f)
for f in glob.glob(os.path.join(onnx_model_dir, 'test_data_set*')):
shutil.rmtree(f)
onnx_model, c2_init_net, c2_predict_net = caffe2_to_onnx(c2_model_name, os.path.join(caffe2_zoo_dir, c2_model_name))
print('Deleteing old ONNX {} model...'.format(onnx_model_name))
for f in glob.glob(os.path.join(onnx_model_dir, 'model*'.format(onnx_model_name))):
os.remove(f)
print('Serializing generated ONNX {} model ...'.format(onnx_model_name))
with open(os.path.join(onnx_model_dir, 'model.onnx'), 'wb') as file:
file.write(onnx_model.SerializeToString())
print('Verifying model {} with ONNX model checker...'.format(onnx_model_name))
onnx.checker.check_model(onnx_model)
total_existing_data_set = 0
print('Verifying model {} with existing test data...'.format(onnx_model_name))
for f in glob.glob(os.path.join(onnx_model_dir, '*.npz')):
test_data = np.load(f, encoding='bytes')
inputs = list(test_data['inputs'])
ref_outputs = list(test_data['outputs'])
onnx_verify(onnx_model, inputs, ref_outputs)
total_existing_data_set += 1
for f in glob.glob(os.path.join(onnx_model_dir, 'test_data_set*')):
inputs = []
inputs_num = len(glob.glob(os.path.join(f, 'input_*.pb')))
for i in range(inputs_num):
tensor = onnx.TensorProto()
with open(os.path.join(f, 'input_{}.pb'.format(i)), 'rb') as pf:
tensor.ParseFromString(pf.read())
inputs.append(numpy_helper.to_array(tensor))
ref_outputs = []
ref_outputs_num = len(glob.glob(os.path.join(f, 'output_*.pb')))
for i in range(ref_outputs_num):
tensor = onnx.TensorProto()
with open(os.path.join(f, 'output_{}.pb'.format(i)), 'rb') as pf:
tensor.ParseFromString(pf.read())
ref_outputs.append(numpy_helper.to_array(tensor))
onnx_verify(onnx_model, inputs, ref_outputs)
total_existing_data_set += 1
starting_index = 0
while os.path.exists(os.path.join(onnx_model_dir, 'test_data_set_{}'.format(starting_index))):
starting_index += 1
if total_existing_data_set == 0 and add_test_data == 0:
add_test_data = 3
total_existing_data_set = 3
print('Generating {} sets of new test data...'.format(add_test_data))
for i in range(starting_index, add_test_data + starting_index):
data_dir = os.path.join(onnx_model_dir, 'test_data_set_{}'.format(i))
os.makedirs(data_dir)
inputs = generate_test_input_data(onnx_model, 255)
ref_outputs = generate_test_output_data(c2_init_net, c2_predict_net, inputs)
onnx_verify(onnx_model, inputs, ref_outputs)
for index, input in enumerate(inputs):
tensor = numpy_helper.from_array(input[1])
with open(os.path.join(data_dir, 'input_{}.pb'.format(index)), 'wb') as file:
file.write(tensor.SerializeToString())
for index, output in enumerate(ref_outputs):
tensor = numpy_helper.from_array(output)
with open(os.path.join(data_dir, 'output_{}.pb'.format(index)), 'wb') as file:
file.write(tensor.SerializeToString())
del onnx_model
del c2_init_net
del c2_predict_net
upload_onnx_model(onnx_model_name, onnx_zoo_dir, backup=False, only_local=only_local)
print('\n\n')