学习实例
参考资料:联邦学习实战{杨强}https://book.douban.com/subject/35436587/
项目地址:https://github.com/FederatedAI/Practicing-Federated-Learning/tree/main/chapter03_Python_image_classification
一、环境准备
GPU安装CUDA、cuDNN
python3.7、PyTorch
数据集选择:CIFAR10
模型:ResNet-18
实验基本流程
基本的流程:
1服务器按配置生成初始化模型,客户端按照顺序将数据集横向不重叠切割
2服务器将全局模型发送给客户端
3客户端接收全局模型(来自服务器)通过本地多次迭代计算本地参数差值返回给服务器
4服务器聚合各个客户端差值更新模型,再评估当前模型性能
5如果性能未达标,则重复2过程,否则结束
二、环境安装【服务器conda安装】
激活虚拟环境
安装需要的库
conda install torchvision
conda install torch
三、项目创建【pycharm进行debug】
新建项目
项目设置
四、项目运行及结果分析【pycharm】
这里设置参数可以直接右键运行(或者在项目所在文件下使用命令行 python main.py -c -c ./utils/conf.json)
运行结果
五、代码分析
配置文件
项目文件夹下./utils/conf.json创建配置文件:
{
"model_name" : "resnet18",
"no_models" : 10,
"type" : "cifar",
"global_epochs" : 20,
"local_epochs" : 3,
"k" : 6,
"batch_size" : 32,
"lr" : 0.001,
"momentum" : 0.0001,
"lambda" : 0.1
}
model_name:模型名称
no_models:客户端总数量
type:数据集信息
global_epochs:全局迭代次数,即服务端与客户端的通信迭代次数
local_epochs:本地模型训练迭代次数
k:每一轮迭代时,服务端会从所有客户端中挑选k个客户端参与训练。
batch_size:本地训练每一轮的样本数
本地训练的超参数设置
lr:学习率,用于控制模型参数更新的步长,这里是0.001
momentum:动量,用于加速训练过程中的参数更新,这里是0.0001
lambda:正则化参数,用于控制模型的复杂度,这里是0.1
数据集代码
datasets.py
import torchvision as tv
# 获取数据集
def get_dataset(dir, name):
if name == 'mnist':
# root: 数据路径
# train参数表示是否是训练集或者测试集
# download=true表示从互联网上下载数据集并把数据集放在root路径中
# transform:图像类型的转换
train_dataset = tv.datasets.MNIST(dir, train=True, download=True, transform=tv.transforms.ToTensor())
eval_dataset = tv.datasets.MNIST(dir, train=False, transform=tv.transforms.ToTensor())
elif name == 'cifar':
# 设置两个转换格式
# transforms.Compose 是将多个transform组合起来使用(由transform构成的列表)
transform_train = tv.transforms.Compose([
# transforms.RandomCrop: 切割中心点的位置随机选取
tv.transforms.RandomCrop(32, padding=4), tv.transforms.RandomHorizontalFlip(),
tv.transforms.ToTensor(),
# transforms.Normalize: 给定均值:(R,G,B) 方差:(R,G,B),将会把Tensor正则化
tv.transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
transform_test = tv.transforms.Compose([
tv.transforms.ToTensor(),
tv.transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
train_dataset = tv.datasets.CIFAR10(dir, train=True, download=True, transform=transform_train)
eval_dataset = tv.datasets.CIFAR10(dir, train=False, transform=transform_test)
return train_dataset, eval_dataset
服务端的主要功能是模型的聚合、评估,最终的模型也是在服务器上生成
定义全局联邦平均FedAvg聚合函数:
FedAvg算法的公式如下:
server.py
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"])
# 生成一个测试集合加载器
self.eval_loader = torch.utils.data.DataLoader(eval_dataset,
# 设置单个批次大小32
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():
# 更新每一层乘上正则化参数
update_per_layer = weight_accumulator[name] * self.conf["lambda"]
# 累加和
if data.type() != update_per_layer.type():
# 因为update_per_layer的type是floatTensor,所以将起转换为模型的LongTensor(有一定的精度损失)
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):
data, target = batch
# 获取所有的样本总量大小
dataset_size += data.size()[0]
if torch.cuda.is_available():
# 存储到gpu
data = data.cuda()
target = target.cuda()
# 加载到模型中
output = self.global_model(data)
# 聚合所有的损失 cross_entropy交叉熵函数计算损失
total_loss += torch.nn.functional.cross_entropy(output, target,
reduction='sum').item() # sum up batch loss
# 获取最大的对数概率的索引值, 即在所有预测结果中选择可能性最大的作为最终的分类结果
pred = output.data.max(1)[1] # get the index of the max log-probability
# 统计预测结果与真实标签target的匹配总个数
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
客户端
客户端的主要功能是:
接受服务器下发的指令和全局模型
利用本地数据进行局部模型训练
client.py
import models, torch, copy
class Client(object):
# 构造函数
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)))
data_len = int(len(self.train_dataset) / self.conf['no_models'])
train_indices = all_range[id * data_len: (id + 1) * data_len]
# 生成一个数据加载器
self.train_loader = torch.utils.data.DataLoader(
# 制定父集合
self.train_dataset,
# batch_size每个batch加载多少个样本
batch_size=conf["batch_size"],
# 指定子集合
# sampler定义从数据集中提取样本的策略
sampler=torch.utils.data.sampler.SubsetRandomSampler(
train_indices))
#本地模型训练函数:采用交叉熵作为本地训练的损失函数,并使用梯度下降来求解参数
def local_train(self, model):
# 整体的过程:拉取服务器的模型,通过部分本地数据集训练得到局部模型
for name, param in model.state_dict().items():
# 客户端首先用服务器端下发的全局模型覆盖本地模型
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))
# 本地训练模型开始
self.local_model.train()
for e in range(self.conf["local_epochs"]):
for batch_id, batch in enumerate(self.train_loader):
data, target = batch
# 加载到gpu
if torch.cuda.is_available():
data = data.cuda()
target = target.cuda()
# 梯度
optimizer.zero_grad()
# 训练预测
output = self.local_model(data)
# 计算损失函数 cross_entropy交叉熵误差
loss = torch.nn.functional.cross_entropy(output, target)
# 反向传播
loss.backward()
# 更新参数
optimizer.step()
print("Epoch %d done." % e)
# 创建差值字典(结构与模型参数同规格),用于记录差值
diff = dict()
for name, data in self.local_model.state_dict().items():
# 计算训练后与训练前的差值
diff[name] = (data - model.state_dict()[name])
# print(diff[name])
# 客户端返回差值
return diff
主函数
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')
# 获取所有的参数
args = parser.parse_args()
# 读取配置文件
with open(args.conf, 'r') as f:
conf = json.load(f)
# 获取数据集, 加载描述信息
train_datasets, eval_datasets = datasets.get_dataset("./data/", conf["type"])
# 开启服务器
server = Server(conf, eval_datasets)
# 客户端列表
clients = []
# 添加no_models客户端到列表
for c in range(conf["no_models"]):
clients.append(Client(conf, server.global_model, train_datasets, c))
print("\n\n")
# 全局模型训练
for e in range(conf["global_epochs"]):
# 每次训练都是从clients列表中随机采样k个进行本轮训练
candidates = random.sample(clients, conf["k"])
# 权重累计
weight_accumulator = {}
# 初始化空模型参数weight_accumulator
for name, params in server.global_model.state_dict().items():
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():
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))