《联邦学习实战—杨强》之使用Python从零开始实现一个简单的横向联邦学习模型
文章目录
- 前言
- 环境准备
- 完整代码
- 配置文件(conf.json)
- 获取数据集(datasets.py)
- 获取PyTorch中自带深度学习网络预训练模型(models.py)
- 客户端(client.py)
- 服务端(server.py)
- main.py
- 运行
- 知识点补充
-
- argparse基本用法
- tensor.copy_()
- [Python dict() 函数](https://www.runoob.com/python/python-func-dict.html)
- 什么是状态字典：state_dict?
- view_as()函数
前言
本文通过阅读《联邦学习实战—杨强》中第3章“用Python实现横向联邦图像分类”入门横向联邦。
核心思想: 使用Python在本地模拟多个客户端，然后由服务器统一管理进行联邦学习，客户端在本地用自己的数据对模型进行训练，服务器将训练结果聚合更新模型并分发给客户端，客户端继续训练。
用横向联邦学习来实现对cifar10图像数据集的分类，模型使用的是ResNet-18。方式为在本地以循环的方式来模拟。
多多调试来理解代码
环境准备
本实验基于Python实现，使用机器学习库PyTorch。
- Python、PyTorch
- 编译器使用PyCharm
- 数据集：cifar10
- 模型：ResNet-18
基本流程：
- 服务器按配置文件生成初始化模型，客户端按照自己的ID将数据集横向不重叠切割
- 服务器将全局模型发送给客户端
- 客户端接收全局模型（来自服务器）并通过本地多次迭代计算本地参数差值返回给服务器
- 服务器聚合各个客户端差值更新模型，再评估当前模型性能
- 如果性能未达标，则重复2过程，否则结束
CIFAR10数据集:
每个批处理文件都包含一个包含以下元素的字典，加起来有6万张图片：
- 数据 —— 60000x3072(32x32x3)的uint8 numpy数组。数组的每一行存储一个 32x32 彩色图像。前 1024 个条目包含红色通道值，接下来的 1024 个条目包含绿色通道值，最后 1024 个条目包含蓝色通道值。图像以行优先顺序存储，因此数组的前 32 个条目是图像第一行的红色通道值。
- labels —— 0-9 范围内的 60000 个数字的列表。索引i处的数字表示数组数据中第i个图像的标签。
完整代码
链接: https://pan.baidu.com/s/1PPCd_YB6w537OYif9TQZPg?pwd=s987 提取码: s987
配置文件(conf.json)
联邦训练配置: 一共10台客户端设备(no_models=10)，每一轮任意挑选其中的5台参与训练(k=5)，每一次本地训练迭代次数为3次(local_epochs=3)，全局迭代次数为20次(global_epochs=20)
- 训练的客户端数量(no_models): 每一轮的选代，服务端会首先从所有的客户端中挑选部分客户端进行本地训练。每一次迭代只选取部分客户端参与，并不会影响全局收敛的效果，且能够提升训练的效率。
- 全局迭代次数(global_epochs): 即服务端和客户端的通信次数。通常会设置一个最大的全局迭代次数，但在训练过程中，只要模型满足收敛的条件，那么训练也可以提前终止。
- 本地模型的选代次数(local_epochs)：即每一个客户端在进行本地模型训练时的迭代次数。每一个客户端的本地模型的迭代次数可以相同，也可以不同。
- 本地训练相关的算法配置：本地模型进行训练时的参数设置，如学习率(lr)、训练样本大小、使用的优化算法等。
- 模型信息：即当前任务我们使用的模型结构。在本案例中，我们使用 ResNet-18 图像分类模型。
- 数据信息：联邦学习训练的数据。在本案例中，我们将使用 cifar10 数据集。为了模拟横向建模，数据集将按样本维度，切分为多份不重叠的数据(本文将5万张训练集图片数据切分为10份，分给10个客户端)，每一份(5000张图片)放置在每一个客户端中作为本地训练数据。
- 其他的配置信息，比如可能使用到的加密方案、是否使用差分隐私、模型是否需要检查点文件（checkpoint ）、模型聚合的策略等，都可以根据实际需要自行添加或者修改。
我们将上面的信息以json 格式记录在配置文件中以便修改，如下所示。
conf.json代码:
{
"model_name" : "resnet18",
"no_models" : 10,
"type" : "cifar",
"global_epochs" : 20,
"local_epochs" : 3,
"k" : 5,
"batch_size" : 32,
"lr" : 0.001,
"momentum" : 0.0001,
"lambda" : 0.1
}
获取数据集(datasets.py)
import torch
from torchvision import datasets, transforms
# 获取数据集
def get_dataset(dir, name):
if name=='mnist':
# root(这里是dir): 数据路径
# train参数表示是否是训练集(True)或者测试集(False)
# download=true表示从互联网上下载数据集并把数据集放在root路径中
# transform：图像类型的转换
# 将图片转化为张量对象。通过ToTensor实例将图像数据从PIL类型变换成32位浮点数格式，并除以255使得所有像素的数值均在 0 到 1 之间
train_dataset = datasets.MNIST(dir, train=True, download=True, transform=transforms.ToTensor())
eval_dataset = datasets.MNIST(dir, train=False, transform=transforms.ToTensor())
elif name=='cifar':
# transforms.Compose(图片预处理步骤)是将多个transform组合起来使用（由transform构成的列表）
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4), # 从图片中随机裁剪出尺寸为size的图片。size: 所需裁剪出的图片尺寸。padding: 设置填充大小，当为a时，上下左右均填充a个像素。
transforms.RandomHorizontalFlip(), # 水平翻转(左右翻转图像通常不会改变对象的类别。这是最早且最广泛使用的图像增广方法)
transforms.ToTensor(), # 将图片转化为张量对象。通过ToTensor实例将图像数据从PIL类型变换成32位浮点数格式，并除以255使得所有像素的数值均在 0 到 1 之间
# transforms.Normalize: 标准化图像的每个通道。第一个参数: 均值; 第二个参数: 方差
# 对于RGB(红、绿、蓝)颜色通道，我们分别标准化每个通道。具体而言，该通道的每个值减去该通道的平均值，然后将结果除以该通道的标准差。
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
train_dataset = datasets.CIFAR10(dir, train=True, download=True,
transform=transform_train)
eval_dataset = datasets.CIFAR10(dir, train=False, transform=transform_test)
return train_dataset, eval_dataset
获取PyTorch中自带深度学习网络预训练模型(models.py)
import torch
from torchvision import models
def get_model(name="vgg16", pretrained=True):
# 卷积神经网络的训练是耗时的，很多场合不可能每次都从随机初始化参数开始训练网络。
# pytorch中自带几种常用的深度学习网络预训练模型，如VGG、ResNet等。往往为了加快学习的进度，在训练的初期我们直接加载pre-train模型中预先训练好的参数
if name == "resnet18":
# 使用在ImageNet数据集上预训练的ResNet-18作为源模型。指定pretrained=True自动下载预训练的模型参数
model = models.resnet18(pretrained=pretrained)
elif name == "resnet50":
# 使用在ImageNet数据集上预训练的ResNet-50作为源模型。指定pretrained=True自动下载预训练的模型参数
model = models.resnet50(pretrained=pretrained)
elif name == "densenet121":
model = models.densenet121(pretrained=pretrained)
elif name == "alexnet":
model = models.alexnet(pretrained=pretrained)
elif name == "vgg16":
model = models.vgg16(pretrained=pretrained)
elif name == "vgg19":
model = models.vgg19(pretrained=pretrained)
elif name == "inception_v3":
model = models.inception_v3(pretrained=pretrained)
elif name == "googlenet":
model = models.googlenet(pretrained=pretrained)
if torch.cuda.is_available():
return model.cuda()
else:
return model
客户端(client.py)
横向联邦学习的客户端的主要功能是接收服务端的下发指令和全局模型(参数)，并利用本地数据进行局部模型训练
# 在项目文件夹下创建client.py文件，客户端的主要功能是接受服务器传来的全局模型，并利用本地数据对模型进行训练后返回"差值"，包括构造函数、本地训练函数。
import models, torch, copy
# 客户端类
class Client(object): # clients.append(Client(conf, server.global_model, train_datasets, c))
# 构造函数
def __init__(self, conf, model, train_dataset, id = -1):
# 读取配置文件
self.conf = conf
# 根据配置文件获取客户端本地模型（一般由服务器传输）
self.local_model = models.get_model(self.conf["model_name"])
# 客户端ID
self.client_id = id
# 客户端本地数据集
self.train_dataset = train_dataset
# 按ID对数据集集合进行拆分
all_range = list(range(len(self.train_dataset))) # all_range列表: 0~49999
data_len = int(len(self.train_dataset) / self.conf['no_models']) # 50000/总客户端数量=data_len:5000
train_indices = all_range[id * data_len: (id + 1) * data_len] # 将切分数据集，每一份有5000张图片
# 若id为0: [0:5000]; 若id为1: [5000, 10000]; 若id为2: [10000,15000],......; 若id为9: [45000, 50000]
self.train_loader = torch.utils.data.DataLoader(self.train_dataset, batch_size=conf["batch_size"],
sampler=torch.utils.data.sampler.SubsetRandomSampler(train_indices))
# SubsetRandomSampler用来打乱数据
# 模型本地训练函数
def local_train(self, model):
# 客户端获取服务器的模型，然后通过部分本地数据集进行训练
for name, param in model.state_dict().items():
# 用服务器下发的全局模型参数(weight、bias)覆盖本地模型参数(weight、bias)。
# 本文服务端、客户端采用的都是ResNet-18模型，所以其实参数时一样的，但是按照代码完整性，还是覆盖一下。
self.local_model.state_dict()[name].copy_(param.clone())
#print(id(model))
# 定义最优化函数器用户本地模型训练
optimizer = torch.optim.SGD(self.local_model.parameters(), lr=self.conf['lr'],
momentum=self.conf['momentum'])
#print(id(self.local_model))
# 本地训练模型
# 设置开启模型训练，如果模型中有BN层(Batch Normalization）和Dropout，需要在训练时添加local_model.train()，在测试时添加local_model.eval()
self.local_model.train()
# 开始训练模型，本文每个客户端都是本地训练3次
for e in range(self.conf["local_epochs"]):
# 每次(共3次)本地训练训练5000张图片，每次下面的循环需要循环157次
for batch_id, batch in enumerate(self.train_loader): # self.train_loader有5000张图片，分成了5000/32=157份
data, target = batch
if torch.cuda.is_available(): # 如果可以的话加载到gpu
data = data.cuda()
target = target.cuda()
optimizer.zero_grad() # 在计算参数的梯度之前，通常需要清零梯度信息。清零模型参数的梯度值
output = self.local_model(data) # # 训练预测
loss = torch.nn.functional.cross_entropy(output, target) # 计算损失函数cross_entropy交叉熵误差
loss.backward() # 利用自动求导函数loss.backward()求出模型中所有参数的梯度信息"loss对权重或偏置求导"，这些梯度会自动保存在每个张量的grad成员变量中
optimizer.step() # 根据梯度下降算法更新参数。w'=w-lr*grad; b'=b-lr*grad
print(f'客户端: {self.client_id} 的', "Epoch %d done." % e)
# 创建差值字典（结构与模型参数同规格），用于记录差值
diff = dict()
# 此时self.local_model.state_dict()和model.state_dict()不一样了
for name, data in self.local_model.state_dict().items(): # ResNet-18有122个层需要更新参数，所以这里执行122次循环(通过调试理解)
# 计算训练后与训练前的差值
diff[name] = (data - model.state_dict()[name])
#print(diff[name])
print("客户端 %d 本地训练结束" % self.client_id)
#客户端返回差值
return diff
服务端(server.py)
横向联邦学习的服务端的主要功能是将被选择的客户端上传的本地模型进行模型聚合(使用FedAvg算法)
import models, torch
# 服务器类
class Server(object):
def __init__(self, conf, eval_dataset): # 定义构造函数
self.conf = conf # 将配置信息拷贝到服务端中
self.global_model = models.get_model(self.conf["model_name"]) # 按照配置中的模型信息获取模型，返回类型为model(nn.Module)
# self.eval_loader有 313 份测试集batch(32张图片为一份batch)
self.eval_loader = torch.utils.data.DataLoader(eval_dataset, batch_size=self.conf["batch_size"], shuffle=True)
# 模型聚合函数
# weight_accumulator 存储了每个客户端上传参数的变化值
def model_aggregate(self, weight_accumulator):
# 遍历服务器的全局模型
for name, data in self.global_model.state_dict().items(): # ResNet-18有122个层需要更新参数，所以这里执行122次循环(通过调试理解)
# 书第44页FedAvg算法公式的右半部分
update_per_layer = weight_accumulator[name] * self.conf["lambda"]
# 累加
if data.type() != update_per_layer.type():
# 因为update_per_layer的type是floatTensor，所以将其转换为模型的LongTensor（损失精度）
# 这里就是整个书44页的FedAvg算法公式，结果为聚合之后的全局模型的weight、bias
data.add_(update_per_layer.to(torch.int64))
else:
data.add_(update_per_layer)
# 模型评估函数
def model_eval(self):
# 开启模型评估模式
self.global_model.eval()
total_loss = 0.0
correct = 0
dataset_size = 0
# 遍历评估数据集合
for batch_id, batch in enumerate(self.eval_loader): # self.eval_loader有 313 份测试集batch(32张图片为一份batch)
data, target = batch
# 获取所有样本总量大小
dataset_size += data.size()[0] # data.size()=torch.Size([32, 3, 32, 32])
if torch.cuda.is_available(): # 如果可以的话存储到gpu
data = data.cuda()
target = target.cuda()
# 加载到模型中训练
output = self.global_model(data) # output的shape为(32, 1000)
# 聚合所有损失 cross_entropy 交叉熵函数计算损失
total_loss += torch.nn.functional.cross_entropy(output, target,
reduction='sum').item() # sum up batch loss
# 获取最大的对数概率的索引值，即在所有预测结果中选择可能性最大的作为最终结果
# 解读output.data.max(1): 首先output的shape是(32, 1000)。通过output的第二个维度找到每行最大值(共32个)，返回是值和索引(通过调试理解)
# 解读output.data.max(1)[1]: 提取出"索引"
pred = output.data.max(1)[1] # get the index of the max log-probability
# 统计预测结果与真实标签的匹配个数
correct += pred.eq(target.data.view_as(pred)).cpu().sum().item()
# 计算准确率
acc = 100.0 * (float(correct) / float(dataset_size))
# 计算总损失值
total_l = total_loss / dataset_size
return acc, total_l
server.py涉及的调试截图:
- output.data.max(1)
解读output.data.max(1): 首先output的shape是(32, 1000)。通过output的第二个维度找到每行最大值(共32个)，返回是值和索引(通过调试理解)
- 解读output.data.max(1)[1]: 提取出"索引"
验证一下:
main.py
import argparse, json
import datetime
import os
import logging
import torch, random
from server import *
from client import *
import models, datasets
if __name__ == '__main__':
# 设置命令行程序
parser = argparse.ArgumentParser(description='Federated Learning')
parser.add_argument('-c', '--conf', dest='conf') # argparse默认的变量名是--或-后面的字符串，
# 但是你也可以通过dest=xxx来设置参数的变量名，然后在代码中用args.xxx来获取参数的值。
args = parser.parse_args()
# 终端执行python main.py -c ./utils/conf.json。此时，args.conf = ./utils/conf.json
# 这里我使用文件路径地址替换了"arg.conf"，可以直接运行main.py
with open('./utils/conf.json', 'r') as f: # args.conf
conf = json.load(f)
# conf = {"model_name" : "resnet18","no_models" : 10,"type" : "cifar",
# "global_epochs" : 20,"local_epochs" : 3,"k" : 5,"batch_size" : 32,
# "lr" : 0.001,"momentum" : 0.0001,"lambda" : 0.1}
# eval_datasets是"字典"，键为"数据"——一个 10000x3072(32*32*3=3072)的uint8 numpy数组，值为"labels" -- 0-9 范围内的 10000 个数字的列表。
# conf中的 "type" : "cifar"
train_datasets, eval_datasets = datasets.get_dataset("./data/", conf["type"])
server = Server(conf, eval_datasets) # Server类实例化。self.eval_loader有 10000/32=313 份测试集batch(32张图片为一份batch)
clients = [] # 定义客户端列表
# clients列表里面存放10个客户端类的实例
for c in range(conf["no_models"]): # c = 0~9。"c"是客户端的id号
clients.append(Client(conf, server.global_model, train_datasets, c))
print("\n\n")
# 全局训练
for e in range(conf["global_epochs"]): # e = 0~19
print(f"当前是第 {e} 次大循环Global Epoch")
# 每次训练从clients列表中随机抽取k个进行训练。candidates也是列表，里面有5个客户端实例
candidates = random.sample(clients, conf["k"])
print("selected clients are: ")
for c in candidates:
print('client_id: ', c.client_id)
# 累加参数变化。把每一个大循环(global_epochs)的差值加起来
weight_accumulator = {}
# 初始化空模型参数weight_accumulator
for name, params in server.global_model.state_dict().items():
# 这里务必调试理解！！！
# name分别为: 'conv1.weight'、'bn1.weight'、'bn1.bias'等等(调试看看);
# params分别为对应各个name的参数值
weight_accumulator[name] = torch.zeros_like(params) # 生成一个和参数矩阵大小相同的零矩阵
# 遍历选中的客户端，每个客户端本地进行训练
for c in candidates:
diff = c.local_train(server.global_model)
# 根据客户端返回的参数差值字典更新总体权重
for name, params in server.global_model.state_dict().items(): # ResNet-18有122个层需要更新参数，所以这里执行122次循环(通过调试理解)
weight_accumulator[name].add_(diff[name])
# 模型参数聚合
server.model_aggregate(weight_accumulator) # 执行完这行代码后，模型全局参数就更新了
# 模型评估
acc, loss = server.model_eval()
print("Epoch %d, acc: %f, loss: %f\n" % (e, acc, loss))
main.py涉及的调试截图:
-
初始化candidates列表(从clients列表里随机选择5个)
-
server.global_model.state_dict().items():
-
初始化字典weight_accumulator
weight_accumulator[name] = torch.zeros_like(params) # 生成一个和参数矩阵大小相同的零矩阵
运行
在PyCharm里直接运行main.py便可
知识点补充
argparse基本用法
- 文章 1
- 文章 2
tensor.copy_()
tensor.copy_(src)
将src中的元素复制到tensor中并返回这个tensor; 两个tensor应该有相同shape
例子:
x = torch.tensor([[1,2], [3,4], [5,6]])
y = torch.rand((3,2))
print(y)
y.copy_(x)
print(y)
输出:
tensor([[0.1604, 0.0176],
[0.3737, 0.2009],
[0.1438, 0.8394]])
tensor([[1., 2.],
[3., 4.],
[5., 6.]])
[Finished in 1.9s]
Python dict() 函数
什么是状态字典：state_dict?
在PyTorch中，torch.nn.Module模型的可学习参数（即权重和偏差）包含在模型的参数中，（使用model.parameters()可以进行访问）。 state_dict是Python字典对象，它将每一层映射到其参数张量。注意，只有具有可学习参数的层（如卷积层，线性层等）的模型才具有state_dict这一项。目标优化torch.optim也有state_dict属性，它包含有关优化器的状态信息，以及使用的超参数。
因为state_dict的对象是Python字典，所以它们可以很容易的保存、更新、修改和恢复，为PyTorch模型和优化器添加了大量模块。
下面通过从简单模型训练一个分类器中来了解一下state_dict的使用。
# 定义模型
class TheModelClass(nn.Module):
def __init__(self):
super(TheModelClass, self).__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 16 * 5 * 5)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
# 初始化模型
model = TheModelClass()
# 初始化优化器
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# 打印模型的状态字典
print("Model's state_dict:")
for param_tensor in model.state_dict():
print(param_tensor, "\t", model.state_dict()[param_tensor].size())
# 打印优化器的状态字典
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
print(var_name, "\t", optimizer.state_dict()[var_name])
输出:
Model's state_dict:
conv1.weight torch.Size([6, 3, 5, 5])
conv1.bias torch.Size([6])
conv2.weight torch.Size([16, 6, 5, 5])
conv2.bias torch.Size([16])
fc1.weight torch.Size([120, 400])
fc1.bias torch.Size([120])
fc2.weight torch.Size([84, 120])
fc2.bias torch.Size([84])
fc3.weight torch.Size([10, 84])
fc3.bias torch.Size([10])
Optimizer's state_dict:
state {}
param_groups [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [4675713712, 4675713784, 4675714000, 4675714072, 4675714216, 4675714288, 4675714432, 4675714504, 4675714648, 4675714720]}]
view_as()函数
函数定义：view_as(tensor) [参数为一个Tensor张量]
该函数的作用是将调用函数的变量，转变为同参数tensor同样的形状。
例如：
data1 = [[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 0], [10, 11]]]
t1 = torch.Tensor(data1).long() # size=2,3,2
data2 = [[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 0], [10, 11]]]
t2 = torch.Tensor(data2).long()
print(t1.size())
print(t2.size())
print("-------view_as()-------")
t2=t2.view_as(t1)
print(t2)
print(t2.size())
输出结果：
可以看出经过view_as()操作后，t2 Tensor转变为了与t1 相同的形状。（需要重新对t2赋值，这是因为不是进行的原地操作）