ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

PyTorch入门:从环境搭建到DataLoader数据装载完整指南

PyTorch入门:从环境搭建到DataLoader数据装载完整指南

聊PyTorch入门,我观察到大多数新人卡住的第一关其实不是神经网络概念,而是环境搭建加数据装载这一连串脏活。照着网上的教程敲两三行pip命令,torch是装上了,一跑却发现CUDA不可用;好不容易装好环境,数据下载下来了,对着文件夹却又不知道怎么写代码才能把它变成模型真正能吃的东西。这些问题我当年全踩过,所以这篇就把从零配环境到数据实战的完整路径走一遍:安装前怎么做环境决策、怎么用Anaconda隔离Python环境、装完PyTorch之后怎么验证,再到数据环节的Dataset封装和DataLoader装配。适合两类人看:一是刚入门PyTorch、想一条路顺畅把环境和数据跑通的;二是已经装好了但数据处理一直用得别扭、想系统性梳理一遍的开发者。

1. 正式动手前必须先做的环境决策:CPU版和GPU版的选择逻辑

很多人一上来就照着GPU版教程装,装完才发现电脑根本没有NVIDIA显卡,或者显卡型号太老,折腾半天全部白费。所以我一直强调:安装PyTorch的第一步不是敲命令,而是先想清楚自己的机器到底属于哪种情况,这个决定直接影响后面所有安装命令和执行结果。

1.1 先判断你的电脑属于哪种情况

不要盲目追求GPU版本。PyTorch的CPU版和GPU版在代码写法上几乎一样,差别只在于底层计算资源和运行速度。如果你属于下面几种情况,直接用CPU版不仅不丢人,反而更合适:

  • 电脑是Mac,或者是没有独立NVIDIA显卡的Windows笔记本;
  • 有NVIDIA显卡但显存小于4GB,训练稍大一点的模型就很容易显存溢出,先用CPU版练手反而少折腾;
  • 主要做数据分析和简单实验,暂时不训大规模模型。

反过来讲,如果你的电脑是Windows或Linux,并且有一块显存在6GB以上的NVIDIA显卡,那直接上GPU版。这里特别提醒一下:AMD显卡和Intel显卡走不了PyTorch原生的CUDA加速。虽然有一些分支方案(比如ROCm、DirectML)能硬上,但入门阶段加这些额外复杂度,属实不值当。就选CPU版或者NVIDIA显卡的GPU版,够用就行。

1.2 版本匹配表:Python、CUDA、PyTorch三方对齐

环境搭建里90%的问题都出在版本不匹配。网上的安装教程很多是几年前写的,里面的Python版本、PyTorch版本早就过时了。我建议安装前先去PyTorch官网查一下最新支持情况,这里给一张当前比较稳定的参考表:

PyTorch版本支持Python版本可选的CUDA计算版本
2.2.x3.8 - 3.1111.8 / 12.1
2.3.x3.8 - 3.1211.8 / 12.1 / 12.4
2.4.x3.8 - 3.1211.8 / 12.1 / 12.4
2.5.x3.9 - 3.1311.8 / 12.4
2.6.x3.9 - 3.1311.8 / 12.4

对大多数入门场景,我的建议是Python选3.10或3.11,PyTorch选当前最新的稳定版,CUDA选11.8或12.1。这个组合在Windows、原生Linux、WSL2上都比较稳。不要一上来就追最新Python大版本,部分第三方包还没来得及适配,装的时候容易碰壁。

1.3 NVIDIA显卡驱动和CUDA的关系,一句话说清

这里有个非常普遍的误解:很多人以为用GPU版PyTorch就必须单独安装CUDA Toolkit,于是去NVIDIA官网下载了一个2.几GB的完整工具包,结果装完也没解决任何问题。实际上PyTorch官方安装包已经内置了CUDA运行库,你不需要单独安装CUDA Toolkit。你要做的其实只有一件事:把NVIDIA显卡驱动更新到能支持目标CUDA版本的程度。

判断方法很直接。在命令行输入nvidia-smi,看右上角显示的“CUDA Version”,这个数字代表当前驱动支持的最高CUDA版本。比如显示CUDA Version: 12.4,那么安装CUDA 12.1或者11.8版本的PyTorch都没问题。注意驱动版本向下兼容,但PyTorch要求的CUDA版本最好不要超过驱动支持的版本,否则就会出现装上了却用不了GPU的情况。

2. 用Anaconda给PyTorch建个“独立房间”:创建与激活环境的完整操作

现在进入实际操作。先说结论:不要拿系统自带的Python直接装PyTorch,强烈建议先装Anaconda(或者轻量一点的Miniconda)。这一步我当年偷懒跳过,后来花了很多时间擦屁股。

2.1 为什么非要Anaconda,系统Python不行吗?

第一是环境隔离。比如一个项目需要PyTorch 2.2,另一个项目需要PyTorch 2.5,用conda建两个环境互不干扰;如果共用一个系统环境,装来装去大概率变成后装的覆盖先装的,之前的项目说挂就挂。

第二是避免依赖冲突。PyTorch底层依赖numpy等一批数值计算库,而numpy恰恰是版本兼容问题的重灾区。conda在安装时会把依赖关系一起梳理好,比逐个pip硬装省心很多。

第三是卸载方便。环境不要了,直接conda env remove -n 环境名,删得干干净净不留垃圾。要是装在系统Python里,哪天想清理会发现不知道哪些包是哪个项目留下的。

2.2 安装Anaconda时最容易忽略的选项

去官网下载Anaconda最新版就行,安装过程有两个点容易被忽略。一是Windows下安装时,建议勾选把Anaconda加入PATH环境变量。如果不勾,之后在cmd或者PowerShell里输入conda就会提示找不到命令,只能每次用Anaconda Prompt,很麻烦。二是安装路径尽量用默认位置,不要放在带中文或空格的目录下,否则后边的pip、conda命令都可能踩编码坑。

Linux或WSL下安装其实更简单。下载.sh安装包,执行这个命令:

bash Anaconda3-2024.10-1-Linux-x86_64.sh

安装过程中一路输入yes,最后source ~/.bashrc让环境变量生效。装完以后在终端输入conda --version能输出版本号就算成功。

2.3 创建环境与激活环境的完整命令

打开终端,执行下面两条命令:

conda create -n pytorch_env python=3.11 conda activate pytorch_env

解释一下每个参数。-n后面跟的是环境名,这里叫pytorch_env,具体名字随你,但建议用项目名,不要叫test或者test1这种毫无辨识度的名字。等环境建好后,安装PyTorch和后续跑训练,都要先conda activate进入这个环境再操作。

激活成功后,命令行开头会出现(pytorch_env)字样,这就说明已经进入独立环境了。如果你在Windows上发现conda activate不好使,多半是conda版本太旧,先执行conda update conda升级一下再试。顺带一提,conda create命令还可以加-y参数跳过确认提示,比如conda create -n pytorch_env python=3.11 -y,一条命令直接建好。

2.4 换源加速:国内环境下的必要动作

如果直接就着conda官方源创建环境,下载速度大概率让人抓狂。建议直接配置国内镜像源,以清华源为例,Linux和macOS下执行:

conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free conda config --set show_channel_urls yes

Windows用户在命令行配置方式也是一样的。pip下载慢的话也可以换源,临时指定最方便:

pip install xxx --index-url https://pypi.tuna.tsinghua.edu.cn/simple

这里有个容易踩的坑:如果你安装PyTorch用的是官方下载地址(download.pytorch.org),那么即使你配了pip清华源也没用,因为官方安装包的下载域名根本不走PyPI。这种情况要么直连,要么换网络环境,这个区别下面安装部分会再强调。

3. 安装PyTorch后必须做的三件事:版本验证、CUDA检测与常见故障定位

环境创建好,终于进入安装环节。安装本身不难,难的是装完之后怎么确认“真的装对了”,以及出了问题怎么快速定位。

3.1 从官网获取匹配的安装命令

最省事的方式是打开PyTorch官网首页(pytorch.org),在Install页面里选好操作系统、包管理器(推荐Pip)、Python版本和CUDA版本,页面会自动生成对应的安装命令。比如在Linux下选Pip和CUDA 12.1,会得到:

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

如果是CPU版,命令是:

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

顺手解释一下这三个包的关系。torch是核心深度学习框架;torchvision用于处理图像数据,提供常见数据集和预训练模型;torchaudio用于处理音频。如果是入门做图像方向,至少装torch和torchvision这两个,绝大多数图像数据操作都要用到torchvision。别为了省那点下载时间只装一个torch,后面用到的时候还得补。

3.2 三个验证步骤,缺一不可

装完后不要急着写模型,先把下面三行代码跑一遍:

import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else "未检测到GPU")

第一行确认torch确实装进来了。第二行是重中之重,GPU版安装成功后这一行应返回True。第三行会输出你显卡的型号名称。如果是CPU版,第二步返回False是正常现象,只要第一步输出了版本号,安装就算成功,不要自己吓自己。

验证通过后,我建议顺手做一件事:把这三行存成一个check_env.py,放在项目根目录下。以后换了电脑、换了环境,先跑一遍这个脚本,几秒钟就知道当前环境撑不撑得起训练任务。

3.3 GPU版显示不可用的排查思路

如果你明明装的是cu121之类的GPU版本,torch.cuda.is_available()却返回False,按顺序排查下面几个位置:

一是显卡驱动没有正常安装。在命令行执行nvidia-smi,如果提示command not found,说明驱动没装好。去NVIDIA官网下载对应型号的驱动,装完重开一个命令行窗口再验证。

二是驱动版本太老,撑不起PyTorch要求的CUDA计算能力。比如驱动只支持CUDA 10.2,你却装了CUDA 12.1版本的PyTorch,自然跑不动。这种情况要么降级PyTorch,要么升级驱动。

三是激活错了环境。很多人电脑上装了多个conda环境,或者系统Python和conda环境里都有torch。以为自己激活的是项目环境,实际import到的却是另一个位置里的CPU版torch。排查方法:

python -c "import torch, sys; print(torch.__file__); print(sys.executable)"

看输出的torch路径和你激活的环境是否一致。不一致,说明当前Python解释器就不是环境里的那个。

四是WSL2用户重点关注。在WSL2里想用GPU,必须先在Windows宿主系统上安装NVIDIA驱动,WSL内部不需要也不能装显卡驱动。确认WSL版本,在Windows PowerShell里执行wsl --set-version Ubuntu-20.04 2可以切换到WSL2。

3.4 一个很常见的报错:找不到适配的torch安装包

如果在pip安装阶段出现ERROR: No matching distribution found for torch,大概率是三个原因之一。第一,Python版本太新或者太老,PyTorch官方还没发布对应的wheel包,比如Python刚出大版本时经常遇到;第二,网络连不上下载服务器,一直超时,烧了代理也一样;第三,操作系统位数不对,比如32位系统。

遇到这个报错,先查Python版本,再查网络,最后确认系统位数,一般都能定位到问题。别盲试那些老教程里的“强制安装”参数,很可能装出个残缺版本,后面用起来更难受。

4. 数据实战第一步:把任意目录下的图像数据封装成Dataset类

环境通了,接下来进入数据实战。官方文档里经常会用到torchvision.datasets.ImageFolder这类现成接口,下载即用。但真实项目中的数据很少长成官方数据集的样子,基本上就是一堆散落在不同文件夹里的文件。所以学会写自定义Dataset类,才算是数据实战的入门必修课。

4.1 Dataset类到底在做什么

可以把Dataset理解成一个“知道怎么取数据、也知道自己有多长数据的长条”。DataLoader在训练过程中会反复向它提两个问题:一是总共有多少条样本?二是请把第N条样本取出来给我。只要Dataset能准确回答这两个问题,后边的数据流水线就顺畅了。

因此自定义Dataset类只需要实现三个方法,缺一不可:

  • __init__:初始化数据路径、标签映射、transform变换等;
  • __len__:返回整个数据集的样本总数;
  • __getitem__:根据索引返回一条样本和对应的标签。

4.2 一个可直接复用的图像分类Dataset写法

假设现在数据目录结构是:

data/train/ cat/001.jpg cat/002.jpg dog/001.jpg dog/002.jpg

写一个适用于这种结构的ImageFolderDataset:

from torch.utils.data import Dataset from PIL import Image import os class ImageFolderDataset(Dataset): def __init__(self, root_dir, transform=None): self.root_dir = root_dir self.transform = transform self.image_paths = [] self.labels = [] # 按类别名排序,保证每次运行时同一个类别的标签一致 class_names = sorted(os.listdir(root_dir)) self.class_to_idx = {cls: i for i, cls in enumerate(class_names)} for cls in class_names: cls_dir = os.path.join(root_dir, cls) for fname in os.listdir(cls_dir): if fname.lower().endswith(('.jpg', '.jpeg', '.png')): self.image_paths.append(os.path.join(cls_dir, fname)) self.labels.append(self.class_to_idx[cls]) def __len__(self): return len(self.image_paths) def __getitem__(self, index): img = Image.open(self.image_paths[index]).convert('RGB') label = self.labels[index] if self.transform: img = self.transform(img) return img, label

几个值得注意的细节。一是用sorted(os.listdir(root_dir))给类别排序,这能保证文件夹在磁盘上顺序变了,标签分配依然稳定。二是convert('RGB')把灰度图、带透明通道的PNG统统统一成RGB三通道,防止训练时batch里混入不同通道数的图导致shape错乱。三是在__getitem__里做transform而不是在__init__里提前一次性做完。这样每次按需读取、动态处理,既省内存又能自然地做数据增强。如果一次性把所有图都预处理完存内存,数据量大一点就爆了。

4.3 从PIL图像到张量:transform的正确用法

图像读进来是PIL对象,不能直接扔进神经网络。需要通过transforms.Compose把“改尺寸、转张量、归一化”等步骤串成一条流水线:

from torchvision import transforms transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] ) ]) dataset = ImageFolderDataset(root_dir='data/train', transform=transform)

这里ToTensor()极其关键。它同时做了两件事:把HWC格式的数组转成CHW通道优先格式,再把像素值从0到255缩放到0到1之间。PyTorch内部所有卷积层、池化层都默认输入是CHW格式,很多人把CV读图的HWC直接丢进模型,shape对不上才发现问题。Normalize这步用的是ImageNet数据集的统计均值方差,如果训练自己的小数据集且用了预训练模型,通常保留这个设置即可;如果完全从零训练,也可不归一化或用自己算的均值方差。

5. 数据实战第二步:用DataLoader和transform把数据打包进训练循环

Dataset只是把“取单条数据”的规则定义好了,真正在训练时批量产出数据的是DataLoader。很多人以为数据加载就是把Dataset传进去那么简单,实际用起来参数不对、多进程不开、设备不迁移,到处是坑。

5.1 DataLoader的关键参数,一个都不能乱设

from torch.utils.data import DataLoader dataloader = DataLoader( dataset, batch_size=32, shuffle=True, num_workers=4, pin_memory=True, drop_last=False )

逐个解释一下。

batch_size是每批样本数量,直接决定显存占用和训练稳定程度。新手先用32起步,显存还有余量再加到64,出现OOM就把32降到16。

shuffle=True会在每个epoch开始时打乱样本顺序,防止模型学到数据集原始顺序的假象。验证集和数据测试环节通常设False,保证结果可复现。

num_workers是多进程预加载数据的工作进程数。Windows下建议设2,设高了一些机器会莫名卡住;Linux可以设成CPU核心数。调试阶段建议先设0,让主进程串行加载,数据管道有个什么问题一眼就能看出来;正式训练再开多进程提速。

pin_memory=True配合GPU训练时非常有用,它把数据固定在内存中,减少数据拷到显存时的开销。前提是训练设备是GPU,如果纯CPU训练设不设影响不大。

drop_last控制最后一批样本数不足batch_size时是否直接丢弃。训练阶段一般设False,最后那个不完整的batch也参与训练;但某些用批归一化的场景,或者想保持每批大小一致时,就会设True把不足一批的尾巴丢掉。

5.2 在训练循环中实际使用DataLoader

写一个极简训练循环,看看数据是怎么真正流动起来的:

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = SomeModel() model.to(device) for epoch in range(num_epochs): for images, labels in dataloader: images = images.to(device) labels = labels.to(device) outputs = model(images) loss = criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step()

这段代码里最容易犯的错就是忘了把images和labels挪到GPU上。数据在CPU、模型在GPU,一跑就报RuntimeError: Expected all tensors to be on the same device。更麻烦的是这个问题不一定在第一个epoch出现,有可能数据跑了一半才爆出来,因为模型和数据的设备不匹配经常是动态的。所以最稳妥的写法就是在循环开头统一device变量,然后每个batch都to(device)。这里还有个细节,optimizer.zero_grad()最好在loss.backward()之前调用,避免上一次迭代的梯度残留导致参数更新错误。

5.3 实战中的验证技巧:先打印批次形状再训练

面对一个不熟的数据集,最忌讳的就是直接开训。第一件做的是“看一眼数据长什么样”:

for images, labels in dataloader: print(images.shape, labels.shape) print(labels[:5]) break

正常输出大概是:

torch.Size([32, 3, 224, 224]) torch.Size([32])

意思是一批32张图,3个通道,宽高224x224,标签是32个数值。如果第二维不是3而是1,说明灰度图混进来了,需要检查convert('RGB')那一步;如果第二个值变成了非224的长条,就要确认Resize是否生效。形状确认无误再开始训练,这一步花的几分钟能省下后面几个小时的报错定位时间。

5.4 训练集和验证集的拆分,以及标签不平衡的入门认识

最后说两个真实项目中避不开的问题。第一是训练集和验证集的拆分。最简单的方式是按目录手动拆,数据量大的场景还可以写个脚本按比例随机拆分,random.shuffle打乱它再切分即可。这里必须注意,验证集的transform里不要包含随机翻转、随机裁剪这类数据增强操作,验证集也尽量保持shuffle=False,否则每次验证结果波动很大,不方便判断模型真实水平。

第二是标签不平衡。比如猫图有1000张、狗图只有100张,模型训练完会严重偏向猫这个类别。入门阶段最简单的方案是使用torch.utils.data.WeightedRandomSampler做加权采样,给少数类分配更高采样概率;或者直接用带权重的交叉熵损失函数。这个问题不需要一开始就深入解决,但脑子里得有这根弦,遇到数据集中类别数量差异大的时候,知道往哪个方向找答案。


最后再分享一个我自己的习惯。每次新建PyTorch项目,我第一件事不是写模型,而是先跑通两个验证脚本:环境验证脚本负责确认torch和CUDA可用,数据验证脚本负责打印一个batch的形状和标签。确保从import torch到打出一个批次数据,全链路无报错,再开始写模型代码。这两步看起来基础,却是整个项目最值得建立的一条基准线。环境这个事,前期花一小时理顺,后面省下的时间往往是以天为单位算的。

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进