使用 Docker + NVidia 运行 Pytorch 验证 GPU 是否真的有效

一、引言

在运行 NGC 深度学习框架容器之前，您的 Docker 环境必须支持 NVIDIA GPU。若要运行容器，请发出相应的命令，如运行容器中所述，并指定注册表、存储库和标记。

如果前面的指导文章没注意到，可以考虑先阅读一下我的文章了解先决条件：

• 《大数据及机器学习带 GPU 必须干的第一件事》
• 《Windows 下让 Docker Desktop 关联上 NVidia GPU》

二、拉取镜像

这个过程不需要考虑网络问题，比较顺利，但由于镜像比较大，大约 22GB，大概需要 N 分钟，看你家网络了。不过第一次很值得，后继又不会一直拉取。

$ docker pull nvcr.io/nvidia/pytorch:24.01-py3
...
Digest: sha256:afd682405d620a620f61f38cb9d9bbc6a5230817699a48e9ed193546e81fb2ee
Status: Downloaded newer image for nvcr.io/nvidia/pytorch:24.01-py3
nvcr.io/nvidia/pytorch:24.01-py3

考虑使用它的原因，是因为阅读了它的构建文件，并且里面有两个重要的属性有在扩展知识说明。

ENV CUDA_MODULE_LOADING=LAZY
ENV TORCH_CUDNN_V8_API_ENABLED=1

三、代码测试

代码不多，主要用于初步了解，后续方便对症下药。

3.1、查看 Python 版本

首先核对一下它使用的 python 版本是多少？答案是 3.10 版，一共两种办法；

3.1.1、容器外部执行

$ docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:24.01-py3 python --version

=============
== PyTorch ==
=============

NVIDIA Release 24.01 (build 80741402)
PyTorch Version 2.2.0a0+81ea7a4

Container image Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

Copyright (c) 2014-2023 Facebook Inc.
Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
Copyright (c) 2011-2013 NYU                      (Clement Farabet)
Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
Copyright (c) 2015      Google Inc.
Copyright (c) 2015      Yangqing Jia
Copyright (c) 2013-2016 The Caffe contributors
All rights reserved.

Various files include modifications (c) NVIDIA CORPORATION & AFFILIATES.  All rights reserved.

This container image and its contents are governed by the NVIDIA Deep Learning Container License.
By pulling and using the container, you accept the terms and conditions of this license:
https://developer.nvidia.com/ngc/nvidia-deep-learning-container-license

NOTE: The SHMEM allocation limit is set to the default of 64MB.  This may be
   insufficient for PyTorch.  NVIDIA recommends the use of the following flags:
   docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 ...

Python 3.10.12

3.1.2、容器内部执行

$ docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:24.01-py3
...也会有和上面一样的 Copyright 和 Note 画面
root@xxd4abbe28xx:/workspace# python --version
Python 3.10.12

3.2、测试是否 CUDA 有效

使用容器外部执行，打印出 True，证明结果有效！

$ docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:24.01-py3 python

上面的 exit 要改成 exit() 才能退出 python 环境。

四、扩展知识

4.1、想调节优化共享内存怎么办？

如果对多线程数据加载器使用多处理，则运行容器的默认共享内存段大小可能不够。特别是某些数据库应用或者机器学习应用可能需要较大的共享内存。

因此，应通过发出以下命令之一来增加共享内存大小：（默认值是 64 MB）

4.1.1、和主机一起共享

$ docker run --gpus all --ipc=host -it --rm nvcr.io/nvidia/pytorch:24.01-py3

以我本机分配给 Docker 16GB 为例，该容器将通过 --ipc=host 拥有 16GB。

4.1.2、自定义调整

$ docker run --gpus all --shm-size=2GB -it --rm nvcr.io/nvidia/pytorch:24.01-py3

通过设置 --shm-size=2GB，该容器将拥有 2GB。

4.2、两个重要的环境变量

在 nvcr.io/nvidia/pytorch:24.01-py3 镜像的构建里，有下面的重要的两行代码：

ENV CUDA_MODULE_LOADING=LAZY
ENV TORCH_CUDNN_V8_API_ENABLED=1

这两个环境变量与 CUDA 和 cuDNN 在 PyTorch 中的使用有关。

4.2.1、延迟加载

ENV CUDA_MODULE_LOADING=LAZY

这个环境变量设置了 CUDA 模块的加载方式为 "LAZY"，即延迟加载。

延迟加载意味着 CUDA 模块将在实际需要时才会被加载，而不是在程序启动时就立即加载。

这种延迟加载的方式可以减少不必要的内存占用，并加快程序的启动速度。

当 PyTorch 在运行时实际需要使用 CUDA 功能时，相关的 CUDA 模块才会被动态加载。

对于 Java 开发的朋友们，可以从 Hibernate 的 Entity 或 Spring 的 Bean 懒加载去理解。

4.2.2、GPU 加速

ENV TORCH_CUDNN_V8_API_ENABLED=1

这个环境变量启用了 PyTorch 中的 cuDNN v8 API。

cuDNN（CUDA Deep Neural Network library）是一个用于深度神经网络的 GPU 加速库，它提供了高度优化的实现，用于卷积、池化、归一化等常见的深度学习操作。

在 PyTorch 中，cuDNN 可以显著加速神经网络的训练和推理过程。

通过将 TORCH_CUDNN_V8_API_ENABLED 设置为 1，PyTorch 将使用 cuDNN v8 的 API，以充分利用其性能优化和新特性。

4.2.3、汇总说明

这两个环境变量的设置有助于优化 PyTorch 在 GPU 上的性能。延迟加载 CUDA 模块可以减少不必要的内存占用和加快启动速度，而启用 cuDNN v8 API 则可以利用其性能优化来加速深度学习任务。

PyTorch NGC 容器已经针对 GPU 加速进行了优化，并包含了一组经过验证的库，如 cuDNN 和 NCCL，用于启用和优化 GPU 性能。这使得开发者可以轻松地利用 GPU 加速，而无需手动配置和优化环境。