RGB no more: Minimally Decoded JPEG Vision Transformers

An official code release of the paper RGB no more: Minimally Decoded JPEG Vision Transformers\ This repository contains the pipeline for training JPEG ViT, including DCT augmentations.

If you use this repo, please cite our paper:

BibTex:

@InProceedings{Park_2023_CVPR,
    author    = {Park, Jeongsoo and Johnson, Justin},
    title     = {RGB No More: Minimally-Decoded JPEG Vision Transformers},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    month     = {June},
    year      = {2023},
    pages     = {22334-22346}
}

Plain text:

Jeongsoo Park, and Justin Johnson. "RGB no more: Minimally-decoded JPEG Vision Transformers." In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 22334-22346. 2023.

This readme heavily borrows the format by SwinV2.

Pretrained Models

Pretrained models are available as:

It is possible to download using wget using wget [link]. One example:

wget http://www-personal.umich.edu/~jespark/rgbnomore-2023/imgnetDCTViTTi_ep300_75.1.pth

Note: RGB ViT-S is trained on raw ImageNet to reproduce the recipe by Google. Others are trained on 512x512 resized ImageNet.

Usage

Install

Note that modified libjpeg requires Linux to compile properly. Other OS are not supported.

• Clone this repo:

git clone https://github.com/JeongsooP/RGB-no-more.git
cd RGB-no-more

• Create a conda virtual environment and activate it:

conda create -n rgbnomore python=3.10
conda activate rgbnomore

• Install PyTorch>=1.12.1 CUDA>=11.3 (we used PyTorch=1.12.1 and CUDA=11.3 in our paper):

  conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 "llvm-openmp<16.0" -c pytorch

We recommend using llvm-openmp<16.0 until the multiprocessing issue is fixed.

• Install gcc_linux-64==12.2.0, gxx_linux-64==12.2.0 and other necessary libraries:

  conda install -c conda-forge gcc_linux-64==12.2.0 gxx_linux-64==12.2.0 torchmetrics torchinfo tensorboard einops scipy yacs pandas timm imageio iopath psutil

Make sure the versions for gcc_linux-64 and gxx_linux-64 are exactly 12.2.0.

• Install other requirements:

pip install -r requirements.txt

• Compile dct_manip -- a modified libjpeg handler:

  • Open dct_manip/setup.py

  • Modify include_dirs and library_dirs to include your include and library folder.

  • Modify extra_objects to the path containing libjpeg.so

  • Modify headers to the path containing jpeglib.h

  • Run cd dct_manip

  • Run pip install .

Data preparation

• Download the ImageNet dataset from http://image-net.org/.
• Place the tar files and valprep.sh in the same folder.
• The dataloader will automatically organize the data as following:

  imgnet_data
  ├── train
  │   ├── n01440764
  │   │   ├── n01440764_10026.jpeg
  │   │   ├── n01440764_10027.jpeg
  │   │   └── ...
  │   ├── n01443537
  │   │   ├── n01443537_10007.jpeg
  │   │   └── ...
  │   └── ...
  └── val
    ├── n01440764
    │   ├── ILSVRC2012_val_00000293.jpeg
    │   ├── ILSVRC2012_val_00002138.jpeg
    │   └── ...
    ├── n01443537
    │   ├── ILSVRC2012_val_00000236.jpeg
    │   └── ...
    └── ...

  Images will be resized to 512x512 unless --no_resize flag is used.

Pipeline

Sample bash script for all models are included in ./job_bash. Make sure to modify --datapath, --temp_datapath, --num_gpus, --num_cpus, --savepath, --loadpath, --verbose and any other flags to your setting.

• Training from scratch on ImageNet-1K

To train a JPEG-Ti on ImageNet from scratch, run:

python train.py --model_arch=vitti --embed_type=1 --domain=dct --datapath=/path_to/tarfiles --temp_datapath=/tmp/temp_path --indexpaths=assets/indexbase_train.csv,assets/indexbase_val.csv --num_cpus=16 --savepath=/path_to_save/jpegti_model.pth --verbose=2 --train --eval

• Evaluation

To evaluate a pre-trained JPEG-Ti on ImageNet val, run:

python train.py --model_arch=vitti --embed_type=1 --domain=dct --datapath=/path_to/tarfiles --temp_datapath=/tmp/temp_path --indexpaths=assets/indexbase_train.csv,assets/indexbase_val.csv --num_cpus=16 --loadpath=/model_savedpath/jpegti_model.pth --verbose=2 --eval

You can change the --model_arch option to the corresponding model architecture (vitti, vits, vitb, swinv2).

• Throughput

To measure the throughput for 128 iterations, run:

python train.py --model_arch=vitti --embed_type=1 --domain=dct --datapath=/path_to/tarfiles --temp_datapath=/tmp/temp_path --indexpaths=assets/indexbase_train.csv,assets/indexbase_val.csv --num_cpus=16 --loadpath=/model_savedpath/jpegti_model.pth --verbose=2 --benchmark 128