图像分类的 InternImage github

图像分类的 InternImage

此文件夹包含用于图像分类的 InternImage 的实现。

用法

安装

• 克隆此 repo：  git clone https://github.com/OpenGVLab/InternImage.git cd InternImage
• 创建 conda 虚拟环境并激活：  conda create -n internimage python=3.7 -y conda activate internimage
• CUDA>=10.2按照cudnn>=7官方进行安装
• 安装PyTorch>=1.10.0并torchvision>=0.9.0使用CUDA>=10.2：

例如，要安装 torch==1.11 和 CUDA==11.3：  pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113  -f https://download.pytorch.org/whl/torch_stable.html

• 安装timm==0.6.11并mmcv-full==1.5.0：  pip install -U openmim mim install mmcv-full==1.5.0 pip install timm==0.6.11 mmdet==2.28.1
• 安装其他要求：  pip install opencv-python termcolor yacs pyyaml scipy
• 编译 CUDA 运算符  cd ./ops_dcnv3 sh ./make.sh # unit test (should see all checking is True) python test.py
• 安装操作员

数据准备

下载。我们提供了以下两种方式来加载数据：

对于标准文件夹数据集，将验证图像移动到标记的子文件夹。文件结构应如下所示：  $ tree data imagenet ├── train │   ├── class1 │   │   ├── img1.jpeg │   │   ├── img2.jpeg │   │   └── ... │   ├── class2 │   │   ├── img3.jpeg │   │   └── ... │   └── ... └── val ├── class1 │   ├── img4.jpeg │   ├── img5.jpeg │   └── ... ├── class2 │   ├── img6.jpeg │   └── ... └── ...

为了提高从大量小文件中读取图像时的速度，我们还支持压缩的 ImageNet，其中包括四个文件：

• train.zip，val.zip：存储用于训练和验证分割的压缩文件夹。
• train.txt，val.txt：分别存储相应 zip 文件和真实值标签的相对路径。确保数据文件夹如下所示：  $ tree data data └── ImageNet-Zip ├── train_map.txt ├── train.zip ├── val_map.txt └── val.zip $ head -n 5 meta_data/val.txt ILSVRC2012_val_00000001.JPEG 65 ILSVRC2012_val_00000002.JPEG 970 ILSVRC2012_val_00000003.JPEG 230 ILSVRC2012_val_00000004.JPEG 809 ILSVRC2012_val_00000005.JPEG 516 $ head -n 5 meta_data/train.txt n01440764/n01440764_10026.JPEG 0 n01440764/n01440764_10027.JPEG 0 n01440764/n01440764_10029.JPEG 0 n01440764/n01440764_10040.JPEG 0 n01440764/n01440764_10042.JPEG 0

对于 ImageNet-22K 数据集，创建一个名为 的文件夹fall11_whole，并将所有图像移动到此文件夹中标记的子文件夹中。然后下载 train-val 分割文件（ & ），并将它们放在 的父目录中fall11_whole。文件结构应如下所示：    $ tree imagenet22k/ imagenet22k/ └── fall11_whole ├── n00004475 ├── n00005787 ├── n00006024 ├── n00006484 └── ...

评估

InternImage要评估在 ImageNet val 上进行预训练的版本，请运行：  python -m torch.distributed.launch --nproc_per_node <num-of-gpus-to-use> --master_port 12345 main.py --eval \ --cfg <config-file> --resume <checkpoint> --data-path <imagenet-path>

例如，InternImage-B使用单个 GPU 进行评估：  python -m torch.distributed.launch --nproc_per_node 1 --master_port 12345 main.py --eval \ --cfg configs/internimage_b_1k_224.yaml --resume internimage_b_1k_224.pth --data-path <imagenet-path>

在 ImageNet-1K 上从头开始训练

论文结果是使用 中的配置训练的模型获得的configs/without_lr_decay。

InternImage要从头开始在 ImageNet 上训练，请运行：  python -m torch.distributed.launch --nproc_per_node <num-of-gpus-to-use> --master_port 12345  main.py \ --cfg <config-file> --data-path <imagenet-path> [--batch-size <batch-size-per-gpu> --output <output-directory> --tag <job-tag>]

使用 Slurm 管理工作。

例如，要InternImage在单个节点上使用 8 个 GPU 进行 300 个 epoch 的训练，请运行：

InternImage-T：  GPUS=8 sh train_in1k.sh <partition> <job-name> configs/internimage_t_1k_224.yaml --resume internimage_t_1k_224.pth --eval

InternImage-S：  GPUS=8 sh train_in1k.sh <partition> <job-name> configs/internimage_s_1k_224.yaml --resume internimage_s_1k_224.pth --eval

InternImage-XL：  GPUS=8 sh train_in1k.sh <partition> <job-name> configs/internimage_xl_22kto1k_384.pth --resume internimage_xl_22kto1k_384.pth --eval

使用 Deepspeed 进行训练

我们支持利用来降低训练大型模型的内存成本，例如具有超过 10 亿个参数的 InternImage-H。要使用它，首先安装以下要求：  pip install deepspeed==0.8.3

然后，您可以在具有 8 个 GPU 的 slurm 系统中启动训练，如下所示（以小型和大型 GPU 为例）。默认零阶段为 1，可以通过命令行参数进行配置--zero-stage。

<span style="background-color:var(--bgColor-muted, var(--color-canvas-subtle))"><span style="color:#1f2328"><span style="color:var(--fgColor-default, var(--color-fg-default))"><span style="background-color:var(--bgColor-muted, var(--color-canvas-subtle))"><code>GPUS=8 GPUS_PER_NODE=8 sh train_in1k_deepspeed.sh vc_research_4 train configs/internimage_t_1k_224.yaml --batch-size 128 --accumulation-steps 4 
GPUS=8 GPUS_PER_NODE=8 sh train_in1k_deepspeed.sh vc_research_4 train configs/internimage_t_1k_224.yaml --batch-size 128 --accumulation-steps 4 --eval --resume ckpt.pth
GPUS=8 GPUS_PER_NODE=8 sh train_in1k_deepspeed.sh vc_research_4 train configs/internimage_t_1k_224.yaml --batch-size 128 --accumulation-steps 4 --eval --resume deepspeed_ckpt_dir
GPUS=8 GPUS_PER_NODE=8 sh train_in1k_deepspeed.sh vc_research_4 train configs/internimage_h_22kto1k_640.yaml --batch-size 16 --accumulation-steps 4 --pretrained ckpt/internimage_h_jointto22k_384.pth
GPUS=8 GPUS_PER_NODE=8 sh train_in1k_deepspeed.sh vc_research_4 train configs/internimage_h_22kto1k_640.yaml --batch-size 16 --accumulation-steps 4 --pretrained ckpt/internimage_h_jointto22k_384.pth --zero-stage 3
</code></span></span></span></span>

🤗 Huggingface 加速与 Deepspeed 的整合

或者，您可以使用我们的集成来使用 deepspeed。  pip install accelerate==0.18.0    accelerate launch --config_file configs/accelerate/dist_8gpus_zero3_wo_loss_scale.yaml main_accelerate.py  --cfg configs/internimage_h_22kto1k_640.yaml --data-path /mnt/lustre/share/images --batch-size 16 --pretrained ckpt/internimage_h_jointto22k_384.pth --accumulation-steps 4 accelerate launch --config_file configs/accelerate/dist_8gpus_zero3_offload.yaml main_accelerate.py  --cfg configs/internimage_t_1k_224.yaml --data-path /mnt/lustre/share/images --batch-size 128 --accumulation-steps 4 --output output_zero3_offload accelerate launch --config_file configs/accelerate/dist_8gpus_zero1.yaml main_accelerate.py  --cfg configs/internimage_t_1k_224.yaml --data-path /mnt/lustre/share/images --batch-size 128 --accumulation-steps 4

内存成本

以下是具有 8 个 GPU 的 InternImage-H 的参考 GPU 内存成本。

• 总批次大小 = 512，每个 GPU 的批次大小为 16，梯度累积步长 = 4。

转换检查点

要将 deepspeed 检查点转换为 pytorch fp32 检查点，您可以使用以下代码片段。  from deepspeed.utils.zero_to_fp32 import convert_zero_checkpoint_to_fp32_state_dict convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir, 'best.pth', tag='best')

然后，你可以best.pth照常使用，例如，model.load_state_dict(torch.load('best.pth'))

由于缺乏计算资源，deepspeed 训练脚本目前仅针对前几个 epoch 进行了验证。如果您在重现整个训练时遇到问题，请提交问题。

提取中间特征

要提取中间层的特征，您可以使用extract_feature.py。

例如，b.png从层中提取特征patch_embed并将levels.0.downsample其保存到“b.pth”。  python extract_feature.py --cfg configs/internimage_t_1k_224.yaml --img b.png --keys patch_embed levels.0.downsample --save --resume internimage_t_1k_224.pth

出口

要从 PyTorch导出InternImage-T到 ONNX，请运行：  python export.py --model_name internimage_t_1k_224 --ckpt_dir /path/to/ckpt/dir --onnx

要从 PyTorch导出InternImage-T到 TensorRT，请运行：  python export.py --model_name internimage_t_1k_224 --ckpt_dir /path/to/ckpt/dir --trt