VisionDK 开源项目使用教程

VisionDK 开源项目使用教程

1. 项目介绍

VisionDK 是一个基于 PyTorch 的强大基线，用于图像分类和人脸识别任务。该项目提供了多种先进的模型架构和训练策略，旨在帮助开发者快速构建和训练高性能的图像分类和人脸识别模型。VisionDK 支持多种模型，如 ResNet、EfficientNet 和 Swin Transformer，并且提供了丰富的训练选项和工具，以满足不同应用场景的需求。

2. 项目快速启动

安装

首先，建议创建一个独立的虚拟环境来安装 VisionDK：

conda create -n vision python=3.10
conda activate vision

然后，安装 PyTorch 和 VisionDK 的依赖：

conda install pytorch torchvision torchaudio cpuonly -c pytorch  # CPU 版本
# 或者
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia  # CUDA 版本

安装 VisionDK 的依赖包：

pip install -r requirements.txt

训练模型

使用 VisionDK 进行训练非常简单。以下是一个单 GPU 训练的示例：

python main.py --cfgs configs/task/pet.yaml

如果需要使用多 GPU 进行训练，可以使用以下命令：

CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node 4 main.py --cfgs configs/classification/pet.yaml --sync_bn

3. 应用案例和最佳实践

图像分类

VisionDK 支持多种图像分类任务，包括但不限于：

• 动物分类
• 物体识别
• 场景分类

通过配置文件，用户可以轻松调整模型的架构和训练参数，以适应不同的分类任务。

人脸识别

VisionDK 提供了多种人脸识别模型和损失函数，如 ArcFace、CircleLoss 和 MegFace。这些模型和损失函数可以帮助用户在人脸识别任务中取得更好的性能。

4. 典型生态项目

数据增强工具

VisionDK 提供了丰富的数据增强工具，可以帮助用户在训练过程中提高模型的泛化能力。例如，使用以下命令可以进行数据增强的可视化：

cd visiondk
python -m tools.test_augment

模型评估工具

VisionDK 还提供了模型评估工具，用户可以使用这些工具来评估训练好的模型的性能：

python validate.py --cfgs configs/task/pet.yaml

数据集分割工具

为了方便用户准备训练数据，VisionDK 提供了数据集分割工具：

python tools/data_prepare.py --postfix <jpg or png> --root <input your data realpath> --frac <train segment ratio>

通过这些工具，用户可以轻松地准备和处理数据，从而提高训练效率和模型性能。