如何构建集成pytorch的gpu docker

万丈高楼平地起~

每次都会因为环境的问题纠结许久，故写个教程，避免来回折腾；

本人的环境为ubuntu20.04；

docker的版本为20.10.21 ,docker版本需要大于19；

推荐一个集成GPU的dockerhub链接，如下：

https://hub.docker.com/r/anibali/pytorch/tags?page=1&ordering=last_updated

选择需要的镜像即可，本人选择的是：anibali/pytorch:1.10.0-cuda11.3-ubuntu20.04

可以把镜像pull下来，节省时间；

sudo docker run -it --gpus all --network=host --shm-size=32G --name my_gpu_env -v /data:/data -e NVIDIA_DRIVER_CAPABILITIES=compute,utility -e NVIDIA_VISIBLE_DEVICES=all anibali/pytorch:1.10.0-cuda11.3-ubuntu20.04

中间尝试减掉--gpus选项，结果是cpu版本的pytorch；

进入环境后，查看python的版本：

查看相应的torch版本，以及是否支持GPU；

顺便贴上小编的nvidia环境；

水平有限，有问题欢迎随时交流；