Update README.md

README.md

@@ -29,10 +29,11 @@ ViT主要包括patch embeding、transformer encoder、MLP head三部分：以图
 ### Docker（方法一）
 ```
 docker pull image.sourcefind.cn:5000/dcu/admin/base/jax:0.4.23-ubuntu20.04-dtk24.04.1-py3.10
-docker run -it --network=host --privileged=true --name=docker_name --device=/dev/kfd --device=/dev/dri --group-add video --shm-size=32G  --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v /path/your_code_data/:/path/your_code_data/ imageID /bin/bash
+docker run -it --network=host --privileged=true --name=docker_name --device=/dev/kfd --device=/dev/dri --group-add video --shm-size=32G  --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v /path/your_code_data/:/path/your_code_data/ <imageID> /bin/bash  # <imageID>为以上拉取
 
 cd /your_code_path/vision_transformer
-pip install -r requirements.txt  # flax会强制安装某版本ai包
+pip install flax==0.6.9 # flax会强制安装某版本ai包
+pip install -r requirements.txt
 pip uninstall tensorflow
 pip uninstall jax
 pip uninstall jaxlib
@@ -44,12 +45,12 @@ pip install jaxlib-0.4.23+das1.1.git387bd43.abi1.dtk2404-cp310-cp310-manylinux_2
 ```
 ### Dockerfile（方法二）
 ```
-cd ./docker
-docker build --no-cache -t vision_transformer:latest .
-docker run -it --network=host --privileged=true --name=docker_name --device=/dev/kfd --device=/dev/dri --group-add video --shm-size=32G  --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v /path/your_code_data/:/path/your_code_data/ imageID /bin/bash
+docker build --no-cache -t vit:latest .
+docker run -it --network=host --privileged=true --name=docker_name --device=/dev/kfd --device=/dev/dri --group-add video --shm-size=32G  --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v /path/your_code_data/:/path/your_code_data/ vit /bin/bash
 
 cd /your_code_path/vision_transformer
-pip install -r requirements.txt  # flax会强制安装某版本ai包
+pip install flax==0.6.9 # flax会强制安装某版本ai包
+pip install -r requirements.txt
 pip uninstall tensorflow
 pip uninstall jax
 pip uninstall jaxlib
@@ -71,6 +72,7 @@ jax：0.4.23
 2、其他非特殊库直接按照下面步骤进行安装
 ```
 cd /your_code_path/vision_transformer
+pip install flax==0.6.9 # flax会强制安装某版本ai包
 pip install -r requirements.txt
 pip install tensorflow-cpu==2.13.1
 ```