mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: https://github.com/pytorch/pytorch/pull/106052 Approved by: https://github.com/albanD, https://github.com/Skylion007
25 lines
653 B
Python
25 lines
653 B
Python
import torch
|
|
import torchvision
|
|
|
|
print(torch.version.__version__)
|
|
|
|
resnet18 = torchvision.models.resnet18(pretrained=True)
|
|
resnet18.eval()
|
|
resnet18_traced = torch.jit.trace(resnet18, torch.rand(1, 3, 224, 224)).save(
|
|
"app/src/main/assets/resnet18.pt"
|
|
)
|
|
|
|
resnet50 = torchvision.models.resnet50(pretrained=True)
|
|
resnet50.eval()
|
|
torch.jit.trace(resnet50, torch.rand(1, 3, 224, 224)).save(
|
|
"app/src/main/assets/resnet50.pt"
|
|
)
|
|
|
|
mobilenet2q = torchvision.models.quantization.mobilenet_v2(
|
|
pretrained=True, quantize=True
|
|
)
|
|
mobilenet2q.eval()
|
|
torch.jit.trace(mobilenet2q, torch.rand(1, 3, 224, 224)).save(
|
|
"app/src/main/assets/mobilenet2q.pt"
|
|
)
|