#! /usr/bin/env python from __future__ import absolute_import from __future__ import division from __future__ import print_function from __future__ import unicode_literals import onnx.backend 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', 'zfnet512': 'zfnet512', } 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')