PyTorch 深度学习实战（33）：联邦学习与隐私保护

在上一篇文章中，我们探讨了多模态学习与CLIP模型的应用。本文将深入介绍联邦学习（Federated Learning）这一新兴的分布式机器学习范式，它能够在保护数据隐私的前提下实现多方协作的模型训练。我们将使用PyTorch实现一个基础的联邦学习框架，并在图像分类任务上进行验证。

---

一、联邦学习基础

联邦学习是一种分布式机器学习方法，其核心思想是数据不动，模型动——参与方的数据保留在本地，仅通过交换模型参数或梯度来实现协同训练。

1. 联邦学习的核心组件

• 中心服务器（Coordinator）：负责协调训练过程，聚合各客户端模型

• 客户端（Client）：拥有本地数据，执行本地训练

• 通信协议：定义参数交换格式和加密方式

2. 联邦学习的数学表达

典型的联邦学习优化目标可以表示为：

3. 联邦学习的优势

优势	说明
隐私保护	原始数据始终保留在本地
数据多样性	利用分布在不同设备上的异构数据
降低通信成本	仅传输模型参数而非原始数据
合规性	满足GDPR等数据保护法规

4. 联邦学习的类型

1. 横向联邦学习：客户端拥有相同特征空间的不同样本

2. 纵向联邦学习：客户端拥有相同样本的不同特征

3. 联邦迁移学习：客户端间数据和特征空间都不同

---

二、联邦学习实战：图像分类

我们将实现一个基于CIFAR-10数据集的横向联邦学习系统，模拟5个客户端协作训练图像分类模型。

1. 环境配置

首先安装必要库：

pip install torch torchvision cryptography

2. 基础实现

2.1 数据分区与客户端模拟

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Subset
import numpy as np
from collections import OrderedDict
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
import os
import base64
​
# 设置随机种子
torch.manual_seed(42)
np.random.seed(42)
​
# 设备配置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")
​
# 数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
​
# 加载完整数据集
full_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
​
# 客户端数量
NUM_CLIENTS = 5
​
# 非IID数据划分（每个客户端只获取2类数据）
def non_iid_split(dataset, num_clients):
    class_indices = {i: [] for i in range(10)}
    for idx, (_, label) in enumerate(dataset)