VisionGRU

VisionGRU: A Linear-Complexity RNN Model for Efficient Image Analysis

Kaixuan Yin¹ *, Shicheng Yin¹ *, Weixin Chen¹, Yang Liu^{1 📧}

¹ Sun Yat-sen University

(*) equal contribution, (^📧) corresponding author.

paper

News

Abstract

Convolutional Neural Networks (CNNs) and Vision Transformers (ViTs) are two dominant models for image analysis. While CNNs excel at extracting multi-scale features and ViTs effectively capture global dependencies, both suffer from high computational costs, particularly when processing high-resolution images. Recently, state-space models (SSMs) and recurrent neural networks (RNNs) have attracted attention due to their efficiency. However, their performance in image classification tasks remains limited. To address these challenges, this paper introduces VisionGRU, a novel RNN-based architecture designed for efficient image classification. VisionGRU leverages a simplified Gated Recurrent Unit (minGRU) to process large-scale image features with linear complexity. It divides images into smaller patches and progressively reduces the sequence length while increasing the channel depth, thus facilitating multi-scale feature extraction. A hierarchical 2DGRU module with bidirectional scanning captures both local and global contexts, improving long-range dependency modeling, particularly for tasks like semantic segmentation. Experimental results on the ImageNet and ADE20K datasets demonstrate that VisionGRU outperforms ViTs, significantly reducing memory usage and computational costs, especially for high-resolution images. These findings underscore the potential of RNN-based approaches for developing efficient and scalable computer vision solutions.

Overview

Envs. for Pretraining

• Python 3.11.8

  • conda create -n your_env_name python=3.11.8

• torch 2.5.1 + cu121

  • pip install torch==2.5.1+cu121 torchvision==0.20.1+cu121 torchaudio==2.5.1+cu121 --index-url https://download.pytorch.org/whl/cu121

• Requirements: requirements.txt

  • pip install -r vig/requirements.txt

Train Your VisionGRU

run with this:

NC={num_gpus} B={batch_size} data_dir={path_to_imagenet} torchrun --nproc_per_node={num_gpus} vig/train.py

or export some config before training like this:

export data_dir={path_to_imagenet}

Model Weights

Model	#param.	Top-1 Acc.	Checkpoint
VisionGRU-Ti	30M	82.0	ckpt

Evaluation on Provided Weights

To evaluate VisionGRU-Ti on ImageNet-1K, run:

data_dir={path_to_imagenet} ckpt={path_to_checkpointfile} torchrun --nproc_per_node={num_gpus} vig/val.py

Citation

If you use this code for your research, please cite our paper.

@misc{yin2024visiongrulinearcomplexityrnnmodel,
      title={VisionGRU: A Linear-Complexity RNN Model for Efficient Image Analysis}, 
      author={Shicheng Yin and Kaixuan Yin and Weixing Chen and Enbo Huang and Yang Liu},
      year={2024},
      eprint={2412.18178},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2412.18178}, 
}