Pytorch-Tutorial/test.py at main · Doraemonblogs/Pytorch-Tutorial · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
import torchvision
import torch
from PIL import Image
from model import classify_model,vgg_16
from torch import nn

# 定义CIFAR-10的类别标签
class_labels = [
    "airplane",
    "automobile",
    "bird",
    "cat",
    "deer",
    "dog",
    "frog",
    "horse",
    "ship",
    "truck"
]

image_path = 'img/cat1.jpg'
img = Image.open(image_path)

transforms = torchvision.transforms.Compose([torchvision.transforms.Resize((32,32)),
                                             torchvision.transforms.ToTensor()])

img = transforms(img)

# 加载模型
model = vgg_16()
model.load_state_dict(torch.load("check_point/vgg16_1/classify_model_98.pth", weights_only=True))

# model = classify_model()
# model.load_state_dict(torch.load("check_point/classify_model/classify_model_20.pth", weights_only=True))
# print(model)

img = torch.reshape(img, (1, 3, 32, 32))
model.eval()
with torch.no_grad():
    output = model(img)

pre_idx = torch.argmax(output, dim=1).item()
pre_label = class_labels[pre_idx]
confidence = torch.nn.functional.softmax(output, dim=1)[0][pre_idx].item()

print("预测标签为 {} ,置信度为 {}".format(pre_label, confidence))