PyTorch 用法指南

明月停辉123

398人浏览 · 2026-05-17 15:32:20

明月停辉123 · 2026-05-17 15:32:20 发布

PyTorch 是一个开源的深度学习框架，广泛应用于研究和生产环境。它基于张量（tensor）操作，提供了灵活的 API 和高效的 GPU 加速。本指南将逐步介绍 PyTorch 的核心功能，包括属性访问、基础运算、索引操作、张量转换、设备切换以及 YOLOv8 推理流程。

一、张量属性

张量是 PyTorch 的核心数据结构，类似于多维数组。以下是常见属性及其用法：

dtype: 数据类型（如 float32、int64）。
ndim: 维度数量（即秩）。
shape: 张量的形状，表示为元组（如 (3, 4) ）。
size(): 同 shape，返回形状元组。
device: 张量所在的设备（如 CPU 或 GPU）。
type(): 返回张量的类型（但通常使用 dtype 更直接）。

代码示例：

import torch

# 创建一个张量
tensor = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)

# 访问属性
print("dtype:", tensor.dtype)  # 输出: torch.float32
print("ndim:", tensor.ndim)    # 输出: 2
print("shape:", tensor.shape)  # 输出: torch.Size([2, 2])
print("size():", tensor.size()) # 输出: torch.Size([2, 2])
print("device:", tensor.device) # 输出: cpu
print("type():", tensor.type()) # 输出: torch.FloatTensor

二、基础运算

PyTorch 支持多种数学运算，适用于张量。以下是常用操作：

sum(): 计算张量所有元素的和。
min(): 返回最小值。
max(): 返回最大值。
argmin(): 返回最小值的索引。
argmax(): 返回最大值的索引。
mean(): 计算平均值。
transpose(dim0, dim1): 交换两个维度（如行和列）。
permute(*dims): 重新排列维度顺序。
T: 对 2D 张量进行转置（等价于 transpose(0, 1)）。
unsqueeze(dim): 在指定维度增加一个大小为 1 的维度。
squeeze(dim): 移除指定维度（如果大小为 1）。
clip(min, max): 将值裁剪到指定范围。

代码示例：

# 创建一个张量
tensor = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)

# 运算示例
print("sum:", tensor.sum())          # 输出: tensor(10.)
print("min:", tensor.min())          # 输出: tensor(1.)
print("max:", tensor.max())          # 输出: tensor(4.)
print("argmin:", tensor.argmin())    # 输出: tensor(0)
print("argmax:", tensor.argmax())    # 输出: tensor(3)
print("mean:", tensor.mean())        # 输出: tensor(2.5)
print("transpose:", tensor.transpose(0, 1)) # 输出: tensor([[1., 3.], [2., 4.]])
print("permute:", tensor.permute(1, 0)) # 同 transpose
print("T:", tensor.T)                # 输出: tensor([[1., 3.], [2., 4.]])
print("unsqueeze:", tensor.unsqueeze(0)) # 输出: tensor([[[1., 2.], [3., 4.]]])
print("squeeze:", tensor.unsqueeze(0).squeeze(0)) # 还原
print("clip:", tensor.clip(2, 3))    # 输出: tensor([[2., 2.], [3., 3.]])

三、索引

索引操作允许访问张量的特定元素或子集，类似于 Python 列表。支持整数索引、切片和布尔索引。

代码示例：

# 创建一个张量
tensor = torch.tensor([[1, 2, 3], [4, 5, 6]])

# 索引示例
print("索引第一个元素:", tensor[0, 0])    # 输出: tensor(1)
print("切片第一行:", tensor[0, :])      # 输出: tensor([1, 2, 3])
print("布尔索引:", tensor[tensor > 3])  # 输出: tensor([4, 5, 6])

四、张量转换

PyTorch 提供了与 NumPy 互操作的方法，以及类型转换和维度重排功能：

numpy(): 将张量转换为 NumPy 数组。
from_numpy(): 从 NumPy 数组创建张量。
float(): 将张量转换为浮点类型。
int(): 将张量转换为整数类型。
permute(): 用于维度重排，常用于图像格式转换（如 HWC 到 CHW）。

代码示例：

import numpy as np

# 创建张量和 NumPy 数组
tensor = torch.tensor([1, 2, 3])
np_array = np.array([4, 5, 6])

# 转换示例
print("to numpy:", tensor.numpy())           # 输出: [1 2 3]
print("from numpy:", torch.from_numpy(np_array)) # 输出: tensor([4, 5, 6])
print("float:", tensor.float())              # 输出: tensor([1., 2., 3.])
print("int:", tensor.int())                  # 输出: tensor([1, 2, 3])

# 图像格式转换示例 (HWC to CHW)
# 假设有一个图像张量 (高度, 宽度, 通道)
image_hwc = torch.randn(224, 224, 3)
image_chw = image_hwc.permute(2, 0, 1)  # 变为 (通道, 高度, 宽度)
print("HWC shape:", image_hwc.shape)     # 输出: torch.Size([224, 224, 3])
print("CHW shape:", image_chw.shape)     # 输出: torch.Size([3, 224, 224])

五、设备切换

PyTorch 允许在 CPU 和 GPU 之间切换设备，以利用硬件加速：

to(device): 移动张量到指定设备。
cpu(): 移动张量到 CPU。
cuda(): 移动张量到 GPU（需 CUDA 支持）。

代码示例：

# 检查 GPU 可用性
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 创建张量
tensor = torch.tensor([1, 2, 3])

# 设备切换示例
tensor_gpu = tensor.to(device)  # 移动到 GPU (如果可用)
print("device after to:", tensor_gpu.device)

tensor_cpu = tensor_gpu.cpu()   # 移回 CPU
print("device after cpu:", tensor_cpu.device)

# 或者直接使用 cuda()
if torch.cuda.is_available():
    tensor_cuda = tensor.cuda()
    print("device after cuda:", tensor_cuda.device)

六、YOLOv8 推理流程

YOLOv8 是一个实时目标检测模型，可通过 ultralytics 库实现。以下是完整推理流程：

安装 ultralytics 库：
```
pip install ultralytics
```
加载预训练模型：使用 YOLO() 加载模型（如 'yolov8n.pt' 表示 nano 版本）。
单图推理：调用 model.predict() 对单张图像进行推理。
检测结果解析：结果对象包含：
- boxes.cls: 检测到的类别索引。
- boxes.conf: 置信度分数。
- names: 类别名称映射。

代码示例：

from ultralytics import YOLO
import cv2

# 安装后导入库，并加载模型
model = YOLO('yolov8n.pt')  # 加载预训练模型

# 读取图像
image_path = 'path/to/image.jpg'
image = cv2.imread(image_path)

# 单图推理
results = model.predict(image)  # 直接传入图像数据或路径

# 解析结果
for result in results:
    boxes = result.boxes  # 获取边界框信息
    print("类别索引:", boxes.cls)  # 输出: tensor([0., 2., ...]) 等
    print("置信度:", boxes.conf)  # 输出: tensor([0.95, 0.89, ...])
    print("类别名称:", result.names)  # 输出: 字典映射，如 {0: 'person', 1: 'car', ...}
    
    # 可视化结果 (可选)
    result.show()  # 显示带检测框的图像