mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08: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
23 lines
683 B
Python
23 lines
683 B
Python
import torch
|
|
from torchvision import models
|
|
|
|
print(torch.version.__version__)
|
|
|
|
resnet18 = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
|
|
resnet18.eval()
|
|
resnet18_traced = torch.jit.trace(resnet18, torch.rand(1, 3, 224, 224)).save(
|
|
"app/src/main/assets/resnet18.pt"
|
|
)
|
|
|
|
resnet50 = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
|
|
resnet50.eval()
|
|
torch.jit.trace(resnet50, torch.rand(1, 3, 224, 224)).save(
|
|
"app/src/main/assets/resnet50.pt"
|
|
)
|
|
|
|
mobilenet2q = 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"
|
|
)
|