torchvision 是一个基于 PyTorch 的计算机视觉库，它提供了一些常用的工具和功能，帮助用户更轻松地进行计算机视觉任务的开发和研究。torchvision 主要包含以下几个方面：
1. 数据集（Datasets）
torchvision 提供了一些常见的计算机视觉数据集的预加载接口，比如：
- MNIST：手写数字数据集
- CIFAR-10/CIFAR-100：小物体图像数据集
- ImageNet：大规模图像分类数据集
- COCO：目标检测、分割和字幕数据集
- 其他多个流行的数据集
这些数据集可以通过简单的代码加载并用于训练和评估模型。
2. 模型（Models）
torchvision 提供了许多预训练的模型，这些模型可以直接用于迁移学习或进行微调。常见的预训练模型包括：
- ResNet 系列（如 ResNet18、ResNet50 等）
- VGG 系列
- AlexNet
- MobileNet
- DenseNet
- Faster R-CNN 和 Mask R-CNN（用于目标检测和分割）
- 以及其他很多经典的计算机视觉模型。
这些模型在大规模数据集（如 ImageNet）上预训练过，可以直接加载并在其他任务中使用。
3. 图像转换（Transforms）
torchvision.transforms 提供了一些常见的图像预处理和数据增强的功能，如：
- 缩放、裁剪、旋转
- 随机翻转、颜色变化、裁剪等
- 归一化和标准化
这些转换操作对于数据预处理、数据增强和训练过程中的优化都非常重要。
4. 工具函数（Utilities）
torchvision 还提供了其他一些实用的工具和函数，例如：
- 用于图像的读取和写入
- 计算模型精度、IoU（交并比）等指标
- 可视化工具（如绘制边界框等）
安装 torchvision
如果你已经安装了 PyTorch，可以通过以下命令安装 torchvision：
pip install torchvision
或者，如果你使用 conda，可以通过以下命令安装：
conda install -c pytorch torchvision
例子：加载数据集并使用预训练模型
以下是一个简单的例子，展示如何使用 torchvision 加载一个预训练模型，并进行推理：
import torch
from torchvision import models, transforms
from PIL import Image
# 加载预训练的ResNet50模型
model = models.resnet50(pretrained=True)
model.eval()
# 加载并预处理图像
img = Image.open("image.jpg")
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
img_tensor = transform(img).unsqueeze(0)
# 进行推理
with torch.no_grad():
output = model(img_tensor)
# 打印预测结果
_, predicted_class = output.max(1)
print(f"Predicted class: {predicted_class.item()}")
