联邦学习实验

《联邦学习实战》第3章：用Python从零实现横向联邦图像分类

参考：Practicing-Federated-Learning/chapter03_Python_image_classification at main · FederatedAI/Practicing-Federated-Learning · GitHub

联邦学习实战：用python从零实现横向联邦图像分类_python_image_classification_FR888的博客-CSDN博客

《联邦学习实战》：从零开始通过联邦学习实现图像分类_def process_selftrain(clients, server, local_epoch_MMashiro的博客-CSDN博客

一、环境配置

已配好环境：参考cpu版的Pytorch环境配置_pytorchcpu版本_萤石ym的博客-CSDN博客

anaconda python3.8 torch(cpu)  在pytorch1

二、导入项目

1、File》setting：

 2、命令行 python main.py -c ./utils/conf.json

cifar10数据集将会被下载到data下：

正在跑，0 1 2指的是客户端1训练了三轮，用的cpu跑，特别慢

Epoch 45, acc: 76.170000, loss: 0.692326
Epoch 46, acc: 76.360000, loss: 0.681031
Epoch 47, acc: 76.700000, loss: 0.670521
Epoch 49, acc: 77.140000, loss: 0.658076
50轮已接近收敛了

三、分析

3.1 conf.json配置文件：根据需要修改

//注：json文件不能这样添加注释
{
	"model_name" : "resnet18",
	"no_models" : 10,     <!--客户端数量-->
	"type" : "cifar",    //数据集信息
	"global_epochs" : 20,    //全局迭代次数，即服务端与客户端的通信迭代次数
	"local_epochs" : 3,    //本地模型训练迭代次数
	"k" : 5,               //每一轮迭代时，服务端会从所有客户端中挑选k个客户端参与训练。
	"batch_size" : 32,    //本地训练每一轮的样本
	//下边三个为超参数设置
	"lr" : 0.001,
	"momentum" : 0.0001,
	"lambda" : 0.1
}

3.2 数据集：datasets.py

from torchvision.transforms import transforms
from torchvision import datasets
# 获取数据集
def get_dataset(dir, name):
    if name == 'mnist':
        # root: 数据路径
        # train参数表示是否是训练集或者测试集
        # download=true表示从互联网上下载数据集并把数据集放在root路径中
        # transform：图像类型的转换
        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： 切割中心点的位置随机选取
            transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            # transforms.Normalize： 给定均值：(R,G,B) 方差：（R，G，B），将会把Tensor正则化
            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

3.3 客户端：client.py

客户端类包括两种函数：构造函数 + 本地训练函数
客户端主要工作包括：
首先，将配置信息复制到客户端中；
然后，按照配置中的模型信息获取模型，通常由服务端将模型参数传递给客户端，客户端将该全局模型覆盖掉本地模型；
最后，配置本地训练数据
这里通过torchvision中的datasets模块获取cifar10数据集后，按客户端ID进行切分，不同客户端拥有不同的子数据集，相互之间没有交集

3.3.1 定义构造函数

class Client(object):
    def __init__(self, conf, model, train_dataset, id = 1):
       # 配置文件
       self.conf = conf
       # 客户端本地模型
       self.local_model = model
       # 客户端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'])
       indices = all_range[id * data_len: (id + 1) * data_len]
       self.train_loader = torch.utils.data.DataLoader(
           self.train_dataset,
           batch_size=conf["batch_size"],
           # sampler定义从数据集中提取样本的策略
           sampler=torch.utils.data.sampler.SubsetRandomSampler(indices)
       )

3.3.2 定义模型本地训练函数

这里使用交叉熵作为本地模型的损失函数，利用梯度下降求解并更新参数值

def local_train(self, model):
    for name, param in model.state_dict().items():
        # 客户端首先用服务器端下发的全局模型覆盖本地模型
        self.local_model.state_dict()[name].copy_(param.clone())
    # 定义最优化函数器，用于本地模型训练
    optimizer = torch.optim.SGD(self.local_model.parameters(), lr=self.conf['lr'], momentum=self.conf['momentum'])

    # 本地训练模型
    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()
            # 梯度清0
            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("Client %d local train done" % self.client_id)
    return diff

3.4 服务端 server.py

服务端类Server主要包括三种函数：

1. 定义构造函数
2. 模型聚合函数
3. 模型评估函数

3.4.1定义构造函数

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,
	    # 设置单个批次大小
	    batch_size=self.conf["batch_size"],
	    # 打乱数据集
	    shuffle=True
	  )

3.4.2 模型聚合函数

这里使用的是经典的FedAvg算法

def model_aggregate(self, weight_accumulator):
	# weight_accumulatot存储了每一个客户端的上传参数变化值
	# 遍历服务器的全局模型
	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)

3.4.3 模型评估函数

对当前的全局模型，利用评估数据评估当前的全局模型性能

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]
        # 将数据存储到gpu
        if torch.cuda.is_available():
            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()
        # 获取最大的对数概率的索引值，在预测结果中选择可能性最大的作为最终的分类结果
        pred = output.data.max(1)[1]
        # 统计预测结果与真实标签相同的总个数
        correct += pred.eq(target.data.view_as(pred)).cpu().sum().item()
    # 计算准确率
    acc = 100.0 * (float(correct) / float(dataset_size))
    # 计算损失值    
    total_1 = total_loss / dataset_size                     
    return acc, total_1

3.5模型文件models.py

项目文件夹下models.py文件，用于各种机器学习模型使用

import torch
from torchvision import models


# 各种机器学习模型
def get_model(name="vgg16", pretrained=True):
    if name == "resnet18":
        model = models.resnet18(pretrained=pretrained)
    elif name == "resnet50":
        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

3.6 主文件main.py

首先，读取配置文件信息

# 读取配置文件
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 = []

# 添加10个客户端到列表
for c in range(conf["no_models"]):
    clients.append(Client(conf, server.global_model, train_datasets, c))

每一轮迭代，服务器会从当前客户端中随机挑选一部分参与本轮训练，被选中的客户端用本地训练接口local_train进行本地训练，最后服务端调用模型聚合函数更新全局模型

for e in range(conf["global_epochs"]):
    print("Global Epoch %d" % e)
    # 随机采样k个客户端参与本轮联邦训练
    candidates = random.sample(clients, conf["k"])
    print("select clients is: ")
    for c in candidates:
        print(c.client_id)

    weight_accumulator = {}

    # 初始化模型参数weight_accumulator
    for name, params in server.global_model.state_dict().items():
        # 生成一个和参数矩阵大小相同的0矩阵
        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))