Linux下深度学习虚拟环境的搭建与模型训练

曲途光未央2024-05-02 21:08

在深度学习实践中，环境配置是十分重要且免不了的一步。本文以 YOLOv4 模型，介绍在Linux下虚拟环境配置到模型训练的过程。

安装Miniconda：

Miniconda是Anaconda的一个轻量级版本，非常适合用于科学计算和数据处理。

复制代码

wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
bash Miniconda3-latest-Linux-x86_64.sh

创建并激活Python环境：

为YOLOv4创建一个隔离的Python环境，避免依赖冲突。

复制代码

conda create -n yolov4 python=3.8.8
conda activate yolov4

使用以下命令检查PyTorch是否正确安装：

复制代码

import torch
print(torch.cuda.is_available())
print(torch.cuda.device_count())
print(torch.version.cuda)

配置镜像源：

使用清华大学提供的镜像源可以加快包的下载速度。

复制代码

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

安装PyTorch及依赖：

复制代码

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

YOLOv4的训练和测试可以通过多种方式实现，这里我们将介绍使用Darknet框架训练的方法。

YOLOv4在Darknet的源代码可以在GitHub上找到，下载地址是https://github.com/AlexeyAB/darknet，数据集的构建与划分请参考官方的文档进行操作。

YOLOv4 模型训练与测试

复制代码

darknet detector train data/voc2021.data cfg/yolov4.cfg yolov4.conv.137 -map

模型剪枝优化

复制代码

python prune.py --percent 0.5 --weights ./runs/train/exp20/weights/last.pt --data ./data/VOC.yaml

PyTorch多GPU训练

提高训练效率，缩短训练周期。

复制代码

python -m torch.distributed.launch --nproc_per_node 4 --master_port 9527 train.py --workers 8 --device 0,1,2,3 --sync-bn --batch-size 128 --data data/voc.yaml --img 640 640 --cfg cfg/training/yolov7.yaml --weights 'weights/yolov7.pt' --name yolov7 --hyp data/hyp.scratch.p5.yaml

模型测试

对训练好的模型进行效果验证。

复制代码

python test.py --weights weights/best.pt --data data/voc.yaml --img 640 --conf 0.5 --iou 0.65

一些高级特性和工具可以帮助更好地开发和监控模型性能

WandB（Weights & Biases）：用于模型训练过程中的性能监控和可视化。

复制代码

python
import wandb
wandb login

Visdom：另一种数据可视化工具，适合实时数据监控

复制代码

conda install visdom -i https://pypi.douban.com/simple/
python -m visdom.server