Vision Transformer的tensorflow环境配置简要说明

Vision Transformer

Tensorflow-GPU版本与CUDA版本对应关系

tensorflow版本	Python版本	cuDnn	CUDA
tensorflow_gpu-2.6.0	3.6-3.9	8.1	11.2
tensorflow_gpu-2.5.0	3.6-3.9	8.1	11.2
tensorflow_gpu-2.4.0	3.6-3.8	8.0	11.0
tensorflow_gpu-2.1.0~2.3.0	3.5-3.8	7.6	10.1
tensorflow_gpu-2.0.0	3.5-3.7	7.4	10
tensorflow_gpu-1.13.0~1.15.0	3.5-3.7	7.4	10
tensorflow_gpu-1.5.0~1.11.0	3.5-3.6	7	9

Vision Transformer使用到了MultiHeadAttention，对应tensorflow-gpu最低版本为2.4.0

安装tensorflow-gpu版本

conda install cudatoolkit=11.0
conda install cudnn==8.0.5.39 -c conda-forge

pip install  tensorflow==2.4.0

指定pip源

# 阿里源
pip install xxx -i https://mirrors.aliyun.com/pypi/simple/
# 腾讯源
pip install xxx -i https://mirrors.tencent.com/pypi/simple/
# 豆瓣源
pip install xxx -i https://pypi.douban.com/simple/
# 清华源
pip install xxx -i https://pypi.tuna.tsinghua.edu.cn/simple/

Vision Transformer tensorflow实现