25 lines
1.1 KiB
Python
25 lines
1.1 KiB
Python
import sys
|
|
from onnx_caffe import frontend
|
|
import argparse
|
|
import logging
|
|
|
|
#logging.basicConfig(level=logging.DEBUG)
|
|
parser = argparse.ArgumentParser(description='Convert a caffe model into an onnx file.')
|
|
#parser.add_argument('kfile', metavar='KerasFile', help='an input hdf5 file')
|
|
parser.add_argument('-n', metavar='prototxt', help='an input prototxt file')
|
|
parser.add_argument('-w', metavar='caffemodel', help='an input caffemodel file')
|
|
parser.add_argument('-o', '--output', dest='ofile', type=str, default="model.onnx", help='the output onnx file')
|
|
parser.add_argument('-c', '--custom', dest='custom', type=str, default=None, help='the customized layers definition file')
|
|
parser.add_argument('-D', '--debug', action='store_true', default=False, help='whether do various optimizations')
|
|
args = parser.parse_args()
|
|
|
|
# If in debug mode, output debug message
|
|
if args.debug:
|
|
logging.basicConfig(level=logging.DEBUG)
|
|
|
|
# Convert it
|
|
converter = frontend.CaffeFrontend()
|
|
converter.loadFromFile(args.n, args.w)
|
|
onnx_model = converter.convertToOnnx()
|
|
converter.saveToFile(args.ofile)
|