接下来介绍使用 Python 的 argparse
模块为 YOLOv5 目标检测模型的训练脚本设置命令行参数解析器。以下是对每个参数的详细解释:
模型配置参数
parser.add_argument('--weights', type=str, default='yolov5s.pt', help='initial weights path')
- 含义:指定预训练模型权重文件的路径。
- 作用:用于模型初始化,如使用
yolov5s.pt
(小模型)、yolov5m.pt
(中等模型)等。
parser.add_argument('--cfg', type=str, default='', help='model.yaml path')
- 含义:模型配置文件(
.yaml
)的路径。 - 作用:定义模型结构(如层数、通道数),为空时使用权重文件自带的配置。
parser.add_argument('--data', type=str, default='../MaskDataSet/data.yaml', help='data.yaml path')
- 含义:数据集配置文件路径。
- 作用:指定训练 / 验证集路径、类别数和类别名称。
parser.add_argument('--hyp', type=str, default='data/hyp.scratch.yaml', help='hyperparameters path')
- 含义:超参数配置文件路径。
- 作用:定义学习率、动量等训练超参数。
训练参数
parser.add_argument('--epochs', type=int, default=300)
- 含义:训练的总轮数。
- 作用:控制训练迭代次数。
parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')
- 含义:所有 GPU 上的总批量大小。
- 作用:每轮训练同时处理的样本数,影响内存使用和训练稳定性。
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes')
- 含义:训练和测试时的输入图像尺寸
[训练尺寸, 测试尺寸]
。 - 作用:图像会被缩放至指定大小后输入模型。
训练策略参数
parser.add_argument('--rect', action='store_true', help='rectangular training')
- 含义:是否使用矩形训练。
- 作用:根据图像长宽比调整批次内图像尺寸,减少填充,提高训练效率。
parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
- 含义:是否恢复最近一次的训练。
- 作用:中断训练后可继续之前的进度。
parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
- 含义:是否仅保存最终模型。
- 作用:减少中间模型的存储空间占用。
parser.add_argument('--notest', action='store_true', help='only test final epoch')
- 含义:是否仅在最后一轮测试。
- 作用:加速训练过程。
parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check')
- 含义:是否禁用自动锚框计算。
- 作用:默认会根据数据集计算合适的锚框尺寸,禁用后使用预设值。
parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters')
- 含义:是否进行超参数进化搜索。
- 作用:自动寻找最优超参数组合。
数据处理参数
parser.add_argument('--cache-images', action='store_true', help='cache images for faster training')
- 含义:是否缓存图像到内存 / 磁盘。
- 作用:加速后续训练轮次的数据加载。
parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
- 含义:是否使用图像权重采样。
- 作用:根据类别分布对图像加权,平衡样本。
设备与分布式训练参数
parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
- 含义:指定训练设备(GPU 编号或 CPU)。
- 作用:如
0
表示使用 GPU0,0,1
表示多 GPU 训练。
parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
- 含义:是否使用多尺度训练。
- 作用:在训练过程中随机调整图像尺寸(±50%),增强模型泛化能力。
parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
- 含义:是否将数据集视为单类别。
- 作用:所有目标均视为同一类别。
parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer')
- 含义:是否使用 Adam 优化器。
- 作用:默认使用 SGD,Adam 通常收敛更快但可能泛化稍差。
parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
- 含义:是否使用跨 GPU 的同步 BatchNorm。
- 作用:仅在分布式训练(DDP)模式下可用。
parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')
- 含义:分布式训练中当前进程的 GPU 编号。
- 作用:PyTorch DDP 自动设置,用户无需手动修改。
其他参数
parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
- 含义:Google Cloud Storage 存储桶路径。
- 作用:用于云端训练时保存模型。
parser.add_argument('--name', default='', help='renames experiment folder apple_exp{N} to apple_exp{N}_{name} if supplied')
- 含义:实验名称后缀。
- 作用:自定义训练结果保存目录的名称。
parser.add_argument('--logdir', type=str, default='runs/', help='logging directory')
- 含义:日志保存目录。
- 作用:存储训练过程中的指标和可视化结果。
parser.add_argument('--workers', type=int, default=0, help='maximum number of dataloader workers')
- 含义:数据加载器的工作线程数。
- 作用:Windows 用户建议设为 0,避免多线程问题。
总结
这些参数覆盖了 YOLOv5 训练的各个方面,包括模型配置、训练策略、数据处理和设备设置。用户可通过命令行灵活调整这些参数,以适应不同的数据集和硬件环境。例如:
python train.py --weights yolov5m.pt --epochs 100 --batch-size 32 --device 0,1
上述命令使用中等模型(yolov5m.pt
)在两个 GPU 上训练 100 轮,每批 32 张图像。