Matryoshka Representation Learning 项目安装与配置指南

Matryoshka Representation Learning 项目安装与配置指南

MRL Code repository for the paper - "Matryoshka Representation Learning" 项目地址: https://gitcode.com/gh_mirrors/mrl/MRL

1. 项目基础介绍

Matryoshka Representation Learning（MRL）是一个用于训练、评估和分析Matryoshka表示的开源项目。该项目旨在解决传统学习表征在维度上信息扩散和部署不灵活的问题。项目主要使用Python编程语言，依赖于深度学习框架PyTorch。

2. 项目使用的关键技术和框架

• Python: 项目的主要编程语言。
• PyTorch: 用于构建和训练神经网络的主流深度学习框架。
• FFCV: 用于高效图像数据加载和预处理的库。
• Matryoshka Linear Layer: MRL中对ResNet50架构进行的微小修改。

3. 项目安装和配置的准备工作与详细步骤

准备工作

• 确保您的系统已安装Python（建议使用Python 3）。

• 安装PyTorch：访问PyTorch官方网站，根据您的系统配置选择合适的安装命令。

• 克隆项目仓库到本地环境：

  git clone https://github.com/RAIVNLab/MRL.git
  cd MRL
  

安装步骤

安装依赖

在项目根目录下执行以下命令安装所需的Python依赖项：

pip install -r requirements.txt

准备数据集

• 将ImageNet数据集转换为FFCV格式，需要设置环境变量IMAGENET_DIR指向包含ImageNet数据集的目录，并设置WRITE_DIR为数据输出目录。

• 执行以下命令准备数据集：

  cd train/
  export IMAGENET_DIR=/path/to/pytorch/format/imagenet/directory/
  export WRITE_DIR=/your/path/here/
  ./write_imagenet.sh 500 0.50 90
  

  请将/path/to/pytorch/format/imagenet/directory/和/your/path/here/替换为实际路径。

训练模型

• 设置CUDA_VISIBLE_DEVICES以指定可用的GPU。

• 使用以下命令开始训练（以下命令假设使用2个A100 GPU，如需修改请相应调整配置）：

  export CUDA_VISIBLE_DEVICES=0,1
  python train_imagenet.py --config-file rn50_configs/rn50_40_epochs.yaml --model.mrl=1 \
  --data.train_dataset=$WRITE_DIR/train_500_0.50_90.ffcv --data.val_dataset=$WRITE_DIR/val_500_uncompressed.ffcv \
  --data.num_workers=12 --data.in_memory=1 --logging.folder=trainlogs --logging.log_level=1 \
  --dist.world_size=2 --training.distributed=1 --lr.lr=0.425
  

  根据需要修改$WRITE_DIR为之前设置的输出目录。

评估模型

• 训练完成后，使用pytorch_inference.py脚本来评估模型。以下是一个评估示例：

  cd inference
  python pytorch_inference.py -path <final_weight.pt> -dataset V1 --mrl
  

  请将<final_weight.pt>替换为实际训练好的模型权重文件。

遵循以上步骤，您可以顺利安装和配置Matryoshka Representation Learning项目，并开始自己的实验。

MRL Code repository for the paper - "Matryoshka Representation Learning" 项目地址: https://gitcode.com/gh_mirrors/mrl/MRL