Docker或者Kubernetes中使用PyTorch训练深度学习模型共享内存不足的问题
异常信息
ERROR: Unexpected bus error encountered in worker. This might be caused by insufficient shared memory (shm)
问题原因
在PyTorch中使用DataLoader加载数据集的时候，由于使用多进程加载数据能够提升模型训练的速度。在物理机上面运行没有任务问题，但是在Docker容器或者Kubernetes的Pod中运行就会出现上面的异常情况。
具体原因如下：
PyTorch使用共享内存在进程之间共享数据，因此如果使用torch多进程（例如，对于多进程加载数据的程序），则容器运行时使用的默认共享内存段大小是不够的，默认情况下，Docker容器（或Kubernetes的Pod）共享内存大小为64M，您应该使用--ipc=host或--shm size命令行选项增加共享内存大小，以运行nvidia-docker。
关于DataLoader的说明如下
DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, num_workers=0, collate_fn=default_collate, pin_memory=False, drop_last=False) 复制代码
参数说明
dataset：加载的数据集(Dataset对象)
batch_size：batch size
shuffle：是否将数据打乱
sampler： 样本抽样，后续会详细介绍
num_workers：使用多进程加载的进程数，0代表不使用多进程
collate_fn： 如何将多个样本数据拼接成一个batch，一般使用默认的拼接方式即可
pin_memory：是否将数据保存在pin memory区，pin memory中的数据转到GPU会快一些
drop_last：dataset中的数据个数可能不是batch_size的整数倍，drop_last为True会将多出来不足一个batch的数据丢弃
方案一：加载数据使用单进程
我们可以将num_workers设置为0。这样的确可以解决多进程通信使用共享内存不足的问题，但是这也大大降低了训练的速度。
方案二：修改Docker容器或者Kubernetes的Pod的共享内存
修改Docker容器的shm-size
启动Docker容器时，指定--shm-size。
# 启动docker容器，并进入交互模式 docker run \ --cpus=16 \ --memory=64g \ --gpus '"device=1"' \ --shm-size 8G \ -v /home/junzhi.fan:/junzhi.fan -it harbor.gd.io/test/ocr_:v1.0 \ /bin/bash 复制代码
验证是否生效：
# 再docker容器的交互式命令行查看共享内存 df -h | grep shm # 结果如下： shm 8.0G 0 8.0G 0% /dev/shm 复制代码
修改Kubernetes中Pod的共享内存
使用emptyDir卷来设置共享内存。
apiVersion: v1 kind: Pod metadata: name: test-pd-shm spec: containers: - image: centos name: centos command: [ "sleep", "1000000" ] imagePullPolicy: "IfNotPresent" volumeMounts: - mountPath: /dev/shm name: cache-volume volumes: - emptyDir: medium: Memory sizeLimit: 512Mi name: cache-volume 复制代码
验证是否生效：
# 进入kubernetes集群的pod的交互模式 kubectl exec -it test-pd-shm-cbc944c56-xlbbc /bin/bash # 查看共享内存 df -h # 结果如下： Filesystem Size Used Avail Use% Mounted on overlay 500G 180G 321G 36% / tmpfs 64M 0 64M 0% /dev tmpfs 63G 0 63G 0% /sys/fs/cgroup /dev/sda3 50G 11G 40G 21% /etc/hosts /dev/sda6 500G 180G 321G 36% /etc/hostname shm 512M 0 64M 0% /dev/shm 复制代码
总结
在机器学习训练或需要高效率运行的其他应用场景中，应该根据实际情况调整shm的大小。设置太小，不能够满足高效率的要求，但是，一味地设置过大，容易导致宿主机内存被占用过大，严重时会出现集群雪崩的问题。
因此，在生产环境中，在前期设计的过程中需要好好考虑，建议shm设置为容器分配内存的1/2。
作者：吃果冻不吐果冻皮
链接：https://juejin.cn/post/7031684646114951181