报错

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