File size: 913 Bytes
330e909 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 | """Fix batch/spatial dimensions explicitly before the direct FlatBuffer export."""
import json,subprocess,sys
from pathlib import Path
import onnx
source,out=map(Path,sys.argv[1:3])
spatial_size=int(sys.argv[3])
model=onnx.load(str(source),load_external_data=False)
shapes=[]
for inp in model.graph.input:
dims=[]
for axis,dim in enumerate(inp.type.tensor_type.shape.dim):
if dim.dim_value>0: dims.append(dim.dim_value)
elif axis==0:dims.append(1)
elif len(inp.type.tensor_type.shape.dim)==4 and axis in {2,3}:dims.append(spatial_size)
else:raise ValueError(f'Unresolved input dimension: {inp.name} axis={axis}')
shapes.append(inp.name+':'+','.join(map(str,dims)))
out.mkdir(parents=True,exist_ok=True)
(out/'input-shapes.json').write_text(json.dumps(shapes,indent=2)+'\n')
subprocess.run([str(Path(sys.executable).parent/'onnx2tf'),'-i',str(source),'-o',str(out),'-b','1','-ois',*shapes],check=True)
|