Merge pull request #135 from SeanSong-amd/update_instructions_readme

update env setup instructions

README.md

@@ -41,12 +41,17 @@ Recently the state space models (SSMs) with efficient hardware-aware designs, i.
 
 ## Envs. for Pretraining
 
-- Python 3.10.13
+- NVIDIA GPUs:
+  - Python 3.10.13
 
-  - `conda create -n your_env_name python=3.10.13`
+    - `conda create -n your_env_name python=3.10.13`
+
+  - torch 2.1.1 + cu118
+    - `pip install torch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1 --index-url https://download.pytorch.org/whl/cu118`
+
+- AMD GPUs:
+  - A [Docker image](https://hub.docker.com/r/rocm/pytorch/tags) is recommended, such as `rocm/pytorch:rocm6.2_ubuntu20.04_py3.9_pytorch_release_2.1.2`. For step-by-step instructions, please refer to [Vision Mamba on AMD GPU with ROCm](https://rocm.blogs.amd.com/artificial-intelligence/vision-mamba/README.html#vision-mamba).
 
-- torch 2.1.1 + cu118
-  - `pip install torch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1 --index-url https://download.pytorch.org/whl/cu118`
 
 - Requirements: vim_requirements.txt
   - `pip install -r vim/vim_requirements.txt`
@@ -73,6 +78,7 @@ Recently the state space models (SSMs) with efficient hardware-aware designs, i.
 | [Vim-tiny<sup>+</sup>](https://huggingface.co/hustvl/Vim-tiny-midclstok)    |       7M       |   78.3   | 94.2 | https://huggingface.co/hustvl/Vim-tiny-midclstok |
 | [Vim-small](https://huggingface.co/hustvl/Vim-small-midclstok)    |       26M       |   80.5   | 95.1 | https://huggingface.co/hustvl/Vim-small-midclstok |
 | [Vim-small<sup>+</sup>](https://huggingface.co/hustvl/Vim-small-midclstok)    |       26M       |   81.6   | 95.4 | https://huggingface.co/hustvl/Vim-small-midclstok |
+| [Vim-base](https://huggingface.co/hustvl/Vim-base-midclstok)    |       98M       |   81.9   | 95.8 | https://huggingface.co/hustvl/Vim-base-midclstok |
 
 **Notes:**
 - <sup>+</sup> means that we finetune at finer granularity with short schedule.
@@ -88,7 +94,7 @@ This project is based on Mamba ([paper](https://arxiv.org/abs/2312.00752), [code
 If you find Vim is useful in your research or applications, please consider giving us a star 🌟 and citing it by the following BibTeX entry.
 
 ```bibtex
-@inproceedings{vim,
+ @inproceedings{vim,
   title={Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model},
   author={Zhu, Lianghui and Liao, Bencheng and Zhang, Qian and Wang, Xinlong and Liu, Wenyu and Wang, Xinggang},
   booktitle={Forty-first International Conference on Machine Learning}