PyTorch 是另一个广泛使用的开源深度学习框架,由 Facebook 开发。它以动态计算图为核心特性,在研究和工业界都有大量应用。以下从基础概念、使用流程、代码示例等方面为你介绍 PyTorch。
基础概念
- 张量(Tensors):类似于 NumPy 的多维数组,但可以在 GPU 上加速运算。张量是 PyTorch 中数据的基本载体,模型的输入、输出以及参数都以张量形式存在。
- 自动求导(Autograd):PyTorch 的自动求导机制能够自动计算张量的梯度,这使得构建和训练神经网络变得非常方便。在定义计算图时,PyTorch 会自动记录所有的操作,当调用反向传播时,它会根据这些记录自动计算梯度。
- 动态计算图(Dynamic Computational Graphs):与 TensorFlow 的静态计算图不同,PyTorch 的计算图是动态的,这意味着在运行时可以根据需要改变计算图的结构,使得代码更加灵活,调试也更加容易。
在 AI 解决方案流程中使用 PyTorch
- 数据准备 :PyTorch 提供了
torch.utils.data
模块,其中包含Dataset
和DataLoader
类。Dataset
用于封装数据集,而DataLoader
则用于批量加载数据,支持多线程加载,提高数据处理效率。 - 模型选择与训练 :可以使用
torch.nn
模块构建神经网络模型,通过继承torch.nn.Module
类来定义自定义的模型结构。同时,PyTorch 提供了丰富的优化器(如torch.optim.SGD
、torch.optim.Adam
)和损失函数(如torch.nn.CrossEntropyLoss
)用于模型训练。 - 部署与集成:PyTorch 支持将训练好的模型导出为多种格式,如 ONNX(Open Neural Network Exchange),方便在不同的硬件平台和推理引擎上进行部署。
PyTorch 示例代码
以下是一个使用 PyTorch 构建并训练简单的全连接神经网络来对 MNIST 手写数字数据集进行分类的示例:
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
数据预处理
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
加载MNIST数据集
train_dataset = datasets.MNIST(root='./data', train=True,
download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False,
download=True, transform=transform)
创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)
定义神经网络模型
class SimpleNet(nn.Module):
def init(self):
super(SimpleNet, 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 = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
model = SimpleNet()
定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
训练模型
for epoch in range(5):
running_loss = 0.0
for i, (images, labels) in enumerate(train_loader):
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
print(f'Epoch {epoch + 1}, Loss: {running_loss / len(train_loader)}')
评估模型
correct = 0
total = 0
with torch.no_grad():
for images, labels in test_loader:
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'Test accuracy: {100 * correct / total}%')
代码解释
- 数据加载与预处理 :使用
torchvision
库加载 MNIST 数据集,并进行归一化处理。 - 模型定义 :定义一个简单的全连接神经网络
SimpleNet
,包含两个线性层。 - 损失函数和优化器:使用交叉熵损失函数和 Adam 优化器。
- 模型训练:通过多个 epoch 对模型进行训练,每个 epoch 中遍历训练数据,计算损失并进行反向传播更新模型参数。
- 模型评估:在测试数据集上评估模型的准确率。