ML-Inference/prepare.py at master · 1997alireza/ML-Inference · GitHub

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
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
import torch
import torchvision.models as models
import coremltools as ct
import onnx
from onnxconverter_common import float16


# Load pretrained model
model = models.resnet50(weights='IMAGENET1K_V1')
model.eval()

dummy_input = torch.rand(1, 3, 224, 224)

# Trace the model
traced_model = torch.jit.trace(model, dummy_input)

# Export CoreML model
mlmodel = ct.convert(
    traced_model,
    inputs=[ct.TensorType(name="input", shape=(1, 3, 224, 224))],
    compute_units=ct.ComputeUnit.ALL,
    compute_precision=ct.precision.FLOAT32
)
mlmodel.save("models/resnet50_fp32.mlpackage")

mlmodel_fp16 = ct.convert(
    traced_model,
    inputs=[ct.TensorType(name="input", shape=(1, 3, 224, 224))],
    compute_units=ct.ComputeUnit.ALL,
    compute_precision=ct.precision.FLOAT16
)
mlmodel_fp16.save("models/resnet50_fp16.mlpackage")


pipeline = ct.PassPipeline.DEFAULT
pipeline.remove_passes({
    "common::merge_consecutive_transposes",  # not many in ResNet
    "common::cast_optimization",             # avoid float<->int casts
    "common::add_int16_cast"                 # not relevant to float models
})

mlmodel = ct.convert(
    traced_model,
    inputs=[ct.TensorType(name="input", shape=(1, 3, 224, 224))],
    compute_units=ct.ComputeUnit.ALL,
    compute_precision=ct.precision.FLOAT32,
    pass_pipeline=pipeline
)
mlmodel.save("models/resnet50_fp32_custom.mlpackage")

pipeline = ct.PassPipeline.DEFAULT
pipeline.remove_passes({
    "common::merge_consecutive_transposes",
    "common::add_int16_cast"
})

mlmodel_fp16 = ct.convert(
    traced_model,
    inputs=[ct.TensorType(name="input", shape=(1, 3, 224, 224))],
    compute_units=ct.ComputeUnit.ALL,
    compute_precision=ct.precision.FLOAT16,
    pass_pipeline=pipeline
)
mlmodel_fp16.save("models/resnet50_fp16_custom.mlpackage")

# Export ONNX model
torch.onnx.export(
    model,
    dummy_input,
    "models/resnet50_fp32.onnx",
    export_params=True,
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
)

model = onnx.load("models/resnet50_fp32.onnx")
model_fp16 = float16.convert_float_to_float16(model)
onnx.save(model_fp16, "models/resnet50_fp16.onnx")