端云协同推理：Split Learning实战

# 更新系统包
sudo apt-get update
sudo apt-get upgrade

# 安装Python和PyTorch
sudo apt-get install python3 python3-pip
pip3 install torch torchvision

# 安装OpenSSL
sudo apt-get install openssl libssl-dev

# 安装CUDA（以CUDA 11.4为例）
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
sudo mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub
sudo add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/ $(lsb_release -cs)"
sudo apt-get update
sudo apt-get -y install cuda-11-4

# 安装Flask
pip3 install flask

# terminal_device.py
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import socket
import ssl
import pickle
import os

# 定义终端设备侧模型（前几层）
class TerminalModel(nn.Module):
    def __init__(self):
        super(TerminalModel, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)  # 展平特征
        return x

# 数据预处理和加载
transform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# 使用CIFAR-10数据集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = DataLoader(trainset, batch_size=32, shuffle=True)

# 初始化终端设备模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
terminal_model = TerminalModel().to(device)
terminal_model.eval()  # 设置为评估模式

# 加载预训练模型权重（假设已经预训练好并保存）
terminal_model.load_state_dict(torch.load("terminal_model.pth"))

# 数据加密函数
def encrypt_data(data, key):
    # 使用AES加密数据
    from Crypto.Cipher import AES
    import base64

    # 将张量转换为字节流
    data_bytes = pickle.dumps(data)
    data_len = len(data_bytes)
    padded_data = data_bytes + (AES.block_size - data_len % AES.block_size) * b'\0'

    # 加密
    iv = os.urandom(AES.block_size)
    cipher = AES.new(key, AES.MODE_CBC, iv)
    encrypted_data = iv + cipher.encrypt(padded_data)
    return base64.b64encode(encrypted_data)

# 通信函数
def send_to_cloud(encrypted_data, cloud_host="127.0.0.1", cloud_port=5000):
    # 创建SSL上下文
    context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
    context.load_verify_locations(cafile="ca.crt")

    # 连接云端服务器
    with socket.create_connection((cloud_host, cloud_port)) as sock:
        with context.wrap_socket(sock, server_hostname=cloud_host) as ssock:
            # 发送加密数据
            ssock.sendall(encrypted_data)

            # 接收加密结果
            encrypted_result = ssock.recv(4096)

    return encrypted_result

# 主流程
if __name__ == "__main__":
    # 生成加密密钥
    key = os.urandom(32)  # 256位AES密钥

    # 遍历数据集，模拟推理请求
    for batch_idx, (inputs, labels) in enumerate(trainloader):
        inputs, labels = inputs.to(device), labels.to(device)

        # 前向推理，提取特征
        features = terminal_model(inputs)

        # 加密特征向量
        encrypted_features = encrypt_data(features, key)

        # 发送到云端进行后向推理
        encrypted_result = send_to_cloud(encrypted_features)

        # 解密结果（这里简化处理，实际应使用对应的解密函数）
        # ...

        # 打印部分结果
        if batch_idx % 10 == 0:
            print(f"Batch {batch_idx}, Encrypted Result Length: {len(encrypted_result)}")

        # 限制处理批次数量（测试用）
        if batch_idx >= 100:
            break

# cloud_server.py
import torch
import torch.nn as nn
import ssl
import socket
import base64
import pickle
from flask import Flask, request, jsonify
from Crypto.Cipher import AES

# 定义云端侧模型（后几层）
class CloudModel(nn.Module):
    def __init__(self, num_classes=10):
        super(CloudModel, self).__init__()
        self.classifier = nn.Sequential(
            nn.Linear(128 * 8 * 8, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(256, num_classes)
        )

    def forward(self, x):
        x = self.classifier(x)
        return x

# 初始化云端模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
cloud_model = CloudModel().to(device)
cloud_model.eval()  # 设置为评估模式

# 加载预训练模型权重（假设已经预训练好并保存）
cloud_model.load_state_dict(torch.load("cloud_model.pth"))

# 解密数据函数
def decrypt_data(encrypted_data, key):
    # 使用AES解密数据
    encrypted_data_bytes = base64.b64decode(encrypted_data)
    iv = encrypted_data_bytes[:AES.block_size]
    encrypted_data_bytes = encrypted_data_bytes[AES.block_size:]

    cipher = AES.new(key, AES.MODE_CBC, iv)
    decrypted_data = cipher.decrypt(encrypted_data_bytes)
    # 去除填充
    decrypted_data = decrypted_data.rstrip(b'\0')
    # 将字节流转换回张量
    data = pickle.loads(decrypted_data)
    return data

# 加密结果函数
def encrypt_result(result, key):
    # 将张量转换为字节流
    result_bytes = pickle.dumps(result)
    data_len = len(result_bytes)
    padded_data = result_bytes + (AES.block_size - data_len % AES.block_size) * b'\0'

    # 加密
    iv = os.urandom(AES.block_size)
    cipher = AES.new(key, AES.MODE_CBC, iv)
    encrypted_result = iv + cipher.encrypt(padded_data)
    return base64.b64encode(encrypted_result)

# Flask应用
app = Flask(__name__)

# 用于存储与终端设备协商的密钥（实际应用中应使用更安全的密钥交换机制）
pre_shared_key = b'0123456789abcdef0123456789abcdef'  # 256位预共享密钥

@app.route('/infer', methods=['POST'])
def infer():
    # 获取加密特征向量
    encrypted_features = request.data

    try:
        # 解密特征向量
        features = decrypt_data(encrypted_features, pre_shared_key)

        # 后向推理
        features = torch.tensor(features).to(device)
        outputs = cloud_model(features)
        _, predicted = torch.max(outputs, 1)

        # 加密结果
        encrypted_result = encrypt_result(predicted.cpu().numpy(), pre_shared_key)

        return encrypted_result

    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == "__main__":
    # 启动Flask服务
    app.run(host='0.0.0.0', port=5000, ssl_context=('cert.pem', 'key.pem'))

# distributed_training.py
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
from terminal_device import TerminalModel
from cloud_server import CloudModel
import torchvision
import torchvision.transforms as transforms
import socket
import ssl
import pickle
import os
import numpy as np

# 初始化分布式训练环境
def init_distributed():
    # 设置分布式训练参数
    rank = int(os.environ['RANK'])
    world_size = int(os.environ['WORLD_SIZE'])
    master_addr = os.environ['MASTER_ADDR']
    master_port = os.environ['MASTER_PORT']

    # 初始化分布式进程组
    dist.init_process_group(
        backend='gloo',
        init_method=f'tcp://{master_addr}:{master_port}',
        rank=rank,
        world_size=world_size
    )

    return rank, world_size

# 数据预处理和加载
transform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# 使用CIFAR-10数据集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=32, shuffle=True)

# 定义联合模型（用于本地训练评估）
class CombinedModel(nn.Module):
    def __init__(self, terminal_model, cloud_model):
        super(CombinedModel, self).__init__()
        self.terminal_model = terminal_model
        self.cloud_model = cloud_model

    def forward(self, x):
        x = self.terminal_model(x)
        x = self.cloud_model(x)
        return x

# 训练函数
def train(rank, world_size, num_epochs=5):
    # 初始化分布式环境
    init_distributed()

    # 初始化终端设备和云端模型
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    terminal_model = TerminalModel().to(device)
    cloud_model = CloudModel().to(device)

    # 合并模型用于本地训练评估
    combined_model = CombinedModel(terminal_model, cloud_model).to(device)

    # 损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(combined_model.parameters(), lr=0.001)

    # 分布式数据采样器
    train_sampler = torch.utils.data.distributed.DistributedSampler(
        trainset,
        num_replicas=world_size,
        rank=rank
    )

    # 分布式数据加载器
    trainloader = torch.utils.data.DataLoader(
        trainset,
        batch_size=32,
        shuffle=False,
        sampler=train_sampler
    )

    # 训练循环
    for epoch in range(num_epochs):
        train_sampler.set_epoch(epoch)
        combined_model.train()

        running_loss = 0.0
        total = 0

        for batch_idx, (inputs, labels) in enumerate(trainloader):
            inputs, labels = inputs.to(device), labels.to(device)

            optimizer.zero_grad()

            # 前向传播
            outputs = combined_model(inputs)
            loss = criterion(outputs, labels)

            # 反向传播和优化
            loss.backward()
            optimizer.step()

            running_loss += loss.item() * inputs.size(0)
            total += inputs.size(0)

            if batch_idx % 10 == 0 and rank == 0:
                print(f"Epoch [{epoch + 1}/{num_epochs}], Batch [{batch_idx + 1}/{len(trainloader)}], Loss: {loss.item():.4f}")

        # 计算平均损失
        avg_loss = running_loss / total
        avg_loss_tensor = torch.tensor(avg_loss).to(device)

        # 在所有进程间同步平均损失
        dist.all_reduce(avg_loss_tensor, op=dist.ReduceOp.SUM)
        avg_loss = avg_loss_tensor.item() / world_size

        if rank == 0:
            print(f"Epoch [{epoch + 1}/{num_epochs}], Average Loss: {avg_loss:.4f}")

    # 保存模型
    if rank == 0:
        torch.save(terminal_model.state_dict(), "terminal_model.pth")
        torch.save(cloud_model.state_dict(), "cloud_model.pth")

    # 清理分布式环境
    dist.destroy_process_group()

if __name__ == "__main__":
    # 设置环境变量（实际应用中应通过其他方式传递）
    os.environ['RANK'] = '0'
    os.environ['WORLD_SIZE'] = '2'
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '29500'

    # 启动训练
    train(rank=0, world_size=2)