xref: /aosp_15_r20/external/pytorch/ios/TestApp/benchmark/trace_model.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1from torchvision import models
2
3import torch
4from torch.utils.mobile_optimizer import optimize_for_mobile
5
6
7model = models.mobilenet_v2(weights=models.MobileNet_V2_Weights.IMAGENET1K_V1)
8model.eval()
9example = torch.rand(1, 3, 224, 224)
10traced_script_module = torch.jit.trace(model, example)
11optimized_scripted_module = optimize_for_mobile(traced_script_module)
12torch.jit.save(optimized_scripted_module, "../models/model.pt")
13exported_optimized_scripted_module = (
14    optimized_scripted_module._save_for_lite_interpreter("../models/model.ptl")
15)
16