一、引言

        近年来随着大模型的大放异彩，AI被推上了新的高度，学习AI是每个人的必修课。AI模型的训练，需要用到很多的计算资源，单机往往不能满足训练的需求。为满足提高训练效率，加快模型迭代的需求，多机多卡训练成为自然而然的选择。如何管理显卡资源，如何调度训练任务，实现多机多卡的集群。本文将详细介绍如何在Kubernetes环境下，基于Python搭建一个高效的分布式深度学习训练平台。

二、环境准备

• 服务器准备：1 master 1 worker
• 操作系统：Centos 7
• Kubernetes：1.18.2
• Docker：19.03.1
• Python：3.9
• 深度学习框架：pytorch 

三、K8S集群搭建

192.168.5.101 k8s-master 

192.168.5.102 k8s-worker01 

Docker安装

• 修改主机名

192.168.5.101机器   hostnamectl set-hostname k8s-master

192.168.5.102机器   hostnamectl set-hostname k8s-worker01 

• 配置本地域名解析

两个机器上都执行

cat >> /etc/hosts << EOF
192.168.5.101 k8s-master
192.168.5.102 k8s-worker01
EOF

• 关闭防火墙并禁止其开机自启动

systemctl stop firewalld & systemctl disable firewalld

• 关闭selinux(关闭selinux以允许容器访问宿主机的文件系统)

sed -i 's/^SELINUX=enforcing$/SELINUX=disabled/' /etc/selinux/config && setenforce 0

• 关闭swap分区(影响k8s内存性能)

sed -i '/ swap / s/^/#/' /etc/fstab

swapoff -a

• 配置yum源

wget -O /etc/yum.repos.d/CentOS-Base.repo http://mirrors.aliyun.com/repo/Centos-7.repo

yum -y install yum-utils

yum-config-manager --add-repo http://mirrors.aliyun.com/docker-ce/linux/centos/docker-ce.repo 

• 安装docker

yum install docker-ce-19.03.1 -y

• 启动一下我们的docker容器并设置为开机自启动

systemctl start docker & systemctl enable docker

执行命令docker version，检查docker是否安装完成，如下则标准成功。

K8S安装

• 添加kubernetes源，这里配置为阿里云的镜像源

#配置kubernetes源

cat <<EOF > /etc/yum.repos.d/kubernetes.repo
[kubernetes]
name=Kubernetes
baseurl=https://mirrors.aliyun.com/kubernetes/yum/repos/kubernetes-el7-x86_64
enabled=1
gpgcheck=0
repo_gpgcheck=0
gpgkey=https://mirrors.aliyun.com/kubernetes/yum/doc/yum-key.gpg \
https://mirrors.aliyun.com/kubernetes/yum/doc/rpm-package-key.gpg
EOF

• 安装kubernetes组件

yum install -y kubelet-1.18.2 kubeadm-1.18.2 kubectl-1.18.2

• 启动kubelet

 systemctl enable kubelet && systemctl start kubelet

• 配置master节点，本步只在master节点执行

kubeadm init --pod-network-cidr=192.168.0.0/16 --kubernetes-version=v1.18.2 --apiserver-advertise-address=192.168.5.101 --image-repository=registry.aliyuncs.com/google_containers

控制台会输出提示，需要建目录，执行如下命令：

mkdir -p $HOME/.kube
sudo cp -i /etc/kubernetes/admin.conf $HOME/.kube/config
sudo chown $(id -u):$(id -g) $HOME/.kube/config

完成该步后，会在控制台打印输出，记住它，拷贝它：

kubeadm join 192.168.23.101:6443 --token 6rc83u.nh3m19pypf024dw7     --discovery-token-ca-cert-hash sha256:cf73c46a33f6abae1bfc9eb6ef29050ddaadc3c7be880ec8cbe26fb4b2c42aaa

• 安装网络插件

kubectl apply -f https://docs.projectcalico.org/v3.10/manifests/calico.yaml

• 检查主节点是否已经在运行

kubectl get nodes

kubectl -n kube-system get pods

若status不是ready状态，可以等待一下，初始化需要时间。

至此，我们的master节点已经配置完成，后面就是怎么把其他的work节点加入到我们集群中。

• 加入work节点

在k8s-worker01上执行上面我们保存的命令

kubeadm join 192.168.5.101:6443 --token 6rc83u.nh3m19pypf024dw7     --discovery-token-ca-cert-hash sha256:cf73c46a33f6abae1bfc9eb6ef29050ddaadc3c7be880ec8cbe26fb4b2c42aaa

• 查看集群状态

kubectl get nodes

可以给work01节点设置角色为worker：

kubectl label nodes k8s-worker01 node-role.kubernetes.io/worker=

查看pod的状态

kubectl -n kube-system get pods -o wide

等等一会后，所有的都ready，择表示机集群部署成功了。

若安装遇到问题，我们有时候需要移除节点，移除节点的方法如下：

#在主节点执行以下命令，这里以work01为例

kubectl drain k8s-worker01 --delete-local-data --force --ignore-daemonsets

kubectl delete node k8s-worker01

#在k8s-worker01节点上执行以下命令，重置节点

kubeadm reset -f

四、训练模型

训练镜像构建

Dockerfile文件内容

FROM python:3.8

RUN pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

WORKDIR /

COPY demo.py demo.py

demo.py

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torchvision import datasets, transforms
from model import SimpleCNN
 
transform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.ToTensor(),
])
train_dataset = datasets.CIFAR10(root='/data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)

loss_values = []

def train(model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 10 == 0:
            loss_values.append(loss.item())
            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} ({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')
 
def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = SimpleCNN().to(device)
    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
    
    for epoch in range(1, 11):  # 训练10个epoch
        train(model, device, train_loader, optimizer, epoch)
    
    torch.save(model.state_dict(), "model.pth")
 
if __name__ == "__main__":
    main()

model.py 模型定义

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torchvision import datasets, transforms
 
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.fc1 = nn.Linear(2304, 128)
        self.fc2 = nn.Linear(128, 10)
 
    def forward(self, x):
        x = self.conv1(x)
        x = nn.functional.relu(x)
        x = nn.functional.max_pool2d(x, 2)
        x = self.conv2(x)
        x = nn.functional.relu(x)
        x = nn.functional.max_pool2d(x, 2)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = nn.functional.relu(x)
        x = self.fc2(x)
        return nn.functional.log_softmax(x, dim=1)
 

构建命令

docker build -t my_train_test_image:v0.1 .

如此，则构建好训练镜像。

创建k8s训练任务

train_job.yaml

apiVersion: batch/v1
kind: Job
metadata:
  name: pytorch-job
spec:
  template:
    spec:
      containers:
      - name: pytorch-container
        image: my_train_test_image:v0.1
        command: ["python3", "/demo.py"]
        args: ["--epochs", "10"]
        volumeMounts:
        - name: pytorch-data
          mountPath: /data  
      restartPolicy: OnFailure
      volumes:
      - hostPath:
          path: /data
          type: DirectoryOrCreate
        name: pytorch-data

提交任务到集群

kubectl apply -f train_job.yaml -n kube-system

通过kubectl get pods -A命令，查看job容器是否完成，在训练中，可已通过kubectl logs -f poxxxx 命令，查看训练的日志输出。

损失可视化

可视化训练中的损失值，可以看到训练的收敛情况

# 可视化损失值
plt.figure(figsize=(10, 5))
plt.plot(loss_values, label='Training Loss')
plt.xlabel('Batch (x100)')
plt.ylabel('Loss')
plt.title('Training Loss Over Time')
plt.legend()
plt.savefig('training_loss.png')  # Save the loss plot
plt.show()

可以看到，训练损失值趋于稳定，模型在逐渐收敛的。

def imshow(img, filename=None):
    plt.figure()  # 为图像创建一个新的图形
    img = img / 2 + 0.5  # 反归一化
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    if filename:
        plt.savefig(filename)  # 保存图像
    plt.show()

# 获取一些随机测试图像
dataiter = iter(testloader)
images, labels = next(dataiter)
imshow(torchvision.utils.make_grid(images), 'test_images.png')
outputs = model(images)
_, predicted = torch.max(outputs, 1)

print('Predicted: ', ' '.join(f'{trainset.classes[predicted[j]]}' for j in range(4)))

下图为部分测试图

下表为每个类别的评估指标的整体评估结果

类别	精确率 (Precision)	召回率 (Recall)	F1 分数 (F1 Score)
airplane	0.72	0.74	0.73
automobile	0.86	0.81	0.84
bird	0.59	0.63	0.61
cat	0.58	0.47	0.52
deer	0.62	0.67	0.65
dog	0.61	0.60	0.61
frog	0.69	0.83	0.76
horse	0.84	0.71	0.77
ship	0.83	0.81	0.82
truck	0.79	0.81	0.80

模型保存

完成后，则会保存模型到/model.pth文件中，供后续的推理脚本或者推理服务使用。

五、模型推理

这一步，加载第四步里面的模型文件，来做推理。

test.py

import torch
import torch.nn.functional as F
from torch.autograd import Variable
from model import SimpleCNN
import numpy as np
import cv2
 
def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = SimpleCNN()
    model.load_state_dict(torch.load("model.pth"))
    model = model.to(device)
    model.eval()

    image = cv2.imread('010029.jpg')
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = cv2.resize(image, (32, 32))
    image = image.astype(np.float32) / 255.0
    image = (image - 0.5) / 0.5
    image = image.transpose(2, 0, 1)
    image_tensor = torch.from_numpy(image).unsqueeze(0)

    classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

    with torch.no_grad():
        output = model(image_tensor)
        pred_class = output.argmax(dim=1).item()
        probability = torch.exp(output[0, pred_class]).item()

    print(f"label_id: {pred_class},label: {classes[pred_class]}, score: {probability:.4f}")

 
if __name__ == "__main__":
    main()

docker run -it --name zzztest111 my_train_test_image:v0.1 python3 /test.py

输出

label_id: 0,label: plane, score: 1.0000

类别序号为0，物体为plane，置信度分值为1.0000

六、总结

        本篇文章我们从0到1实战了从部署k8s集群，到训练一个CV算法模型的小案例，不仅掌握了构建集群和训练CNN模型的基本流程，还学习了如何评估使用模型进行推理，帮助将来使用的时候得心应手。K8S工具功能强大，如何进行计算资源分配，限制，调度等延伸能力，感兴趣的同学可以继续往下深究，如何训练模型，使用模型这些技能都是CV算法工程师的基础能力，若感兴趣可以在案例的基础上进行延伸扩展，试试模型优化参数调优。是不是比较简单，赶紧动手学习起来吧。