Azure ML Studio中使用YOLOv8训练时GPU内存不足问题求助

我在Azure AI Machine Learning Studio的笔记本中，通过pip install ultralytics安装YOLOv8后训练分类模型，出现内存不足错误：

我的配置 #

训练代码 #

import os
from ultralytics import YOLO

# Load a model
model = YOLO('yolov8n-cls.pt')  

# Train the model
results = model.train(
    name='my_classification_run',
    task='classify',
    data='data', 
    batch=16,
    epochs=1,
    imgsz=64,
    device=0,
    cache=True,
    val=True
    )   

数据集详情 #

• 训练集：2385张64x64彩色.jpg，总大小约9.35MB
• 验证集：299张64x64彩色.jpg，总大小约1.16MB

计算资源 #

• 虚拟机规格：Standard_NV12s_v3（12核CPU，112GB系统内存，336GB磁盘）
• GPU：1块NVIDIA Tesla M60处理单元

已尝试的解决方法 #

• 确认CUDA环境可用
• 通过nvidia-smi命令验证NVIDIA驱动已正确安装
• 检查并终止所有闲置GPU进程
• 调整批次大小：尝试过8、16、32及自动批次（autobatch）
• 使用YOLOv8n-cls最小规格预训练模型

错误发生前的关键日志 #

YOLOv8n-cls summary: 99 layers, 1443412 parameters, 1443412 gradients, 3.4 GFLOPs
Transferred 156/158 items from pretrained weights

解决思路 #

1. 禁用数据集缓存：当前开启了cache=True，会将整个数据集解码后加载到内存。虽然原始图片文件小，但解码后的张量占用内存远高于原始文件，加上Azure ML环境本身的后台进程占用，可能触发内存限制。建议将cache=True改为cache=False，改用磁盘实时读取数据。

2. 区分系统内存与GPU显存：注意你提到的112GB是系统内存，而Tesla M60的单卡显存为16GB，错误可能是GPU显存不足而非系统内存。可以在训练前运行nvidia-smi查看初始显存占用，训练时实时监控显存峰值，确认是否是显存耗尽导致的错误。

3. 启用半精度训练：在model.train()中添加half=True参数，使用FP16半精度计算，可大幅降低GPU显存占用，同时几乎不影响模型精度。

4. 进一步降低输入尺寸：尝试将imgsz从64调整为32，进一步减少单批次数据的显存占用。

5. 检查数据集结构：确保data目录完全符合YOLOv8分类数据集要求——必须包含train和val子目录，每个子目录下按类别创建独立文件夹并存放对应类别的图片。结构错误可能导致数据加载逻辑异常，额外消耗内存。

6. 更新ultralytics库：运行pip install --upgrade ultralytics，新版本可能修复了内存泄漏或资源分配相关的bug。

7. 检查Azure ML容器内存限制：Azure ML笔记本可能存在容器级别的内存配额，即使虚拟机有112GB内存，容器可用内存可能被限制。可在笔记本中运行!free -h查看实际可用系统内存，确认是否存在配额限制。

内容的提问来源于stack exchange，提问作者Alex P