# 简介

pytorch框架下VIT模型预训练和微调任务，使用数据集为pytorch版本的ImageNet数据集

[代码地址](https://github.com/facebookresearch/mae)

# 数据集&模型地址

```
#昆山
/public/software/apps/DeepLearning/Data/ImageNet-pytorch
#乌镇
/public/DL_DATA/ImageNet-pytorch
```

模型下载地址（finetune时加载）：

[mae_pretrain_vit_base.pth](https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_base.pth)

[mae_pretrain_vit_large.pth](https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_large.pth)

[mae_pretrain_vit_huge.pth](https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_huge.pth)

# 运行前准备

## 依赖包版本

```
需要安装timm包，且版本必须为0.3.2，一些重要依赖包版本如下：
tensorboard            2.0.2
protobuf               3.20.0
matplotlib             3.5.3
numpy                  1.21.6
absl-py                1.2.0
```

## 报错

```
报错ImportError: cannot import name ‘container_abcs‘ from ‘torch._six‘

这是由于在1.8版本之后container_abcs就已经被移除了，而代码是基于1.8之前别的版本，所以导入方式不同会出现这样的错误：cannot import name ‘container_abcs’ from ‘torch._six’
因此使用不同版本的torch会出现这个问题。

1.8以下版本使用 ：
from torch._six import container_abcs；
1.8以上版本使用（改成这个）：
import collections.abc as container_abcs
int_classes = int
string_classes = str
```

## 环境变量

```
#环境变量
IMAGENET_DIR                  数据集地址
PRETRAIN_CHKPT                模型文件地址
```

# 预训练任务

```
#模型参数
model  模型类型,可选:mae_vit_base_patch16,mae_vit_large_patch16,mae_vit_huge_patch14
预训练不需要加载模型，只需给出模型类型
```

## 单卡

./run_pretrain.sh

```
#!/bin/bash
export HIP_VISIBLE_DEVICE=0
export IMAGENET_DIR=/public/DL_DATA/ImageNet-pytorch
python3 main_pretrain.py  --epochs 1 --model mae_vit_base_patch16 --batch_size 64 --model mae_vit_base_patch16  --data_path ${IMAGENET_DIR}
```

## 多卡

./run_pretrain-4.sh

```
export IMAGENET_DIR=/public/DL_DATA/ImageNet-pytorch
export HIP_VISIBLE_DEVICE=0,1,2,3
OMP_NUM_THREADS=1 python3 -m torch.distributed.launch   --nproc_per_node=4 main_pretrain.py --epochs 1 --accum_iter 4 --model mae_vit_base_patch16 --batch_size 64 --model mae_vit_base_patch16  --data_path ${IMAGENET_DIR}   

```



# 微调任务

```
#模型参数
model  模型类型,可选:vit_base_patch16,vit_large_patch16,vit_huge_patch14,注意这个参数要与PRETRAIN_CHKPT所给模型地址相匹配。例如：vit_base_patch16对应mae_pretrain_vit_base.pth
```

## 单卡

./run-finetune.sh

```
export HIP_VISIBLE_DEVICE=0
export PRETRAIN_CHKPT=/work/home/hepj/model/VIT/mae_pretrain_vit_base.pth
export IMAGENET_DIR=/public/DL_DATA/ImageNet-pytorch 
python3  main_finetune.py \
    --batch_size 32 \
    --model vit_base_patch16 \
    --finetune ${PRETRAIN_CHKPT} \
    --epochs 1 \
    --blr 5e-4 --layer_decay 0.65 \
    --weight_decay 0.05 --drop_path 0.1 --mixup 0.8 --cutmix 1.0 --reprob 0.25 \
    --dist_eval --data_path ${IMAGENET_DIR}

```

## 多卡

./run-finetune-4.sh

```
export HIP_VISIBLE_DEVICE=0,1,2,3
export PRETRAIN_CHKPT=/work/home/hepj/model/VIT/mae_pretrain_vit_base.pth  
export IMAGENET_DIR=/public/DL_DATA/ImageNet-pytorch 
OMP_NUM_THREADS=1 python3 -m torch.distributed.launch --nproc_per_node=4 main_finetune.py \
    --accum_iter 4 \
    --batch_size 32 \
    --model vit_base_patch16 \
    --finetune ${PRETRAIN_CHKPT} \
    --epochs 1 \
    --blr 5e-4 --layer_decay 0.65 \
    --weight_decay 0.05 --drop_path 0.1 --mixup 0.8 --cutmix 1.0 --reprob 0.25 \
    --dist_eval --data_path ${IMAGENET_DIR}
```

# 结果验证

验证使用的模型为mae_finetuned_vit_xxx.pth，下载地址：

[mae_finetuned_vit_base.pth](https://dl.fbaipublicfiles.com/mae/finetune/mae_finetuned_vit_base.pth)

[mae_finetuned_vit_large.pth](https://dl.fbaipublicfiles.com/mae/finetune/mae_finetuned_vit_large.pth)

[mae_finetuned_vit_huge.pth](https://dl.fbaipublicfiles.com/mae/finetune/mae_finetuned_vit_huge.pth)

./run_eval.sh

```
#注意 model 与 resume模型对应
export PRETRAIN_CHKPT=/work/home/hepj/model/VIT/mae_finetuned_vit_base.pth
export IMAGENET_DIR=/public/DL_DATA/ImageNet-pytorch 
python main_finetune.py\
   --eval\
   --resume ${PRETRAIN_CHKPT}\
   --model vit_base_patch16\
   --batch_size 16\
   --data_path ${IMAGENET_DIR}
```