使用GPU模型训练和目标识别
· 阅读需 15 分钟
在 YOLO 算法中，GPU 在性能和可扩展性方面具有明显优势，适合大规模数据的模型训练和实时目标识别任务；而 CPU 则在成本和开发难度方面具有优势，适用于小规模数据和对实时性要求不高的场景。在使用YOLO（You Only Look Once）进行模型训练和目标识别时，使用GPU和CPU的主要区别体现在以下几个方面：
1. 计算速度
- GPU：GPU 拥有大量的计算核心，能够并行处理大规模数据。在 YOLO 模型训练时，GPU 可以同时对多个图像批次进行卷积、池化等操作，显著加快训练速度。例如在处理大规模数据集时，使用 GPU 训练 YOLOv5 模型可能只需要数小时到数天，而使用 CPU 则可能需要数周时间。在目标识别阶段，GPU 也能快速完成图像的特征提取和目标检测，实现实时或接近实时的识别速度，适合应用于视频监控、自动驾驶等对实时性要求较高的场景。
- CPU：CPU 核心数量相对较少，且主要侧重于串行计算，在处理大规模并行计算任务时效率较低。在 YOLO 模型训练和目标识别中，CPU 需要逐个处理图像数据，导致处理速度较慢，难以满足实时性要求高的应用场景。
2. 内存带宽
- GPU：GPU通常具有更高的内存带宽，能够快速读取和写入大量数据，这对于深度学习中的大规模矩阵运算非常重要。
- CPU：CPU的内存带宽相对较低，处理大规模数据时可能会成为瓶颈。
3. 并行计算能力
- GPU：GPU专为并行计算设计，能够同时处理多个任务，非常适合深度学习中的批量数据处理和矩阵运算。
- CPU：CPU虽然也能进行并行计算，但并行能力远不及GPU，尤其是在处理深度学习任务时。
4. 能耗和成本
- GPU：GPU的能耗较高，且高端GPU价格昂贵。然而，对于深度学习任务，GPU的高效计算能力通常能够抵消其能耗和成本。
- CPU：CPU的能耗相对较低，成本也较低，但在处理深度学习任务时效率较低。
5. 适用场景
- GPU：适合大规模数据集、复杂模型的训练和推理任务，尤其是在需要快速处理大量数据时。
- CPU：适合小规模数据集、简单模型的训练和推理任务，或者在没有GPU的情况下使用。
6. 硬件要求
- GPU：需要支持CUDA（NVIDIA的并行计算平台）的GPU，并且需要安装相应的驱动和深度学习框架（如TensorFlow、PyTorch）的GPU版本。
- CPU：不需要特殊的硬件支持，普通的CPU即可运行，但效率较低。
7. 训练时间
- GPU：训练时间显著缩短，尤其是在处理大规模数据集时，GPU可以在几分钟或几小时内完成训练，而CPU可能需要几天甚至几周。
- CPU：训练时间较长，适合小规模实验或调试。
8. 推理速度
- GPU：目标识别的推理速度非常快，能够实时处理视频流或大量图像。
- CPU：推理速度较慢，可能无法满足实时处理的需求。
GPU环境安装
深度学习框架采用 PyTorch，GPU 为 NVIDIA 显卡，以叉车识别为例，在Windows下安装GPU环境大致为：
1. 确认硬件与软件兼容性
- 显卡：确保你的显卡是 NVIDIA GPU，且计算能力不低于 3.0。可以通过 NVIDIA 官网查询显卡型号的计算能力。
- CUDA：依据显卡驱动版本，挑选适配的 CUDA 版本。通常，显卡驱动版本越高，支持的 CUDA 版本也越高。
- cuDNN：cuDNN 版本要和 CUDA 版本相匹配。
2. 安装显卡驱动
- 访问 NVIDIA 驱动下载页面。
- 按照页面提示，选择显卡产品类型、产品系列、产品型号、操作系统等信息，然后点击 “搜索” 按钮。
- 下载适合你显卡的最新驱动程序，并运行安装程序，按照提示完成驱动安装。
3. 安装 CUDA
- 访问 NVIDIA CUDA Toolkit Archive，选择与显卡驱动兼容的 CUDA 版本进行下载。
- 运行 CUDA 安装程序，按照提示进行安装。在安装过程中，可以选择自定义安装路径和组件。
4. 安装 PyTorch
-
卸载当前 PyTorch（如果不先卸载原有库，安装PyTorch可以存在兼容问题）：
pip uninstall torch torchvision torchaudio -
安装 PyTorch 2.5.x：
进入 https://pytorch.org/get-started/locally/ 可以检查当前环境匹配的PyTorch库，复制页面中的安装命令，如：
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
-
验证
执行以下Python脚本，验证环境是否可用
import torch
import numpy as np
print("PyTorch version:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
print("NumPy Version", np.__version__)执行结果：
PyTorch version: 2.6.0+cu118
CUDA available: True
NumPy Version 1.26.4如果CUDA available返回的结果为True，表示环境已安装成功
模型训练
模型训练脚本与CPU下的并无不同，如果上面CUDA available环境可用就使用GPU训练，否则使用CPU训练
# 导入yolo模块
from ultralytics import YOLO
import torch
def train_model():
# 检查是否有可用的 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# 加载 YOLOv8 预训练模型
model = YOLO("./models/yolov8n.pt")
# 将模型移动到 GPU
model.to(device)
# 训练模型
results = model.train(data="forklift.yaml", epochs=100, imgsz=640, device=device)
if __name__ == '__main__':
# 在 Windows 上需要调用 freeze_support()
from multiprocessing import freeze_support
freeze_support()
# 启动训练
train_model()
GPU训练日志
D:\work\trunk\python\project\yolov8\forklift>python train.py
Using device: cuda
engine\trainer: task=detect, mode=train, model=D:/work/trunk/python/project/yolov8/forklift/models/yolov8n.pt, data=forklift.yaml, epochs=100, time=None, patience=100, batch=16, imgsz=640, save=True, save_period=-1, cache=False, device=cuda, workers=8, project=None, name=train9, exist_ok=False, pretrained=True, optimizer=auto, verbose=True, seed=0, deterministic=True, single_cls=False, rect=False, cos_lr=False, close_mosaic=10, resume=False, amp=True, fraction=1.0, profile=False, freeze=None, multi_scale=False, overlap_mask=True, mask_ratio=4, dropout=0.0, val=True, split=val, save_json=False, save_hybrid=False, conf=None, iou=0.7, max_det=300, half=False, dnn=False, plots=True, source=None, vid_stride=1, stream_buffer=False, visualize=False, augment=False, agnostic_nms=False, classes=None, retina_masks=False, embed=None, show=False, save_frames=False, save_txt=False, save_conf=False, save_crop=False, show_labels=True, show_conf=True, show_boxes=True, line_width=None, format=torchscript, keras=False, optimize=False, int8=False, dynamic=False, simplify=True, opset=None, workspace=None, nms=False, lr0=0.01, lrf=0.01, momentum=0.937, weight_decay=0.0005, warmup_epochs=3.0, warmup_momentum=0.8, warmup_bias_lr=0.1, box=7.5, cls=0.5, dfl=1.5, pose=12.0, kobj=1.0, nbs=64, hsv_h=0.015, hsv_s=0.7, hsv_v=0.4, degrees=0.0, translate=0.1, scale=0.5, shear=0.0, perspective=0.0, flipud=0.0, fliplr=0.5, bgr=0.0, mosaic=1.0, mixup=0.0, copy_paste=0.0, copy_paste_mode=flip, auto_augment=randaugment, erasing=0.4, crop_fraction=1.0, cfg=None, tracker=botsort.yaml, save_dir=runs\detect\train9
Overriding model.yaml nc=80 with nc=1
from n params module arguments
0 -1 1 464 ultralytics.nn.modules.conv.Conv [3, 16, 3, 2]
1 -1 1 4672 ultralytics.nn.modules.conv.Conv [16, 32, 3, 2]
2 -1 1 7360 ultralytics.nn.modules.block.C2f [32, 32, 1, True]
3 -1 1 18560 ultralytics.nn.modules.conv.Conv [32, 64, 3, 2]
4 -1 2 49664 ultralytics.nn.modules.block.C2f [64, 64, 2, True]
5 -1 1 73984 ultralytics.nn.modules.conv.Conv [64, 128, 3, 2]
6 -1 2 197632 ultralytics.nn.modules.block.C2f [128, 128, 2, True]
7 -1 1 295424 ultralytics.nn.modules.conv.Conv [128, 256, 3, 2]
8 -1 1 460288 ultralytics.nn.modules.block.C2f [256, 256, 1, True]
9 -1 1 164608 ultralytics.nn.modules.block.SPPF [256, 256, 5]
10 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']
11 [-1, 6] 1 0 ultralytics.nn.modules.conv.Concat [1]
12 -1 1 148224 ultralytics.nn.modules.block.C2f [384, 128, 1]
13 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']
14 [-1, 4] 1 0 ultralytics.nn.modules.conv.Concat [1]
15 -1 1 37248 ultralytics.nn.modules.block.C2f [192, 64, 1]
16 -1 1 36992 ultralytics.nn.modules.conv.Conv [64, 64, 3, 2]
17 [-1, 12] 1 0 ultralytics.nn.modules.conv.Concat [1]
18 -1 1 123648 ultralytics.nn.modules.block.C2f [192, 128, 1]
19 -1 1 147712 ultralytics.nn.modules.conv.Conv [128, 128, 3, 2]
20 [-1, 9] 1 0 ultralytics.nn.modules.conv.Concat [1]
21 -1 1 493056 ultralytics.nn.modules.block.C2f [384, 256, 1]
22 [15, 18, 21] 1 751507 ultralytics.nn.modules.head.Detect [1, [64, 128, 256]]
Model summary: 129 layers, 3,011,043 parameters, 3,011,027 gradients, 8.2 GFLOPs
Transferred 319/355 items from pretrained weights
TensorBoard: Start with 'tensorboard --logdir runs\detect\train9', view at http://localhost:6006/
Freezing layer 'model.22.dfl.conv.weight'
AMP: checks failed ❌. AMP training on NVIDIA T550 Laptop GPU GPU may cause NaN losses or zero-mAP results, so AMP will be disabled during training.
train: Scanning D:\work\trunk\python\project\yolov8\forklift\forklift-detection\train\labels.cache... 1249 images, 116
val: Scanning D:\work\trunk\python\project\yolov8\forklift\forklift-detection\valid\labels.cache... 119 images, 1 backg
Plotting labels to runs\detect\train9\labels.jpg...
optimizer: 'optimizer=auto' found, ignoring 'lr0=0.01' and 'momentum=0.937' and determining best 'optimizer', 'lr0' and 'momentum' automatically...
optimizer: AdamW(lr=0.002, momentum=0.9) with parameter groups 57 weight(decay=0.0), 64 weight(decay=0.0005), 63 bias(decay=0.0)
TensorBoard: model graph visualization added ✅
Image sizes 640 train, 640 val
Using 8 dataloader workers
Logging results to runs\detect\train9
Starting training for 100 epochs...
Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size
1/100 3.5G 0.9992 1.746 1.297 1 640: 100%|██████████| 79/79 [00:50<00:00, 1.
Class Images Instances Box(P R mAP50 mAP50-95): 100%|██████████| 4/4 [00:02<0
all 119 131 0.688 0.573 0.71 0.365
100 epochs completed in 3.080 hours.
Optimizer stripped from runs\detect\train9\weights\last.pt, 6.3MB
Optimizer stripped from runs\detect\train9\weights\best.pt, 6.3MB
Validating runs\detect\train9\weights\best.pt...
Ultralytics 8.3.94 🚀 Python-3.9.4 torch-2.6.0+cu118 CUDA:0 (NVIDIA T550 Laptop GPU, 4096MiB)
Model summary (fused): 72 layers, 3,005,843 parameters, 0 gradients, 8.1 GFLOPs
Class Images Instances Box(P R mAP50 mAP50-95): 100%|██████████| 4/4 [00:09<0
all 119 131 0.916 0.863 0.91 0.7
Speed: 7.3ms preprocess, 12.8ms inference, 0.0ms loss, 22.4ms postprocess per image
CPU训练日志
D:\work\trunk\python\project\yolov8\forklift>python train.py
E:\Program\Python\Python39\lib\site-packages\ultralytics\nn\tasks.py:708: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
ckpt = torch.load(file, map_location="cpu")
New https://pypi.org/project/ultralytics/8.3.78 available 😃 Update with 'pip install -U ultralytics'
Ultralytics YOLOv8.1.15 🚀 Python-3.9.4 torch-2.5.1+cpu CPU (12th Gen Intel Core(TM) i7-1260P)
engine\trainer: task=detect, mode=train, model=D:/work/trunk/python/project/yolov8/forklift/models/yolov8n.pt, data=forklift.yaml, epochs=100, time=None, patience=50, batch=16, imgsz=640, save=True, save_period=-1, cache=False, device=None, workers=8, project=None, name=train2, exist_ok=False, pretrained=True, optimizer=auto, verbose=True, seed=0, deterministic=True, single_cls=False, rect=False, cos_lr=False, close_mosaic=10, resume=False, amp=True, fraction=1.0, profile=False, freeze=None, multi_scale=False, overlap_mask=True, mask_ratio=4, dropout=0.0, val=True, split=val, save_json=False, save_hybrid=False, conf=None, iou=0.7, max_det=300, half=False, dnn=False, plots=True, source=None, vid_stride=1, stream_buffer=False, visualize=False, augment=False, agnostic_nms=False, classes=None, retina_masks=False, embed=None, show=False, save_frames=False, save_txt=False, save_conf=False, save_crop=False, show_labels=True, show_conf=True, show_boxes=True, line_width=None, format=torchscript, keras=False, optimize=False, int8=False, dynamic=False, simplify=False, opset=None, workspace=4, nms=False, lr0=0.01, lrf=0.01, momentum=0.937, weight_decay=0.0005, warmup_epochs=3.0, warmup_momentum=0.8, warmup_bias_lr=0.1, box=7.5, cls=0.5, dfl=1.5, pose=12.0, kobj=1.0, label_smoothing=0.0, nbs=64, hsv_h=0.015, hsv_s=0.7, hsv_v=0.4, degrees=0.0, translate=0.1, scale=0.5, shear=0.0, perspective=0.0, flipud=0.0, fliplr=0.5, mosaic=1.0, mixup=0.0, copy_paste=0.0, auto_augment=randaugment, erasing=0.4, crop_fraction=1.0, cfg=None, tracker=botsort.yaml, save_dir=runs\detect\train2
Overriding model.yaml nc=80 with nc=1
from n params module arguments
0 -1 1 464 ultralytics.nn.modules.conv.Conv [3, 16, 3, 2]
1 -1 1 4672 ultralytics.nn.modules.conv.Conv [16, 32, 3, 2]
2 -1 1 7360 ultralytics.nn.modules.block.C2f [32, 32, 1, True]
3 -1 1 18560 ultralytics.nn.modules.conv.Conv [32, 64, 3, 2]
4 -1 2 49664 ultralytics.nn.modules.block.C2f [64, 64, 2, True]
5 -1 1 73984 ultralytics.nn.modules.conv.Conv [64, 128, 3, 2]
6 -1 2 197632 ultralytics.nn.modules.block.C2f [128, 128, 2, True]
7 -1 1 295424 ultralytics.nn.modules.conv.Conv [128, 256, 3, 2]
8 -1 1 460288 ultralytics.nn.modules.block.C2f [256, 256, 1, True]
9 -1 1 164608 ultralytics.nn.modules.block.SPPF [256, 256, 5]
10 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']
11 [-1, 6] 1 0 ultralytics.nn.modules.conv.Concat [1]
12 -1 1 148224 ultralytics.nn.modules.block.C2f [384, 128, 1]
13 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']
14 [-1, 4] 1 0 ultralytics.nn.modules.conv.Concat [1]
15 -1 1 37248 ultralytics.nn.modules.block.C2f [192, 64, 1]
16 -1 1 36992 ultralytics.nn.modules.conv.Conv [64, 64, 3, 2]
17 [-1, 12] 1 0 ultralytics.nn.modules.conv.Concat [1]
18 -1 1 123648 ultralytics.nn.modules.block.C2f [192, 128, 1]
19 -1 1 147712 ultralytics.nn.modules.conv.Conv [128, 128, 3, 2]
20 [-1, 9] 1 0 ultralytics.nn.modules.conv.Concat [1]
21 -1 1 493056 ultralytics.nn.modules.block.C2f [384, 256, 1]
22 [15, 18, 21] 1 751507 ultralytics.nn.modules.head.Detect [1, [64, 128, 256]]
Model summary: 225 layers, 3011043 parameters, 3011027 gradients, 8.2 GFLOPs
Transferred 319/355 items from pretrained weights
TensorBoard: Start with 'tensorboard --logdir runs\detect\train2', view at http://localhost:6006/
Freezing layer 'model.22.dfl.conv.weight'
E:\Program\Python\Python39\lib\site-packages\ultralytics\engine\trainer.py:271: FutureWarning: `torch.cuda.amp.GradScaler(args...)` is deprecated. Please use `torch.amp.GradScaler('cuda', args...)` instead.
self.scaler = torch.cuda.amp.GradScaler(enabled=self.amp)
train: Scanning D:\work\trunk\python\project\yolov8\forklift\forklift-detection\train\labels.cache... 1249 images, 116 backgrounds, 0 corrupt: 100%|███████
val: Scanning D:\work\trunk\python\project\yolov8\forklift\forklift-detection\valid\labels.cache... 119 images, 1 backgrounds, 0 corrupt: 100%|██████████|
Plotting labels to runs\detect\train2\labels.jpg...
optimizer: 'optimizer=auto' found, ignoring 'lr0=0.01' and 'momentum=0.937' and determining best 'optimizer', 'lr0' and 'momentum' automatically...
optimizer: AdamW(lr=0.002, momentum=0.9) with parameter groups 57 weight(decay=0.0), 64 weight(decay=0.0005), 63 bias(decay=0.0)
TensorBoard: model graph visualization added ✅
Image sizes 640 train, 640 val
Using 0 dataloader workers
Logging results to runs\detect\train2
Starting training for 100 epochs...
Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size
1/100 0G 1.037 1.885 1.329 1 640: 100%|██████████| 79/79 [03:45<00:00, 2.85s/it]
Class Images Instances Box(P R mAP50 mAP50-95): 100%|██████████| 4/4 [00:05<00:00, 1.44s/it]
all 119 131 0.424 0.252 0.235 0.118
100 epochs completed in 8.448 hours.
Optimizer stripped from runs\detect\train2\weights\last.pt, 6.3MB
Optimizer stripped from runs\detect\train2\weights\best.pt, 6.3MB
Validating runs\detect\train2\weights\best.pt...
Ultralytics YOLOv8.1.15 🚀 Python-3.9.4 torch-2.5.1+cpu CPU (12th Gen Intel Core(TM) i7-1260P)
Model summary (fused): 168 layers, 3005843 parameters, 0 gradients, 8.1 GFLOPs
Class Images Instances Box(P R mAP50 mAP50-95): 100%|██████████| 4/4 [00:04<00:00, 1.15s/it]
all 119 131 0.94 0.84 0.892 0.685
Speed: 0.9ms preprocess, 32.7ms inference, 0.0ms loss, 0.3ms postprocess per image
Results saved to runs\detect\train2
Ultralytics YOLOv8.1.15 🚀 Python-3.9.4 torch-2.5.1+cpu CPU (12th Gen Intel Core(TM) i7-1260P)
Model summary (fused): 168 layers, 3005843 parameters, 0 gradients, 8.1 GFLOPs
val: Scanning D:\work\trunk\python\project\yolov8\forklift\forklift-detection\valid\labels.cache... 119 images, 1 backgrounds, 0 corrupt: 100%|██████████|
Class Images Instances Box(P R mAP50 mAP50-95): 100%|██████████| 8/8 [00:04<00:00, 1.77it/s]
all 119 131 0.948 0.84 0.893 0.686
Speed: 0.9ms preprocess, 32.0ms inference, 0.0ms loss, 0.4ms postprocess per image
通过训练日志可以看到， 同样的数据集100 epochs的训练时间，GPU用了3.080 hours，CPU用了8.448 hours
目标识别
# 导入yolo模块
from ultralytics import YOLO
import torch
# 检查是否有可用的 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 加载模型
model = YOLO("./models/best.pt")
model.to(device)
# 使用测试图片或视频文件测试模型
model.predict(source='./test/9.jpg', save=True, conf=0.5 )
GPU识别结果（使用GPU训练的模型）
CPU识别结果（使用CPU训练的模型）
