通俗点说:
PyTorch 是一个用 Python 写深度学习模型的工具,就像乐高积木一样,简单好拼,灵活好玩。
它是 Facebook AI 开发的,后来开源火了,很多科研人员和开发者都在用。
如果你想做 AI 项目、深度学习、NLP、CV,PyTorch 几乎是第一推荐。
和 TensorFlow 一样,PyTorch 的核心数据结构叫 Tensor,就是高维数组:
import torch
x = torch.tensor([1.0, 2.0, 3.0]) # 一维张量(向量)
m = torch.tensor([[1, 2], [3, 4]]) # 二维张量(矩阵)
你可以像操作 Numpy 一样操作它:加法、乘法、矩阵运算都支持。
PyTorch 最大的亮点就是:自动帮你算导数!
只要你告诉它哪些变量需要“求导”,它就能帮你自动追踪梯度,也即是导数:
x = torch.tensor([1.0, 2.0], requires_grad=True)
y = x ** 2 + 3 * x
z = y.sum()
z.backward() # 反向传播
print(x.grad) # 输出梯度
#输出为
#tensor([5., 7.])
#可以看到在x=1和x=2这两个点上的变化快慢。
PyTorch 提供了 torch.nn
这个模块,里面封装了一堆搭建神经网络的积木。
例子:定义一个模型,给它一张图(展开成 784(28*28)个数字),它就会告诉你这张图最可能是哪个数字(输出 10 个概率值)。中间它做了两步计算:先提取一些特征(用 ReLU 激活),再把这些特征判断成 10 个数字中的哪一个,并给出概率,看代码:
import torch.nn as nn #引入 PyTorch 提供的神经网络工具箱,这个工具箱里有各种组件(比如线性层、激活函数等)。
model = nn.Sequential(
nn.Linear(784, 128),# 输入层: 784个特征(如28x28像素图像),压缩成128个特征值
nn.ReLU(),# 激活函数
nn.Linear(128, 10),# 输出层: 把前面的128个特征变成10个类别
nn.Softmax(dim=1) # 输出概率分布
)
PyTorch 里常用的优化器是 SGD
、Adam
,使用方法如下:
#设置 Adam 优化器,学习率为 0.001,用于训练模型参数。
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
Dataset
就是你自定义的数据集,DataLoader
可以一批一批地送数据给模型:
from torch.utils.data import DataLoader, TensorDataset
data = torch.randn(100, 784)# 生成100个784维的随机数据样本 (模拟MNIST数据集中的28x28图像)
labels = torch.randint(0, 10, (100,))# 生成100个随机标签,范围0-9 (10分类问题)
dataset = TensorDataset(data, labels)# 将数据和标签封装成TensorDataset
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)# 创建DataLoader,设置批次大小为32,并启用数据随机打乱
建议到官网选择对应系统和 CUDA 版本安装:https://pytorch.org/get-started/locally/
快速方式(CPU 版本):
pip install torch torchvision
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
# 数据预处理
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
testset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False)
# 定义模型
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(28*28, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = x.view(-1, 28*28)
x = F.relu(self.fc1(x))
x = self.fc2(x)
return x
model = Net()
# 损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# 训练模型
for epoch in range(5):
for images, labels in trainloader:
outputs = model(images)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f"第 {epoch+1} 轮训练完成")
# 保存模型
torch.save(model.state_dict(), 'model.pth')
# 加载模型
model = Net()
model.load_state_dict(torch.load('model.pth'))
model.eval() # 切换为评估模式
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
for images, labels in trainloader:
images, labels = images.to(device), labels.to(device)
...
项目 | PyTorch | TensorFlow |
---|---|---|
是否动态图 | 默认动态图 | 2.x后支持动态图 |
代码简洁程度 | 更符合 Python 思维 | 较复杂 |
社区支持 | 学术界主流 | 工业界主流 |
上手难度 | 易上手 | 需要熟悉 Keras |
调试方便 | Pythonic 调试 | 静态图调试麻烦 |
pytorch/examples
PyTorch 真的是目前最简单、好用、灵活的深度学习框架,尤其适合新手和研究人员。
你能理解它的基本逻辑,再配合一些数学原理和神经网络结构,基本就能做很多有趣的 AI 项目了。