人工智能原理——课程项目二

2025-05-13 Created By Zhang Zhang

项目结构

项目结构如下所示：

.
├── ckpts                           # mae的预训练权重
│   └── ...
├── config_mae.py                   # mae分类器配置文件
├── config_res.py                   # ResNet分类器配置文件
├── data                            # 数据集定义与数据预处理
│   ├── __init__.py
│   ├── dataset.py                  # 数据集定义
│   └── transforms.py               # 数据预处理
├── launch.sh                       # 启动训练的脚本
├── main.py                         # 主函数
├── model                           # 模型定义
│   ├── classifier.py               # 分类器定义
│   ├── mae.py                      # MAE模型定义 **注：代码来源：https://github.com/facebookresearch/mae/blob/main/models_mae.py，对于代码设计思想和框架的分析见实验报告**
│   └── resnet50.py                 # ResNet模型定义    
├── output                          # 输出目录，此处保存检查点，可视化结果与训练日志
│   └── ...
├── README.md                       # 项目文档
├── trainer                         # 封装好的训练器
│   ├── __init__.py 
│   ├── classifier_trainer.py       # mae的分类器训练器
│   ├── mae_trainer.py              # mae训练器
│   └── resnet_trainer.py           # ResNet训练器
├── utils                           # 工具函数
│   ├── __init__.py
│   └── logger.py                   # 日志工具函数
├── environment.yml                 # conda环境设置
└── wandb                           # wandb日志记录结果存放目录
    └── ...

运行代码

以下运行流程在Ubuntu 22.04.5 LTS，CUDA==12.8环境下运行无误，如果遇到其他问题，欢迎联系我：

• email: z-z23@mails.tsinghua.edu.cn
• WeChat: 15393057884

创建conda环境

conda env create -f environment.yml
conda activate garbage_classification

MAE补丁

由于版本兼容性问题，运行MAE分类模型需要打一个补丁，请您按照下面的issue进行修改：

huggingface/pytorch-image-models#420 (comment)

设置wandb环境

本项目使用wandb记录训练日志，您需要注册登录才能使用:

• 在网页https://wandb.ai注册账号
• 在本地运行如下命令进行登录

wandb login

或者，如果您不想使用wandb进行日志记录，您可以修改launch.sh脚本，应用命令行参数--no_wandb，不使用wandb进行日志记录.

下载预训练权重

运行下面的命令下载MAE预训练权重

wget -O ./ckpts/mae_finetuned_vit_base.pth https://dl.fbaipublicfiles.com/mae/finetune/mae_finetuned_vit_base.pth

下载数据集

您需要将数据集下载至/garbage_dataset目录下，在项目根目录下运行如下命令

wget -O ./dataset.zip https://cloud.tsinghua.edu.cn/f/4a89782049fc4f1aa9ce/?dl=1
unzip ./dataset.zip
rm ./data.zip

开始训练！

完成以上环境配置之后，您可以运行下列命令开始训练：

1. 运行ResNet训练：

bash run.sh resnet    # 运行 ResNet

2. 运行MAE训练：

bash run.sh mae       # 运行 MAE

注：MAE模型的训练大约会花费一至两天的时间.

自定义运行

您可以通过1)修改配置文件；2）修改launch.sh中运行主程序的命令行参数的方式自定义您的运行方式.

main.py的命令行参数如下所示：

参数名	默认值	说明
--model_type	mae_classifier	可选：resnet 或 mae_classifier
--config	config_mae.py	配置文件路径
--data_dir	./data	数据集路径
--resume	None	是否从指定模型权重继续训练
--output_dir	./output	输出目录
--mae_checkpoint	None	用于 MAE 微调的预训练权重
--eval_only	关闭状态	若开启，仅进行模型评估
--no_wandb	关闭状态	若开启，不记录 wandb 日志
--seed	42	随机种子