SimMIM

repository·main·Indexed 22 days ago

A framework for Masked Image Modeling (MIM) that achieves representation learning by predicting raw RGB pixels of masked patches using a lightweight prediction head. The repository provides scripts for pre-training models via main_simmim.py, fine-tuning via main_finetune.py, and evaluating pre-trained models on the ImageNet validation set.

Tokens: 1.1K
Snippets: 4
Records: 4
Agent score: 27%

What's inside SimMIM

Fine-tune pre-trained models

main

To fine-tune a model that was previously pre-trained using SimMIM, use the main_finetune.py script. You must provide the configuration file, the ImageNet data path, and the path to the pre-trained checkpoint using the --pretrained flag.

Command Template:

python -m torch.distributed.launch --nproc_per_node <num-of-gpus-to-use> main_finetune.py \
--cfg <config-file> --data-path <imagenet-path> --pretrained <pretrained-ckpt> [--batch-size <batch-size-per-gpu> --output <output-directory> --tag <job-tag>]

python -m torch.distributed.launch --nproc_per_node 16 main_finetune.py \ 
--cfg configs/swin_base__800ep/simmim_finetune__swin_base__img224_window7__800ep.yaml --batch-size 128 --data-path <imagenet-path> --pretrained <pretrained-ckpt> [--output <output-directory> --tag <job-tag>]

Evaluate provided SimMIM models

main

To evaluate a pre-trained and fine-tuned model on the ImageNet validation set, use the main_finetune.py script with the --eval flag. You must provide a configuration file, a checkpoint to resume from, and the path to your ImageNet data.

Command Template:

python -m torch.distributed.launch --nproc_per_node <num-of-gpus-to-use> main_finetune.py \
--eval --cfg <config-file> --resume <checkpoint> --data-path <imagenet-path>

python -m torch.distributed.launch --nproc_per_node 1 main_finetune.py \
--eval --cfg configs/swin_base__800ep/simmim_finetune__swin_base__img224_window7__800ep.yaml --resume simmim_finetune__swin_base__img224_window7__800ep.pth --data-path <imagenet-path>

Install SimMIM
main
Follow these steps to set up the SimMIM environment. Ensure you have CUDA 11.3 and cuDNN 8 installed according to official NVIDIA guides before proceeding.
1. Create and activate a Conda environment with Python 3.8.
2. Install PyTorch dependencies including torchvision, torchaudio, and cudatoolkit=11.3.
3. Install NVIDIA Apex from source to enable optimized training.
4. Clone the SimMIM repository.
5. Install remaining requirements via requirements.txt.
```
# Create environment
conda create -n SimMIM python=3.8 -y
conda activate SimMIM

# Install requirements
conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch -y

# Install apex
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
cd ..

# Clone SimMIM
git clone https://github.com/microsoft/SimMIM
cd SimMIM

# Install other requirements
pip install -r requirements.txt
```
https://github.com/microsoft/simmim/blob/main/README.md

Pre-train models with SimMIM

main

To perform masked image modeling pre-training, use the main_simmim.py script. This requires a configuration file and the path to the ImageNet training directory.

Command Template:

python -m torch.distributed.launch --nproc_per_node <num-of-gpus-to-use> main_simmim.py \ 
--cfg <config-file> --data-path <imagenet-path>/train [--batch-size <batch-size-per-gpu> --output <output-directory> --tag <job-tag>]

python -m torch.distributed.launch --nproc_per_node 16 main_simmim.py \ 
--cfg configs/swin_base__800ep/simmim_pretrain__swin_base__img192_window6__800ep.yaml --batch-size 128 --data-path <imagenet-path>/train [--output <output-directory> --tag <job-tag>]