mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
d94bfaff2e
Enable flake8 plugin for https://github.com/pytorch/test-infra/tree/main/tools/torchfix - TorchFix 0.1.1. Disable TorchFix codes that don't make sense for PyTorch itself. Update deprecated TorchVision APIs to make TorchFix pass. Pull Request resolved: https://github.com/pytorch/pytorch/pull/113403 Approved by: https://github.com/Skylion007, https://github.com/malfet
39 lines
1 KiB
Python
39 lines
1 KiB
Python
import torch
|
|
|
|
from torch.backends._coreml.preprocess import CompileSpec, CoreMLComputeUnit, TensorSpec
|
|
from torchvision import models
|
|
|
|
|
|
def mobilenetv2_spec():
|
|
return {
|
|
"forward": CompileSpec(
|
|
inputs=(
|
|
TensorSpec(
|
|
shape=[1, 3, 224, 224],
|
|
),
|
|
),
|
|
outputs=(
|
|
TensorSpec(
|
|
shape=[1, 1000],
|
|
),
|
|
),
|
|
backend=CoreMLComputeUnit.CPU,
|
|
allow_low_precision=True,
|
|
),
|
|
}
|
|
|
|
|
|
def main():
|
|
model = models.mobilenet_v2(weights=models.MobileNet_V2_Weights.IMAGENET1K_V1)
|
|
model.eval()
|
|
example = torch.rand(1, 3, 224, 224)
|
|
model = torch.jit.trace(model, example)
|
|
compile_spec = mobilenetv2_spec()
|
|
mlmodel = torch._C._jit_to_backend("coreml", model, compile_spec)
|
|
print(mlmodel._c._get_method("forward").graph)
|
|
mlmodel._save_for_lite_interpreter("../models/model_coreml.ptl")
|
|
torch.jit.save(mlmodel, "../models/model_coreml.pt")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|