Docker：使用Nvidia官方的pytorch、tensorflow、TensorRT镜像创建Container容器_业界新闻

发布时间:2024-08-03 15:33

阅读量:0

前言

相信大家在学习新的知识前都遇到过开发环境安装不上，或者环境冲突和版本不匹配的情况，另外当我们想要安装多个版本的支持库时，在本地环境上直接安装往往会导致版本冲突的情况，如果我们使用虚拟机或者WSL技术新建一个完整系统，这又往往需要耗费很长时间，同时在我们学习深度学习等相关技术时，我们需要使用到显卡进行计算，虚拟机调用显卡很不方便，同时CUDA、cuDNN、cuBLAS、TensorRT等GPU计算支持库都有强版本依赖，手动安装需要耗费很长时间，因此本文介绍通过docker的容器技术来实现使用Nvidia官方提供的镜像库创建 container容器。

一、前期准备

本文将以创建一个包含python=3.8、CUDA=12.1、cuBLAS=12.1 、cuDNN=8.9、TensorRT=8.6.1、pytorch的container为例，介绍使用docker创建各种深度学习开发环境的详细流程。下面是Nvidia Container Toolkit的官方链接。

Installing the NVIDIA Container Toolkit — NVIDIA Container Toolkit 1.14.4 documentation编辑https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html本文需要您提前安装好docker，如果您是第一次使用docker，十分建议您安装docker桌面版。

Docker: Accelerated Container Application Development编辑https://www.docker.com/

pytorch环境配置(docker)

拉取pytorch镜像--参考网站(命令--docker pull nvcr.io/nvidia/pytorch:19.01-py3)
启动docker容器--参考网站(命令--sudo nvidia-docker run -it --rm -v local_dir:container_dir nvcr.io/nvidia/pytorch:<xx.xx>-py3)

PyTorch | NVIDIA NGC

二、具体步骤

1.启动容器

docker run -it -e NVIDIA_VISIBLE_DEVICES=0 --gpus "device=0" --name trt-cookbook \ --shm-size 16G --ulimit memlock=-1 --ulimit stack=67108864 \ -v ~:/work \ nvcr.io/nvidia/pytorch:23.04-py3 /bin/bash

这个命令是用于在 Docker 容器中运行一个名为 “trt-cookbook” 的容器，它包含了 NVIDIA 的 CUDA 支持和 PyTorch 深度学习框架。以下是命令的各个部分的解释：

docker run：运行一个 Docker 容器。
-it：在交互模式下运行容器。
-e NVIDIA_VISIBLE_DEVICES=0：设置环境变量 NVIDIA_VISIBLE_DEVICES 的值为 0，这表示仅将第一个 NVIDIA 可见设备（显卡）分配给容器。
--gpus "device=0"：显式指定将第一个 GPU 分配给容器。
--name trt-cookbook：为容器指定一个名称为 “trt-cookbook”。
--shm-size 16G：设置容器内共享内存大小为 16GB。
--ulimit memlock=-1：设置容器的内存锁定限制为无限制，这意味着容器可以使用任意数量的内存。
--ulimit stack=67108864：设置容器的堆栈大小限制为 64MB。
-v ~:/work：将主机的家目录 ~ 挂载到容器的 /work 目录，这样你可以在容器中访问主机上的文件。
nvcr.io/nvidia/pytorch:23.04-py3：使用 NVIDIA 这个 Docker 镜像，它包含了 PyTorch 深度学习框架和其他依赖库的预安装版本。
/bin/bash：在容器中启动一个交互式的 Bash 终端。