smoke_test.py 2.85 KB
Newer Older
1
2
"""Run smoke tests"""

3
import os
4
from pathlib import Path
5
from sys import platform
6

7
import torch
8
import torch.nn as nn
soumith's avatar
soumith committed
9
import torchvision
10
from torchvision.io import read_image
11
from torchvision.models import resnet50, ResNet50_Weights
12

13
14
15
16
17
SCRIPT_DIR = Path(__file__).parent


def smoke_test_torchvision() -> None:
    print(
18
        "Is torchvision usable?",
19
20
21
        all(x is not None for x in [torch.ops.image.decode_png, torch.ops.torchvision.roi_align]),
    )

22

23
24
25
26
27
28
29
30
def smoke_test_torchvision_read_decode() -> None:
    img_jpg = read_image(str(SCRIPT_DIR / "assets" / "encode_jpeg" / "grace_hopper_517x606.jpg"))
    if img_jpg.ndim != 3 or img_jpg.numel() < 100:
        raise RuntimeError(f"Unexpected shape of img_jpg: {img_jpg.shape}")
    img_png = read_image(str(SCRIPT_DIR / "assets" / "interlaced_png" / "wizard_low.png"))
    if img_png.ndim != 3 or img_png.numel() < 100:
        raise RuntimeError(f"Unexpected shape of img_png: {img_png.shape}")

Philip Meier's avatar
Philip Meier committed
31

32
def smoke_test_compile() -> None:
33
34
35
36
37
38
39
40
41
42
43
    try:
        model = resnet50().cuda()
        model = torch.compile(model)
        x = torch.randn(1, 3, 224, 224, device="cuda")
        out = model(x)
        print(f"torch.compile model output: {out.shape}")
    except RuntimeError:
        if platform == "win32":
            print("Successfully caught torch.compile RuntimeError on win")
        else:
            raise
44

Philip Meier's avatar
Philip Meier committed
45

46
47
def smoke_test_torchvision_resnet50_classify(device: str = "cpu") -> None:
    img = read_image(str(SCRIPT_DIR / ".." / "gallery" / "assets" / "dog2.jpg")).to(device)
48
49
50

    # Step 1: Initialize model with the best available weights
    weights = ResNet50_Weights.DEFAULT
51
    model = resnet50(weights=weights).to(device)
52
53
54
55
56
57
58
59
60
61
62
63
64
65
    model.eval()

    # Step 2: Initialize the inference transforms
    preprocess = weights.transforms()

    # Step 3: Apply inference preprocessing transforms
    batch = preprocess(img).unsqueeze(0)

    # Step 4: Use the model and print the predicted category
    prediction = model(batch).squeeze(0).softmax(0)
    class_id = prediction.argmax().item()
    score = prediction[class_id].item()
    category_name = weights.meta["categories"][class_id]
    expected_category = "German shepherd"
66
    print(f"{category_name} ({device}): {100 * score:.1f}%")
67
    if category_name != expected_category:
68
69
        raise RuntimeError(f"Failed ResNet50 classify {category_name} Expected: {expected_category}")

70
71
72

def main() -> None:
    print(f"torchvision: {torchvision.__version__}")
73
    print(f"torch.cuda.is_available: {torch.cuda.is_available()}")
74
75
76
    smoke_test_torchvision()
    smoke_test_torchvision_read_decode()
    smoke_test_torchvision_resnet50_classify()
77
78
    if torch.cuda.is_available():
        smoke_test_torchvision_resnet50_classify("cuda")
79
80
        smoke_test_compile()

81
82
    if torch.backends.mps.is_available():
        smoke_test_torchvision_resnet50_classify("mps")
83

84

85
86
if __name__ == "__main__":
    main()