使用 ResNet 进行图像分类
您已基于预训练的 ResNet18 构建了模型。现在,是时候在一张示例图像上进行测试了。
您将对图像应用预处理变换并进行分类。由于 ResNet18 是在多类别数据集上训练的,您需要先使用 softmax() 层,然后再用 argmax()。
您选择了以下图像用于预测测试:

预处理变换已保存为 preprocess。PIL 图像对象已上传为 img。
本练习是课程的一部分
使用 PyTorch 进行图像深度学习
练习说明
- 对图像应用预处理变换,并使用
.unsqueeze(0)重塑以添加批次维度。 - 将图像传入模型,对输出使用
.squeeze(0)去除批次维度,并添加softmax()层。 - 应用
argmax()选择概率最高的类别。
交互式实操练习
通过完成这段示例代码来试试这个练习。
# Apply preprocessing transforms
batch = ____.____
# Apply model with softmax layer
prediction = ____.____.____
# Apply argmax
class_id = prediction.____.item()
score = prediction[class_id].item()
category_name = weights.meta["categories"][class_id]
print(category_name)