在深度学习的浪潮中,选择一款合适的框架是入门的关键一步。当下主流的深度学习框架里,PyTorch凭借其简洁的API设计、灵活的使用逻辑以及强大的社区支持,成为了科研人员和工程实践者的首选之一。无论是刚接触深度学习的新手,还是想要切换框架的开发者,掌握PyTorch都能为你的技术之路添砖加瓦。本文将从基础概念出发,逐步深入到实战应用,带你完整走过PyTorch的学习旅程,让你真正做到"入门看这一篇就够了"。
可能有刚入门的同学会疑惑,为什么众多框架中要优先选择PyTorch?其实答案很简单,PyTorch最核心的优势就是"好懂好用"。它的设计理念与Python高度契合,强调代码的可读性和简洁性,避免了复杂的隐式逻辑,让开发者能更专注于算法本身而非框架的使用技巧。除此之外,动态计算图带来的灵活性、丰富的预训练模型资源、高效的GPU加速能力,以及活跃的社区生态,都让PyTorch成为了深度学习入门的绝佳选择。接下来,我们就从基础开始,一步步揭开PyTorch的神秘面纱。

一、初识PyTorch:它是什么,为什么值得学
在正式上手操作之前,我们先对PyTorch有一个全面的认识,了解它的发展历程、核心优势以及适用场景,这样能帮助我们在后续的学习中更好地理解其设计逻辑。
1.1 PyTorch的发展历程
PyTorch是由Facebook的人工智能研究团队(FAIR)开发的开源深度学习框架,于2016年正式发布。在它诞生之前,深度学习领域已有TensorFlow、Theano等成熟框架,但这些框架大多采用静态计算图,灵活性较差,对于科研人员来说不够友好。PyTorch的出现打破了这一局面,它创新性地采用了动态计算图作为核心,允许开发者在运行时调整计算流程,这一设计让模型的调试和迭代变得更加高效。
发布初期,PyTorch主要在科研社区传播,凭借其易用性迅速积累了大量用户。2019年,PyTorch 1.0版本正式发布,引入了一系列重要功能,包括对ONNX格式的支持、分布式训练工具以及C++前端支持等,这使得PyTorch不仅能满足科研需求,也具备了大规模工业应用的能力。此后,PyTorch的发展进入快车道,全球用户数量持续增长,GitHub上的星标数量突破50万,成为了深度学习领域最流行的框架之一。如今,无论是顶尖高校的科研项目,还是大型科技公司的工业应用,都能看到PyTorch的身影。
1.2 PyTorch的核心优势
PyTorch之所以能得到广泛认可,离不开其独特的优势,这些优势让它在学术界和工业界都备受青睐。
首先是动态计算图带来的灵活性。与静态计算图需要先定义完整流程再运行不同,动态计算图可以边运行边定义,这对于处理复杂的模型结构或者需要根据输入动态调整逻辑的场景非常友好。比如在自然语言处理任务中,输入文本的长度各不相同,动态计算图能轻松应对这种变长输入,而静态计算图则需要额外的适配工作。
其次是极高的易用性。PyTorch的API设计遵循"Pythonic"原则,简洁直观,与Python原生语法高度兼容。对于熟悉Python的开发者来说,上手PyTorch几乎没有学习门槛,很多操作都和使用NumPy类似。比如创建数组、进行矩阵运算等,PyTorch的语法都非常贴近自然语言,让人一看就懂。
调试便捷也是PyTorch的一大亮点。由于它是动态执行的,开发者可以使用Python的标准调试工具,比如PDB、PyCharm调试器等,直接查看每个步骤的中间结果和变量状态。这一点对于排查模型中的错误至关重要,相比之下,静态框架的调试往往需要额外的日志输出,效率大打折扣。
此外,PyTorch拥有强大的社区支持和丰富的资源。官方文档详细全面,社区论坛、GitHub、Stack Overflow等平台上有大量的用户分享的教程、代码和问题解决方案。同时,PyTorch还提供了丰富的预训练模型,比如ResNet、VGG、BERT等,这些模型可以直接用于项目开发,大大降低了开发成本。最后,PyTorch对GPU的支持非常完善,能够高效利用NVIDIA的CUDA库进行并行计算,同时还支持分布式训练,可在多个GPU或服务器上提升训练速度。
1.3 PyTorch的适用场景
凭借其强大的功能和灵活性,PyTorch在多个深度学习应用领域都有出色的表现。
在计算机视觉领域,PyTorch的torchvision库提供了大量的预训练模型和数据处理工具,可轻松应用于图像分类、目标检测、语义分割、图像生成等任务。比如使用ResNet模型进行图像识别,用YOLO模型进行目标检测,这些都可以通过PyTorch快速实现。
自然语言处理是PyTorch的优势领域之一。动态计算图对变长文本的良好支持,让它在文本分类、情感分析、命名实体识别、机器翻译等任务中表现突出。同时,PyTorch的transformers库提供了BERT、GPT等主流预训练语言模型,让开发者能够快速搭建高性能的NLP系统。
生成对抗网络(GANs)的开发和训练也非常适合使用PyTorch。GANs的结构相对复杂,需要频繁调整模型逻辑,PyTorch的灵活性能够让开发者快速迭代模型结构,提升研发效率。无论是图像生成、风格迁移,还是数据增强,PyTorch都能提供有力的支持。
此外,PyTorch在强化学习、时序数据分析等领域也有广泛的应用。在强化学习中,需要根据环境反馈动态调整模型策略,PyTorch的动态计算图能很好地适配这一需求;在时序数据分析任务中,比如语音识别、时间序列预测,PyTorch提供的RNN、LSTM、GRU等模型,以及对变长序列的支持,都能帮助开发者高效完成任务。
二、PyTorch基础:打好入门的核心基石
了解了PyTorch的基本情况后,我们开始学习它的核心基础内容。这一部分主要包括张量(Tensor)操作、GPU加速以及自动求导机制,这些是使用PyTorch进行深度学习开发的必备知识。
2.1 张量:PyTorch的核心数据结构
张量是PyTorch中最基本的数据结构,可以理解为多维数组或矩阵。它与NumPy的数组非常相似,但最大的区别在于张量可以在GPU上进行运算,从而实现加速。掌握张量的基本操作,是后续学习模型构建和训练的基础。
首先,我们需要导入PyTorch库,这是使用所有功能的前提。导入方式非常简单,只需执行import torch即可。接下来,我们学习几种常见的张量创建方式。
如果想要创建一个未初始化的张量,可以使用torch.empty(shape)方法,比如创建一个5行3列的未初始化矩阵,代码为x = torch.empty(5, 3)。需要注意的是,未初始化的张量的值是随机的,取决于内存中的现有数据,通常用于后续会被覆盖的场景。
创建随机初始化的张量可以使用torch.rand(shape)方法,该方法会生成取值在[0,1)之间的随机数。例如x = torch.rand(5, 3),就能得到一个5x3的随机矩阵。如果需要创建全零张量,可以使用torch.zeros(shape),还可以通过dtype参数指定数据类型,比如x = torch.zeros(5, 3, dtype=torch.long)就创建了一个5x3的长整型零矩阵。
最常用的方式是直接从现有数据创建张量,使用torch.tensor(data)方法,其中data可以是列表、元组等可迭代对象。比如x = torch.tensor([5.5, 3]),就创建了一个包含两个元素的张量。
除了创建张量,我们还需要掌握一些基本的张量操作。比如对张量进行算术运算、形状变换等。这里我们以一个简单的示例来演示:首先创建一个2x2的全一张量,并设置requires_grad=True,这个参数的作用是追踪该张量上的所有操作,以便后续进行自动求导。代码如下:x = torch.ones(2, 2, requires_grad=True)。接着对x进行加法操作,y = x + 2,此时y会带有一个grad_fn属性,用于记录生成该张量的操作,这里y的grad_fn就是AddBackward0,表示它是通过加法操作得到的。
我们还可以对y进行更复杂的操作,比如z = y * y * 3,然后计算z的均值out = z.mean()。通过打印z和out,我们可以看到它们都带有grad_fn属性,这说明它们的计算过程都被追踪了。这些操作看似简单,但却是后续模型训练中梯度计算的基础。
2.2 GPU加速:让训练速度飞起来
在深度学习中,模型训练往往需要处理大量的数据和复杂的计算,仅靠CPU很难满足效率需求。GPU凭借其强大的并行计算能力,能够显著提升训练速度,而PyTorch提供了非常便捷的GPU加速方式。
首先,我们需要判断当前环境是否存在可用的GPU,可以使用torch.cuda.is_available()方法。如果返回True,说明有可用的GPU;如果返回False,则只能使用CPU进行计算。代码示例如下:
python
import torch
if torch.cuda.is_available():
print("存在可用的GPU")
else:
print("无可用的GPU,将使用CPU")
如果存在可用的GPU,我们可以通过.to('cuda')方法将张量移动到GPU上。例如,先创建一个张量x = torch.tensor([1.0, 2.0]),然后执行x = x.to('cuda'),就能将x从CPU转移到GPU。此外,也可以在创建张量时直接指定设备,比如x = torch.tensor([1.0, 2.0], device='cuda'),这样创建的张量会直接存储在GPU上。
在模型训练中,不仅需要将数据转移到GPU,还需要将模型也转移到GPU上,方法同样是使用.to('cuda')。示例如下:
python
# 创建一个简单的线性模型
model = torch.nn.Linear(10, 1)
# 创建数据
data = torch.randn(100, 10)
# 将模型和数据转移到GPU
if torch.cuda.is_available():
model = model.to('cuda')
data = data.to('cuda')
需要注意的是,数据和模型必须在同一设备上才能进行计算,否则会报错。另外,数据在CPU和GPU之间的传输会消耗一定的时间,因此在实际训练中,应尽量减少数据传输的次数,以提升效率。
2.3 自动求导:深度学习的核心动力
深度学习模型的训练本质上是通过梯度下降优化损失函数,而梯度计算是这一过程的核心。PyTorch的自动求导(autograd)机制能够自动计算张量的梯度,极大地简化了模型训练的流程。
自动求导的使用非常简单,只需在创建张量时设置requires_grad=True,PyTorch就会追踪该张量上的所有操作。当完成计算后,调用.backward()方法,就能自动计算并存储梯度,梯度值可以通过张量的.grad属性获取。
我们用之前创建的张量来演示自动求导的过程:
python
import torch
# 创建张量并追踪操作
x = torch.ones(2, 2, requires_grad=True)
# 执行一系列操作
y = x + 2
z = y * y * 3
out = z.mean()
# 反向传播计算梯度
out.backward()
# 打印梯度 d(out)/dx
print(x.grad)
运行这段代码后,我们会得到一个2x2的矩阵,每个元素的值都是4.5。这是因为out是z的均值,z = 3*(x+2)²,out = (1/4)sum(z),对x求导后得到out对x的梯度为 (3 2*(x+2))/4,由于x初始值为1,因此每个元素的梯度值为 (323)/4 = 4.5。
需要注意的是,如果要计算梯度的张量不是标量(即维度大于0),那么在调用.backward()时需要传入一个与该张量同形的权重向量,用于加权求和得到标量后再计算梯度。例如:
python
x = torch.randn(3, requires_grad=True)
y = x * 2
while y.data.norm< 1000:
y = y * 2
# 传入权重向量v
v = torch.tensor([0.1, 1.0, 0.0001], dtype=torch.float)
y.backward(v)
print(x.grad)
自动求导机制是PyTorch的核心特性之一,它让开发者无需手动推导复杂的梯度公式,就能轻松完成模型的训练,极大地降低了深度学习的开发门槛。
三、PyTorch进阶:构建和训练神经网络
掌握了基础的张量操作和自动求导后,我们就可以开始学习如何使用PyTorch构建和训练神经网络了。这一部分将涵盖神经网络的构建、数据加载和处理,以及模型的保存和加载等核心内容。
3.1 构建神经网络:torch.nn的使用
PyTorch提供了torch.nn库,专门用于构建神经网络。该库依赖于autograd机制来计算梯度,核心是nn.Module类,所有的神经网络模型都需要继承这个类,并在__init__方法中定义网络层,在forward方法中定义数据的前向传播流程。
下面我们以一个简单的卷积神经网络(CNN)为例,演示如何构建神经网络。这个网络包含两个卷积层、两个最大池化层和三个全连接层,适用于图像分类任务:
python
import torch
import torch.nn as nn
import torch.nn.functional as F
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
# 卷积层1:输入通道1,输出通道6,卷积核5x5
self.conv1 = nn.Conv2d(1, 6, 5)
# 卷积层2:输入通道6,输出通道16,卷积核5x5
self.conv2 = nn.Conv2d(6, 16, 5)
# 全连接层1:输入维度16*5*5,输出维度120
self.fc1 = nn.Linear(16 * 5 * 5, 120)
# 全连接层2:输入维度120,输出维度84
self.fc2 = nn.Linear(120, 84)
# 全连接层3:输入维度84,输出维度10(对应10分类任务)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
# 卷积1 + ReLU激活 + 2x2最大池化
x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
# 卷积2 + ReLU激活 + 2x2最大池化(窗口为方形时可只写一个维度)
x = F.max_pool2d(F.relu(self.conv2(x)), 2)
# 将特征图展平为一维向量,以便输入全连接层
x = x.view(-1, self.num_flat_features(x))
# 全连接1 + ReLU激活
x = F.relu(self.fc1(x))
# 全连接2 + ReLU激活
x = F.relu(self.fc2(x))
# 全连接3(输出层,无激活)
x = self.fc3(x)
return x
# 辅助函数:计算展平后的特征维度
def num_flat_features(self, x):
# 获取除batch维度外的其他维度
size = x.size()[1:]
num_features = 1
for s in size:
num_features *= s
return num_features
# 创建网络实例
net = Net()
print(net)
运行这段代码后,会打印出网络的结构信息。从代码中可以看到,我们在__init__方法中定义了网络的各个层,在forward方法中定义了数据从输入到输出的传播路径。需要注意的是,backward方法(用于反向传播计算梯度)会由autograd自动生成,我们无需手动实现。
创建网络后,可以使用net.parameters()方法获取网络的所有可学习参数,这些参数会被后续的优化器用于更新。例如,通过params = list(net.parameters())可以将参数转换为列表,查看参数的数量和形状。
3.2 数据加载和处理:torch.utils.data的使用
在深度学习项目中,数据的加载和处理是非常重要的一环。PyTorch提供了torch.utils.data模块,其中的DataLoader类可以帮助我们高效地加载数据,支持批量读取、随机打乱、多进程加载等功能,极大地提升了数据处理的效率。
3.2.1 加载内置数据集
PyTorch的torchvision.datasets模块提供了多种常用的公共数据集,比如MNIST、CIFAR10、ImageNet等,我们可以直接下载并使用这些数据集。下面以CIFAR10数据集为例,演示如何加载数据:
python
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
# 定义数据预处理流程
transform = transforms.Compose([
transforms.ToTensor(), # 将PIL图像转换为Tensor,并归一化到[0,1]
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) # 标准化,均值和标准差均为0.5
])
# 下载并加载训练集
trainset = datasets.CIFAR10(
root='./data', # 数据保存路径
train=True, # 表示加载训练集
download=True, # 如果数据不存在则自动下载
transform=transform # 应用预处理
)
# 创建训练集数据加载器
trainloader = DataLoader(
trainset,
batch_size=4, # 批次大小,每次加载4个样本
shuffle=True, # 每个epoch打乱数据
num_workers=2 # 使用2个进程加载数据
)
# 下载并加载测试集
testset = datasets.CIFAR10(
root='./data',
train=False, # 表示加载测试集
download=True,
transform=transform
)
testloader = DataLoader(
testset,
batch_size=4,
shuffle=False, # 测试集无需打乱
num_workers=2
)
在这段代码中,我们首先定义了数据预处理流程,使用transforms.Compose将多个预处理操作组合起来。然后通过datasets.CIFAR10下载并加载训练集和测试集,最后使用DataLoader创建数据加载器。通过数据加载器,我们可以在训练时按批次获取数据,大大方便了模型训练。
3.2.2 自定义数据集
除了使用内置数据集,我们还经常需要处理自己的数据集。此时可以通过继承Dataset类来实现自定义数据集,需要重写__len__和__getitem__两个方法。__len__方法返回数据集的大小,__getitem__方法根据索引返回一个样本(输入数据和对应的标签)。
下面是一个简单的自定义数据集示例:
python
from torch.utils.data import Dataset, DataLoader
class MyDataset(Dataset):
def __init__(self, x_tensor, y_tensor):
# 初始化数据集,保存输入和标签张量
self.x = x_tensor
self.y = y_tensor
def __getitem__(self, index):
# 根据索引返回一个样本
return (self.x[index], self.y[index])
def __len__(self):
# 返回数据集的样本数量
return len(self.x)
# 创建输入和标签张量
x = torch.arange(10) # 输入:0-9的整数
y = torch.arange(10) + 1 # 标签:1-10的整数
# 创建自定义数据集实例
my_dataset = MyDataset(x, y)
# 创建数据加载器
loader = DataLoader(
my_dataset,
batch_size=4,
shuffle=True,
num_workers=0 # 单进程加载
)
# 遍历数据加载器
for x_batch, y_batch in loader:
print("输入批次:", x_batch, "标签批次:", y_batch)
运行这段代码后,会看到数据被按批次加载,并且每个批次的顺序是随机的(因为设置了shuffle=True)。自定义数据集的灵活性很高,可以根据实际需求处理各种类型的数据,比如图像、文本、音频等。
3.3 模型的保存和加载
在模型训练过程中,我们通常需要保存训练好的模型参数,以便后续继续训练、测试或部署。PyTorch提供了简单易用的API来实现模型的保存和加载,主要有两种方式:保存/加载模型参数和保存/加载整个模型。
3.3.1 保存和加载模型参数
这是最常用的方式,只保存模型的参数(通过state_dict()方法获取),不保存模型的结构。这种方式的优点是占用空间小,并且可以灵活地加载到不同的模型实例中(只要模型结构相同)。
保存模型参数的代码如下:torch.save(model.state_dict(), PATH),其中PATH是保存路径,比如./model_params.pth。
加载模型参数时,需要先创建一个与原模型结构相同的模型实例,然后使用load_state_dict()方法加载参数:
python
# 创建模型实例(结构必须与原模型一致)
model = TheModelClass(*args, **kwargs)
# 加载保存的参数
model.load_state_dict(torch.load(PATH))
# 设置模型为评估模式
model.eval()
需要注意的是,load_state_dict()方法接受的是一个字典对象,而不是文件路径,因此需要先使用torch.load()加载保存的参数文件,再传入该方法。另外,加载参数后,需要调用model.eval()方法将模型设置为评估模式,这会将dropout、batch normalization等层切换到评估状态,避免影响评估结果。
3.3.2 保存和加载整个模型
这种方式会将整个模型(包括结构和参数)保存下来,加载时无需手动创建模型实例。代码如下:
python
# 保存整个模型
torch.save(model, PATH)
# 加载整个模型
model = torch.load(PATH)
model.eval()
这种方式的优点是使用方便,加载时直接得到模型实例。但缺点是占用空间较大,并且模型的兼容性较差,如果后续修改了模型的代码,可能会导致加载失败。因此,在实际应用中,更推荐使用保存/加载模型参数的方式。
四、PyTorch进阶技巧:提升开发和训练效率
掌握了基本的模型构建和训练后,我们可以学习一些进阶技巧,进一步提升开发和训练效率。这一部分主要介绍GPU加速的进阶使用、torchvision的高级功能,以及如何使用TensorBoard进行可视化。
4.1 GPU加速的进阶使用
前面我们已经介绍了GPU加速的基本使用,这里补充一些进阶技巧。首先,可以通过torch.device来统一管理设备,避免在代码中多次判断是否存在GPU:
python
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
之后,无论是张量还是模型,都可以使用.to(device)方法转移到指定设备,这样代码会更加简洁。
另外,如果系统中有多个GPU,可以指定使用特定的GPU,比如使用第0个GPU:device = torch.device("cuda:0")。如果想要使用多个GPU进行分布式训练,可以使用torch.nn.DataParallel包装模型:
python
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model)
model.to(device)
这样,模型会自动在多个GPU上并行训练,提升训练速度。
4.2 torchvision的高级功能
torchvision是PyTorch的官方计算机视觉工具库,除了提供数据集和数据预处理功能外,还包含了大量的预训练模型和图像处理工具,能够极大地提升计算机视觉项目的开发效率。
torchvision.models模块提供了多种预训练的图像分类模型,比如ResNet、VGG、AlexNet、EfficientNet等。我们可以直接加载这些预训练模型,用于特征提取或迁移学习。例如,加载预训练的ResNet18模型:
python
import torchvision.models as models
# 加载预训练模型
resnet18 = models.resnet18(pretrained=True)
# 将模型转移到GPU
resnet18.to(device)
# 设置为评估模式
resnet18.eval()
加载预训练模型后,可以直接用于预测,也可以根据自己的任务微调模型参数。迁移学习是一种非常有效的方法,尤其当自己的数据集较小时,使用预训练模型可以显著提升模型性能。
此外,torchvision.transforms模块还提供了多种数据增强操作,比如随机裁剪、随机翻转、颜色抖动等,这些操作可以增加数据集的多样性,提升模型的泛化能力。例如:
python
transform = transforms.Compose([
transforms.RandomResizedCrop(224), # 随机裁剪并调整大小为224x224
transforms.RandomHorizontalFlip(), # 随机水平翻转
transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), # 颜色抖动
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # ImageNet数据集的标准化参数
])
4.3 使用TensorBoard进行可视化
TensorBoard是TensorFlow官方提供的可视化工具,PyTorch也提供了对TensorBoard的支持,可以通过torch.utils.tensorboard模块使用。TensorBoard可以帮助我们可视化模型结构、训练过程中的损失和准确率、输入数据、特征图等,方便我们监控训练过程、调试模型。
4.3.1 启动TensorBoard
首先,需要在命令行中启动TensorBoard,命令为:tensorboard --logdir=runs,其中runs是保存TensorBoard日志文件的目录。启动后,在浏览器中访问http://localhost:6006,即可看到TensorBoard的界面。
4.3.2 记录和可视化数据
使用TensorBoard的核心是创建SummaryWriter对象,通过该对象的方法记录各种数据。例如,记录训练过程中的损失和准确率:
python
from torch.utils.tensorboard import SummaryWriter
import numpy as np
# 创建SummaryWriter对象,日志保存到runs/experiment1目录
writer = SummaryWriter('runs/experiment1')
# 模拟训练过程,记录损失和准确率
for n_iter in range(100):
# 记录训练损失
writer.add_scalar('Loss/train', np.random.random(), n_iter)
# 记录测试损失
writer.add_scalar('Loss/test', np.random.random(), n_iter)
# 记录训练准确率
writer.add_scalar('Accuracy/train', np.random.random(), n_iter)
# 记录测试准确率
writer.add_scalar('Accuracy/test', np.random.random(), n_iter)
# 关闭SummaryWriter
writer.close()
运行这段代码后,在TensorBoard的Scalar页面可以看到损失和准确率随迭代次数的变化曲线。此外,还可以使用add_graph()方法可视化模型结构:
python
# 创建模型和示例输入
model = Net()
images = torch.randn(1, 1, 32, 32) # 示例输入:1个样本,1个通道,32x32大小
# 记录模型结构
writer.add_graph(model, images)
writer.close()
在TensorBoard的Graph页面,可以看到模型的层级结构和数据流向,非常直观。除此之外,TensorBoard还支持可视化图像、嵌入向量、直方图等,功能十分强大。
五、实战案例:用PyTorch训练CIFAR10图像分类模型
理论学习之后,通过实战案例来巩固知识是最好的方式。下面我们将使用PyTorch完整地训练一个卷积神经网络,用于CIFAR10数据集的图像分类任务。CIFAR10数据集包含10个类别的彩色图像,每个类别有6000个样本,共60000个样本,分为训练集(50000个)和测试集(10000个)。
5.1 环境准备和数据加载
首先,导入所需的库,并加载和预处理CIFAR10数据集。这里我们使用数据增强来提升模型的泛化能力:
python
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
# 定义数据预处理和增强
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4), # 随机裁剪,填充4个像素
transforms.RandomHorizontalFlip(), # 随机水平翻转
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)) # CIFAR10的均值和标准差
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])
# 下载并加载训练集
trainset = datasets.CIFAR10(
root='./data',
train=True,
download=True,
transform=transform_train
)
trainloader = DataLoader(
trainset,
batch_size=64,
shuffle=True,
num_workers=2
)
# 下载并加载测试集
testset = datasets.CIFAR10(
root='./data',
train=False,
download=True,
transform=transform_test
)
testloader = DataLoader(
testset,
batch_size=64,
shuffle=False,
num_workers=2
)
# 定义类别名称
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
# 定义设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
5.2 定义网络模型
我们使用前面定义的卷积神经网络模型,稍作调整以适应CIFAR10的3通道输入:
python
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
# 输入通道3(彩色图像),输出通道6,卷积核5x5
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 16 * 5 * 5)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
# 创建模型并转移到设备
net = Net().to(device)
5.3 定义损失函数和优化器
对于分类任务,我们使用交叉熵损失函数(CrossEntropyLoss),优化器使用随机梯度下降(SGD),并添加动量以提升训练稳定性:
python
# 定义损失函数
criterion = nn.CrossEntropyLoss()
# 定义优化器
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
5.4 训练网络
训练过程包括前向传播、计算损失、反向传播和参数更新。我们训练2个epoch(遍历整个数据集2次),并打印训练过程中的损失信息:
python
print("开始训练...")
for epoch in range(2): # 训练2个epoch
running_loss = 0.0
# 遍历训练集
for i, data in enumerate(trainloader, 0):
# 获取输入数据和标签,并转移到设备
inputs, labels = data[0].to(device), data[1].to(device)
# 梯度清零
optimizer.zero_grad()
# 前向传播
outputs = net(inputs)
# 计算损失
loss = criterion(outputs, labels)
# 反向传播
loss.backward()
# 更新参数
optimizer.step()
# 累加损失
running_loss += loss.item()
# 每2000个批次打印一次损失
if i % 2000 == 1999:
print('[%d, %5d] loss: %.3f' %
(epoch + 1, i + 1, running_loss / 2000))
running_loss = 0.0
print("训练完成!")
运行这段代码后,会看到每2000个批次的损失值逐渐下降,说明模型在不断学习。如果使用GPU训练,会明显感觉到训练速度比CPU快很多。
5.5 测试网络
训练完成后,我们需要在测试集上评估模型的性能,计算模型的准确率:
python
print("开始测试...")
correct = 0
total = 0
# 测试时不需要计算梯度
with torch.no_grad():
for data in testloader:
images, labels = data[0].to(device), data[1].to(device)
# 前向传播获取输出
outputs = net(images)
# 获取预测结果(概率最大的类别)
_, predicted = torch.max(outputs.data, 1)
# 累加总样本数和正确预测数
total += labels.size(0)
correct += (predicted == labels).sum().item()
# 计算并打印准确率
print('模型在10000张测试图像上的准确率:%d %%' % (100 * correct / total))
# 查看每个类别的准确率
class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))
with torch.no_grad():
for data in testloader:
images, labels = data[0].to(device), data[1].to(device)
outputs = net(images)
_, predicted = torch.max(outputs, 1)
c = (predicted == labels).squeeze()
for i in range(len(labels)):
label = labels[i]
class_correct[label] += c[i].item()
class_total[label] += 1
# 打印每个类别的准确率
for i in range(10):
print('类别 %5s 的准确率:%2d %%' % (classes[i], 100 * class_correct[i] / class_total[i]))
运行这段代码后,会得到模型在整个测试集上的准确率,以及每个类别的准确率。通常,经过2个epoch的训练,模型的准确率可以达到50%左右,如果增加训练epoch、调整模型结构或优化参数,准确率还可以进一步提升。
5.6 保存模型
训练完成后,保存模型参数以便后续使用:
python
# 保存模型参数
torch.save(net.state_dict(), './cifar10_net.pth')
print("模型参数已保存到 ./cifar10_net.pth")
六、总结与学习建议
本文从PyTorch的简介入手,逐步深入讲解了核心基础、神经网络构建、数据处理、模型保存与加载、进阶技巧,并通过一个完整的实战案例演示了如何使用PyTorch进行深度学习项目开发。通过学习本文,相信你已经对PyTorch有了全面的认识,并具备了一定的实战能力。
PyTorch的学习是一个循序渐进的过程,想要真正掌握它,仅仅阅读教程是不够的,还需要大量的实践。这里给大家几点学习建议:首先,多动手写代码,把本文中的示例代码逐行敲一遍,理解每个步骤的作用;其次,尝试修改代码,比如调整模型结构、改变训练参数、使用不同的数据集,观察结果的变化;再次,阅读PyTorch的官方文档,官方文档是最权威、最全面的学习资源,能够帮助你解决很多细节问题;最后,参与开源项目,通过实际项目积累经验,提升自己的开发能力。
深度学习的领域非常广阔,PyTorch只是一个工具,更重要的是掌握深度学习的核心思想和算法原理。希望本文能为你的深度学习之路打下坚实的基础,祝你在学习和实践中不断进步,实现自己的技术目标。