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 are available as:
Model | Val Acc (%) | Link |
---|---|---|
JPEG-Ti | 75.1 | link |
ViT-Ti | 74.1 | link |
JPEG-S | 76.5 | link |
ViT-S | 76.5 | link |
SwinV2-T (DCT, window=8) |
79.4 | link |
SwinV2-T (RGB, window=8) |
79.0 | link |
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.
Note that modified libjpeg
requires Linux to compile properly. Other OS are not supported.
git clone https://github.com/JeongsooP/RGB-no-more.git
cd RGB-no-more
conda create -n rgbnomore python=3.10
conda activate rgbnomore
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.
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.
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 .
tar
files and valprep.sh
in the same folder.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.
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.
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
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).
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