Spaces:
Runtime error
Runtime error
File size: 2,248 Bytes
83d8d3c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
import numpy as np
import onnx
import torch
def convert_onnx(net, path_module, output, opset=11, simplify=False):
assert isinstance(net, torch.nn.Module)
img = np.random.randint(0, 255, size=(112, 112, 3), dtype=np.int32)
img = img.astype(np.float)
img = (img / 255.0 - 0.5) / 0.5 # torch style norm
img = img.transpose((2, 0, 1))
img = torch.from_numpy(img).unsqueeze(0).float()
weight = torch.load(path_module)
net.load_state_dict(weight, strict=True)
net.eval()
torch.onnx.export(
net, img, output, input_names=["data"], keep_initializers_as_inputs=False, verbose=False, opset_version=opset
)
model = onnx.load(output)
graph = model.graph
graph.input[0].type.tensor_type.shape.dim[0].dim_param = "None"
if simplify:
from onnxsim import simplify
model, check = simplify(model)
assert check, "Simplified ONNX model could not be validated"
onnx.save(model, output)
if __name__ == "__main__":
import os
import argparse
from backbones import get_model
parser = argparse.ArgumentParser(description="ArcFace PyTorch to onnx")
parser.add_argument("input", type=str, help="input backbone.pth file or path")
parser.add_argument("--output", type=str, default=None, help="output onnx path")
parser.add_argument("--network", type=str, default=None, help="backbone network")
parser.add_argument("--simplify", type=bool, default=False, help="onnx simplify")
args = parser.parse_args()
input_file = args.input
if os.path.isdir(input_file):
input_file = os.path.join(input_file, "model.pt")
assert os.path.exists(input_file)
# model_name = os.path.basename(os.path.dirname(input_file)).lower()
# params = model_name.split("_")
# if len(params) >= 3 and params[1] in ('arcface', 'cosface'):
# if args.network is None:
# args.network = params[2]
assert args.network is not None
print(args)
backbone_onnx = get_model(args.network, dropout=0.0, fp16=False, num_features=512)
if args.output is None:
args.output = os.path.join(os.path.dirname(args.input), "model.onnx")
convert_onnx(backbone_onnx, input_file, args.output, simplify=args.simplify)
|