搭建GPU容器环境

概述

测试环境可以在办公网内访问，本文主要介绍在测试环境的 GPU 节点，通过 Docker 来运行一个 PyTorch 环境的 GPU 容器，用于测试和调试。

操作

假设 10.189.109.88 这个节点是经过跳板机之后在的测试集群的 Master 上 ssh 登录上的，总之在这个节点上启动的 Notebook，可以在办公网访问，这个很重要，否则各种操作将会变得十分麻烦。

已知公司对外部网站有做防火墙隔离，所以一些 Nvidia 的镜像是无法在测试/生产环境拉取的，所以这里可以通过已经开通了防火墙策略的节点先下载到本地，再推到公司的 Harbor。

1
2

# docker run --network host --gpus all -it -v /root/pytorch:/pytorch nvcr.io/nvidia/pytorch:21.12-py3
docker run --network host --gpus all -it -v /mnt/vos:/pytorch docker.xxx.com/library/pytorch:21.12-py3

启动很简单，参考上面的命令，并且执行一个 nvidia-smi 浅浅测试一下，读取 GPU 设备确认是没问题的。

通过下面的命令，在容器内启动一个 Notebook，因为容器已经做了 --network=host，所以容器内启动一个 80 端口的进程是自然会映射到宿主机的。

1

jupyter lab --allow-root --ip 0.0.0.0 --NotebookApp.token='TensorRT' --port 80

现在我们在办公网环境下，在浏览器打开 Notebook，输入 TensorRT，就可以看到下面的界面了。

有了这个 Notebook，相信调试代码，就会变得更加方便，当然也可以根据这个镜像，自定义 PyTorch 的 GPU 镜像，这样就可以做更多的事情。

ResNet 50例子

下面在 Notebook 上实测一个程序，必须要理解这些下载的内容，大部分都在公网上，如果有些网站在公司的防火墙策略以外，那么就无法下载，这些在下载一些模型和数据集就会镜像遇到这样的问题，这个问题在测试、生产环境都经常遇到，不要惊讶，大部分的公司也有这样的限制，但是也可以有比较优雅的解决方法的。

很明显，下面的代码在通过 torchvision 来下载一些模型的时候，就因为网络问题无法继续下去了。

很麻烦是不是？怎么在网络策略不通的情况下，可以解决一下这个模型下载的问题呢？这个问题的解决可能可以很简单，看到 Downloading 这个关键词了吗，你可以在办公网，也就是自己电脑先下载了，然后通过 Notebook 上传，再把文件移到 /root/.cache/torch/hub/ 这个目录（虽然听起来很麻烦就是了）。

1
2
3
4
5

Downloading: "https://github.com/pytorch/vision/archive/v0.10.0.zip" to /root/.cache/torch/hub/v0.10.0.zip
---------------------------------------------------------------------------
SSLCertVerificationError                  Traceback (most recent call last)
/opt/conda/lib/python3.8/urllib/request.py in do_open(self, http_class, req, **http_conn_args)
   1353             try:

细看一下下面的操作，也就是做了上面说的事情。

你以为这样就完事了？其实不是，因为 torch.hub 实际存的是各种模型的下载地址，如果要真的下载到本地，你还得能访问这些下载链接。

为了解决这个问题，那么还是得修改一下代码，如下吧，这么写，就能读到刚刚手动上传的模型了。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

import torch
import torchvision
import torchvision.models as models

torch.hub._validate_not_a_forked_repo=lambda a,b,c: True
# 创建 ResNet50 模型实例
resnet50_model = models.resnet50()

# 加载预先训练好的权重
weights_path = '/workspace/examples/torch_tensorrt/notebooks/resnet50-0676ba61.pth'
resnet50_model.load_state_dict(torch.load(weights_path))
# resnet50_model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet50', pretrained=True)

resnet50_model.eval()

更不要说，Notebook 下面提到的这些图片的下载链接了，这种类型的 test 数据，经常会在各种 Notebook 看到。

上面说到了，这些图片的下载，可以自行下载，然后上传到下面的这个目录就可以了，当然也不是什么大问题，除了麻烦以外。

几经一翻挣扎，终于跑起来了。

总结

所以有没有觉得，仅仅是在公司环境内跑一个 Notebook Demo 就很难了？确实如此。有同学可能会说，都下载到本地好了，然后再上传呗。可以认定的是，这么说的同学肯定不是 AI/算法工程师，还是没有接触到相关的痛点，大家需要的，仅仅是跑通一个 Demo 而已，并不想去研究这些所谓的网络策略，但是如果对公司的网络环境比较熟悉的，也可以找到非常 tricky 的方法来实现一些在公网很容易实现的事情。

参考资料

1. Nvidia PyTorch镜像

💡赞赏支持

微信打赏 支付宝打赏