部署一个 PyTorch 模型

每次训练后，我们会保存若干检查点。其中 resnet50_best.pth.tar 是表现最佳的模型，我们将使用它进行推理。

加载检查点

初始化以下变量并加载我们的检查点：

1import torch
2from PIL import Image
3import torchvision.transforms as transforms
4import timm
5
6# 初始化一个新变量，表示检查点路径
7checkpoint_path = "/app/resnet50_best.pth.tar"
8# 初始化一个新变量，表示GPU可用性
9device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
10
11# 加载本地检查点
12model = timm.create_model(
13    "resnet50",
14    num_classes=4,
15    in_chans=3,
16    pretrained=False,
17    checkpoint_path=checkpoint_path
18)
19
20# 将模型设置为推理模式
21model.eval()
22# 根据GPU可用性将模型移动到对应设备
23model = model.to(device)
24print('模型中类别总数为:', model.num_classes)

注意：请确保设置正确的网络结构和类别数量。检查点的配置必须与模型一致。

预处理图像

使用 PIL 包加载输入图像，通过 transforms.Resize 函数调整其大小，最后将其转换为张量以进行推理。

1import torch
2from PIL import Image
3import torchvision.transforms as transforms
4import timm
5
6# 初始化一个新变量，表示检查点路径
7checkpoint_path = "/app/resnet50_best.pth.tar"
8# 初始化一个新变量，表示GPU可用性
9device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
10
11# 加载本地检查点
12model = timm.create_model(
13    "resnet50",
14    num_classes=4,
15    in_chans=3,
16    pretrained=False,
17    checkpoint_path=checkpoint_path
18)
19
20# 将模型设置为推理模式
21model.eval()
22# 根据GPU可用性将模型移动到对应设备
23model = model.to(device)
24
25# 将图像转换为RGB模式
26img = Image.open('/app/dataset2/val/cloudy/cloudy271.jpg').convert('RGB')
27# 将图像调整为224x224
28resize = transforms.Resize([224, 224])
29img = resize(img)
30# 将图像转换为张量
31to_tensor = transforms.ToTensor()
32tensor = to_tensor(img).unsqueeze(0)
33print('输出张量为:')
34print(tensor)

获取预测结果和概率

创建一个 torch.no_grad 上下文管理器，在其中执行推理。模型将输出一个张量，使用 torch.nn.functional.softmax 函数将其转换为概率值：

1import torch
2from PIL import Image
3import torchvision.transforms as transforms
4import timm
5
6# 初始化一个新变量，表示检查点路径
7checkpoint_path = "/app/resnet50_best.pth.tar"
8# 初始化一个新变量，表示GPU可用性
9device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
10
11# 加载本地检查点
12model = timm.create_model(
13    "resnet50",
14    num_classes=4,
15    in_chans=3,
16    pretrained=False,
17    checkpoint_path=checkpoint_path
18)
19
20# 将模型设置为推理模式
21model.eval()
22# 根据GPU可用性将模型移动到对应设备
23model = model.to(device)
24
25# 将图像转换为RGB模式
26img = Image.open('/app/dataset2/val/cloudy/cloudy271.jpg').convert('RGB')
27# 将图像调整为224x224
28resize = transforms.Resize([224, 224])
29img = resize(img)
30# 将图像转换为张量
31to_tensor = transforms.ToTensor()
32tensor = to_tensor(img).unsqueeze(0)
33
34# 在no_grad上下文中执行推理
35with torch.no_grad():
36    tensor = tensor.to(device)
37    out = model(tensor)
38
39# 将模型输出的张量转换为概率
40probabilities = torch.nn.functional.softmax(out[0], dim=0)
41print('输出的概率为:', probabilities)

获取前k个预测结果

调用 torch.topk 函数获取概率最高的预测结果，然后通过循环打印对应的类别名称和概率值。

1import torch
2from PIL import Image
3import torchvision.transforms as transforms
4import timm
5
6# 初始化一个新变量，表示检查点路径
7checkpoint_path = "/app/resnet50_best.pth.tar"
8# 初始化一个新变量，表示GPU可用性
9device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
10
11# 加载本地检查点
12model = timm.create_model(
13    "resnet50",
14    num_classes=4,
15    in_chans=3,
16    pretrained=False,
17    checkpoint_path=checkpoint_path
18)
19
20# 将模型设置为推理模式
21model.eval()
22# 根据GPU可用性将模型移动到对应设备
23model = model.to(device)
24
25# 将图像转换为RGB模式
26img = Image.open('/app/dataset2/val/cloudy/cloudy271.jpg').convert('RGB')
27# 将图像调整为224x224
28resize = transforms.Resize([224, 224])
29img = resize(img)
30# 将图像转换为张量
31to_tensor = transforms.ToTensor()
32tensor = to_tensor(img).unsqueeze(0)
33
34# 在no_grad上下文中执行推理
35with torch.no_grad():
36    tensor = tensor.to(device)
37    out = model(tensor)
38
39# 将模型输出的张量转换为概率
40probabilities = torch.nn.functional.softmax(out[0], dim=0)
41
42# 获取前4个预测结果
43k = 4
44topk_prob, topk_catid = torch.topk(probabilities, k)
45# 将索引映射到类别名称
46class_names = ['cloudy', 'rain', 'shine', 'sunrise']
47for i in range(topk_prob.size(0)):
48    print(class_names[topk_catid[i]], topk_prob[i].item())

上述推理代码可轻松集成到任意 Python 应用程序中。