mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
f159f12fee
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/64851 ghstack-source-id: 137970229 Test Plan: CircleCI Reviewed By: hanton, cccclai Differential Revision: D30877963 fbshipit-source-id: 7bb8ade1959b85c3902ba9dc0660cdac8f558d64
12 lines
508 B
Python
12 lines
508 B
Python
import torch
|
|
import torchvision
|
|
from torch.utils.mobile_optimizer import optimize_for_mobile
|
|
|
|
model = torchvision.models.mobilenet_v2(pretrained=True)
|
|
model.eval()
|
|
example = torch.rand(1, 3, 224, 224)
|
|
traced_script_module = torch.jit.trace(model, example)
|
|
optimized_scripted_module = optimize_for_mobile(traced_script_module)
|
|
torch.jit.save(optimized_scripted_module, '../models/model.pt')
|
|
exported_optimized_scripted_module = optimized_scripted_module._save_for_lite_interpreter("../models/model_lite.ptl")
|