Add TorchFix to the CI (#113403)

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
This commit is contained in:
Sergii Dymchenko
2023-11-14 01:26:01 +00:00
committed by PyTorch MergeBot
parent e1c872e009
commit d94bfaff2e
8 changed files with 18 additions and 15 deletions

View File

@ -1,23 +1,21 @@
import torch
import torchvision
from torchvision import models
print(torch.version.__version__)
resnet18 = torchvision.models.resnet18(pretrained=True)
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 = torchvision.models.resnet50(pretrained=True)
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 = torchvision.models.quantization.mobilenet_v2(
pretrained=True, quantize=True
)
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"

View File

@ -5,11 +5,11 @@ build script to create a tailored build which only contains these used ops.
"""
import torch
import torchvision
import yaml
from torchvision import models
# Download and trace the model.
model = torchvision.models.mobilenet_v2(pretrained=True)
model = models.mobilenet_v2(weights=models.MobileNet_V2_Weights.IMAGENET1K_V1)
model.eval()
example = torch.rand(1, 3, 224, 224)
# TODO: create script model with `torch.jit.script`