在机器学习中，通常需要一个大规模且标注良好的数据集来解决问题。有许多免费的预训练模型可以直接使用。开始自己的深度学习项目时，推荐先在线上查找已有的模型，看看是否能帮助实现目标。例如，可以探索 NGC 或在 GitHub 上搜索。这些资源中包含了大量现成的模型。

1 工程介绍

目的如下：

1. 使用 TorchVision 加载一个训练良好的预训练模型。
2. 预处理图像以符合预训练模型的输入要求。
3. 使用预训练模型对自己的图像进行准确的推理。

我们将创建一个宠物门，只允许狗进入，但将猫等其他动物留在外面。如果通过手动训练模型实现该功能，需要一个包含多种动物的大型数据集。不过，我们可以直接使用已有的预训练模型实现，VGG 模型（如 VGG16 和 VGG19）最初是在 ImageNet 数据集上的 ILSVRC 挑战中提出，并使用该数据集进行了预训练。训练结果表明，VGG 模型在图像分类任务上表现优异。

ImageNet：

• 很多模型基于这里面数百万张图片进行训练，可将图像分类为 1000 个类别。
• 包含许多动物类别，包括狗和猫的多个品种。

2 代码实现

下面是我们需要使用的库：

import os
from google.colab import drive

# 导入所需库
import torch
import torchvision.transforms.v2 as transforms
import torchvision.io as tv_io
import json

# 检查设备是否支持 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.cuda.is_available()

# 抑制动态计算中的错误提示
import torch._dynamo
torch._dynamo.config.suppress_errors = True

2.1 加载模型

ImageNet 的预训练模型可直接通过 TorchVision 下载并加载。以下是使用 VGG16 模型的示例：

from torchvision.models import vgg16
from torchvision.models import VGG16_Weights

# 加载 VGG16 模型，并使用默认权重
weights = VGG16_Weights.DEFAULT
model = vgg16(weights=weights)

# 将模型加载到指定设备
model.to(device)

2.1.2 输入

输入图像的尺寸需与模型期望的输入尺寸匹配。PyTorch 的模型是动态的，这意味着你可以传入任意尺寸的输入，只要模型架构本身支持这种形状的计算（例如卷积层不会因输入大小而出错）。

• 但实际使用中，预训练模型的权重是基于特定输入形状训练出来的。如果输入的形状与训练时的形状不一致，模型的预测可能会不准确。

  • 预训练权重是针对某个固定的输入形状（如 $224\times 224$）训练得到的。例如，VGG16 使用 ImageNet 数据集训练时，所有图像都被调整为 $224\times 224$ 的分辨率。

  • 预训练权重的效果依赖于输入图像与训练数据形状的匹配，因为这确保了模型提取特征的方式与训练时一致。

weights.transforms() 会根据预训练模型的要求，自动生成一组标准化的预处理步骤。这些步骤是为了将输入数据调整为模型训练时的形状、范围和分布。这些预处理步骤是模型权重（weights）中定义的，它们基于该模型在特定数据集（如 ImageNet）上训练时的预处理需求。

pre_trans = weights.transforms()
pre_trans

输出(下面是参数显示的顺序和执行顺序无关):
ImageClassification(
    crop_size=[224]
    resize_size=[256]
    mean=[0.485, 0.456, 0.406]
    std=[0.229, 0.224, 0.225]
    interpolation=InterpolationMode.BILINEAR
)

• 预训练模型中的参数，比如标准化中均值和标准差的具体值、图像的大小，都可以在TorchVision的官方文档中找到。

等价于以下代码：

IMG_WIDTH, IMG_HEIGHT = (224, 224)

pre_trans = transforms.Compose([
    transforms.ToDtype(torch.float32, scale=True),  # 将 [0, 255] 转为 [0, 1]
    transforms.Resize((IMG_WIDTH, IMG_HEIGHT)),     # 调整大小
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],                 # 归一化均值
        std=[0.229, 0.224, 0.225],                  # 归一化标准差
    ),
    transforms.CenterCrop(224)                      # 中心裁剪
])

2.1.3 输出

模型输出为长度为 1000 的数组，每个值表示图像属于对应类别的概率，即 ImageNet 数据集中的 $1000$ 个类别。

类别解析：

• 狗类别：编号 151 到 268。
• 猫类别：编号 281 到 285。

这些类别可帮助宠物门识别门口的动物类型，并决定是否允许进入。

2.2 加载图片

我们首先加载一张图片并显示。

import matplotlib.pyplot as plt
import matplotlib.image as mpimg

def show_image(image_path):
    image = mpimg.imread(image_path)
    plt.imshow(image)

show_image("happy_dog.jpg")

输出如下：

2.3 图像预处理

接下来，我们对图像进行预处理，使其能够作为输入发送到模型中。这与之前处理手语图像的练习类似。这里需要确保图像的最终形状为 $(1,3,224,224)$（批次大小、通道数、高度、宽度）。我们将使用预训练权重提供的 Transforms 来完成这一任务。

def load_and_process_image(file_path):
    # 打印图像的原始尺寸
    print('Original image shape: ', mpimg.imread(file_path).shape)
    image = tv_io.read_image(file_path).to(device)  # 读取图像并移动到设备
    image = pre_trans(image)  # 使用权重提供的 transforms
    image = image.unsqueeze(0)  # 转换为批次
    return image

以下是对图像预处理的测试结果：

processed_image = load_and_process_image("happy_dog.jpg")
print("Processed image shape: ", processed_image.shape)

输出:
Original image shape:  (1200, 1800, 3)
Processed image shape:  torch.Size([1, 3, 224, 224])

预处理完成后，看起来形状是对的。接下来打印图像以验证预处理效果：

import torchvision.transforms.functional as F

plot_image = F.to_pil_image(torch.squeeze(processed_image))  # 转换为 PIL 图像
plt.imshow(plot_image, cmap='gray')  # 显示图像

输出如下：

这里的颜色异常由于 Normalize 预处理的效果，图像的像素值被调整为标准化的分布（基于均值和标准差），导致颜色看起来不正常，但模型仍然能够正确处理标准化后的图像并进行预测。

2.4 预测

现在图像已经经过正确的预处理，可以将其输入到模型中进行预测。模型输出是一个包含 1000 个元素的数组，直接查看结果会比较困难。我们可以通过加载一个包含所有类别信息的 JSON 文件来将预测结果转化为易于理解的形式。

vgg_classes = json.load(open("imagenet_class_index.json"))

JSON 文件中每个类别使用数字字符串作为键，例如：

vgg_classes["0"]  # 查看类别 "0" 的信息

输出(0对应的动物为鲤鱼tench):
['n01440764', 'tench']

下面我们写一个函数，用于从 VGG 模型中提取预测结果并将其转换为人类可读的形式。这与之前的 predict_letter 函数类似，但这里使用 torch.topk 函数获取预测概率最高的前 3 个结果。

def readable_prediction(image_path):
    # 显示图像
    show_image(image_path)
    # 加载并预处理图像
    image = load_and_process_image(image_path)
    # 进行预测
    output = model(image)[0]  # 去除批次维度
    predictions = torch.topk(output, 3)  # 获取前 3 个预测结果
    indices = predictions.indices.tolist()
    # 将预测结果转化为可读形式
    out_str = "Top results: "
    pred_classes = [vgg_classes[str(idx)][1] for idx in indices]
    out_str += ", ".join(pred_classes)
    print(out_str)

    return predictions

现在试着对几种动物的图片进行分类，观察模型的结果。

readable_prediction("happy_dog.jpg")

输出：

readable_prediction("brown_bear.jpg")

readable_prediction("sleepy_cat.jpg")

2.5 识别狗

现在我们可以用模型的预测结果，通过分类控制出入：让狗进出，把猫留在里面。狗的分类编号是151到268，猫的分类编号是281到285。

import numpy as np

def doggy_door(image_path):
    show_image(image_path)
    image = load_and_process_image(image_path)  # 加载并处理图片
    idx = model(image).argmax(dim=1).item()     # 获取预测的类别索引
    print("Predicted index:", idx)
    if 151 <= idx <= 268:
        print("Doggy come on in!")
    elif 281 <= idx <= 285:
        print("Kitty stay inside!")
    else:
        print("You're not a dog! Stay outside!")

下面看一下效果：

doggy_door("brown_bear.jpg")
doggy_door("happy_dog.jpg")
doggy_door("sleepy_cat.jpg")

输出如下：

3 总结

从本篇文章可以学到，利用一个强大的预训练模型，我们仅用几行代码就创建了一个功能齐全的狗狗门。随着深度学习社区的不断发展，会有更多模型供你在自己的项目中使用，无需大量前期工作就能使用深度学习的乐趣。