联邦学习开山之作代码解读与收获

参考:联邦学习代码解读,超详细_一只揪°的博客-CSDN博客_联邦学习代码

参考文献:[1602.05629] Communication-Efficient Learning of Deep Networks from Decentralized Data (arxiv.org)

参考代码:GitHub - AshwinRJ/Federated-Learning-PyTorch: Implementation of Communication-Efficient Learning of Deep Networks from Decentralized Data


        今天尝试阅读开山之作的代码。

目录

一、加载参数——options.py        

二、数据IID和非IID采样——sampling.py

1.mnist_iid()

2.mnist_noniid()

3.mnist_noniid()

4.cifar_iid()、cifar_noniid()

三、本地模型参数更新——update.py

1.DatasetSplit(Dataset)

2.LocalUpdate(object)

 3.test_inference(self,model)

四、应用集——utils.py

1.get_dataset(args)

2.average_weights(w)

3.exp_details(args)

五、模型设置——models.py

1.MLP多层感知机模型

2.CNN卷积神经网络

3.自创模型

六、主函数——federated_main.py

七、作图

八、个人总结


一、加载参数——options.py        

import argparse


def args_parser():
    parser = argparse.ArgumentParser()

    # federated arguments (Notation for the arguments followed from paper)
    parser.add_argument('--epochs', type=int, default=10,
                        help="number of rounds of training")
    parser.add_argument('--num_users', type=int, default=100,
                        help="number of users: K")
    parser.add_argument('--frac', type=float, default=0.1,
                        help='the fraction of clients: C')
    parser.add_argument('--local_ep', type=int, default=10,
                        help="the number of local epochs: E")
    parser.add_argument('--local_bs', type=int, default=10,
                        help="local batch size: B")
    parser.add_argument('--lr', type=float, default=0.01,
                        help='learning rate')
    parser.add_argument('--momentum', type=float, default=0.5,
                        help='SGD momentum (default: 0.5)')

    # model arguments
    parser.add_argument('--model', type=str, default='mlp', help='model name')
    parser.add_argument('--kernel_num', type=int, default=9,
                        help='number of each kind of kernel')
    parser.add_argument('--kernel_sizes', type=str, default='3,4,5',
                        help='comma-separated kernel size to \
                        use for convolution')
    parser.add_argument('--num_channels', type=int, default=1, help="number \
                        of channels of imgs")
    parser.add_argument('--norm', type=str, default='batch_norm',
                        help="batch_norm, layer_norm, or None")
    parser.add_argument('--num_filters', type=int, default=32,
                        help="number of filters for conv nets -- 32 for \
                        mini-imagenet, 64 for omiglot.")
    parser.add_argument('--max_pool', type=str, default='True',
                        help="Whether use max pooling rather than \
                        strided convolutions")

    # other arguments
    parser.add_argument('--dataset', type=str, default='mnist', help="name \
                        of dataset")
    parser.add_argument('--num_classes', type=int, default=10, help="number \
                        of classes")
    parser.add_argument('--gpu', default=None, help="To use cuda, set \
                        to a specific GPU ID. Default set to use CPU.")
    parser.add_argument('--optimizer', type=str, default='sgd', help="type \
                        of optimizer")
    parser.add_argument('--iid', type=int, default=1,
                        help='Default set to IID. Set to 0 for non-IID.')
    parser.add_argument('--unequal', type=int, default=0,
                        help='whether to use unequal data splits for  \
                        non-i.i.d setting (use 0 for equal splits)')
    parser.add_argument('--stopping_rounds', type=int, default=10,
                        help='rounds of early stopping')
    parser.add_argument('--verbose', type=int, default=1, help='verbose')
    parser.add_argument('--seed', type=int, default=1, help='random seed')
    args = parser.parse_args()
    return args

        这里使用argparse输入了三类参数,分别是联邦参数模型参数其他参数。其中联邦参数:

  • epochs:训练轮数,10
  • num_users:用户数量K,默认100
  • frac:用户选取比例C,默认0.1
  • local_ep:本地训练数量E,默认10
  • local_bs:本地训练批量B,默认10
  • lr:学习率,默认0.01
  • momentum:SGD动量(为什么SGD有动量?),默认0.5

        模型参数:

  • model:模型名称,默认mlp,即全连接神经网络
  • kernel_num:卷积核数量,默认9个
  • kernel_sizes:卷积核大小,默认3,4,5
  • num_channels:图像通道数,默认1
  • norm:归一化方式,可以是BN和LN
  • num_filters:过滤器数量,默认32
  • max_pool:最大池化,默认为True

        其他设置:

  • dataset:选择什么数据集,默认mnist
  • num_class:分类数量,默认10
  • gpu:默认使用,可以填入具体cuda编号
  • optimizer:优化器,默认是SGD算法
  • iid:独立同分布,默认是1,即是独立同分布
  • unequal:是否平均分配数据集,默认0,即是
  • stopping_rounds:停止轮数,默认是10
  • verbose:日志显示,0不输出,1输出带进度条的日志,2输出不带进度条的日志
  • seed: 随机数种子,默认1

        最后args_parser()函数会返回args,里面包含了控制台输入的参数。


二、数据IID和非IID采样——sampling.py

        这个文件从mnist和cifar-10采集IID和非IID的数据。

1.mnist_iid()

def mnist_iid(dataset, num_users):
    """
    Sample I.I.D. client data from MNIST dataset
    :param dataset:
    :param num_users:
    :return: dict of image index
    """
    num_items = int(len(dataset)/num_users)
    dict_users, all_idxs = {}, [i for i in range(len(dataset))]
    for i in range(num_users):
        dict_users[i] = set(np.random.choice(all_idxs, num_items,
                                             replace=False))
        all_idxs = list(set(all_idxs) - dict_users[i])
    return dict_users

        随机给100个用户选600个随机的样本。

2.mnist_noniid()

def mnist_noniid(dataset, num_users):
    """
    Sample non-I.I.D client data from MNIST dataset
    :param dataset:
    :param num_users:
    :return:
    """
    # 60,000 training imgs -->  200 imgs/shard X 300 shards
    num_shards, num_imgs = 200, 300
    idx_shard = [i for i in range(num_shards)]
    dict_users = {i: np.array([]) for i in range(num_users)}
    idxs = np.arange(num_shards*num_imgs)
    labels = dataset.train_labels.numpy()

    # sort labels
    idxs_labels = np.vstack((idxs, labels))
    idxs_labels = idxs_labels[:, idxs_labels[1, :].argsort()]
    idxs = idxs_labels[0, :]

    # divide and assign 2 shards/client
    for i in range(num_users):
        rand_set = set(np.random.choice(idx_shard, 2, replace=False))
        idx_shard = list(set(idx_shard) - rand_set)
        for rand in rand_set:
            dict_users[i] = np.concatenate(
                (dict_users[i], idxs[rand*num_imgs:(rand+1)*num_imgs]), axis=0)
    return dict_users
  • num_shards:把60000个训练集图片分为200份
  • [i for i in range()]:可以生成一个递增list
  • {i: np.array([]) for i in range(num_users)}:以大括号生成100个用户的字典
  • np.vstack((idxs, labels)):把编号和标签堆叠起来,形成一个(2,60000)的数组
  • idxs_labels = idxs_labels[:, idxs_labels[1, :].argsort()]:argsort函数的作用是,输出数组中的元素从小到大排序后的索引数组值

        经过筛选之后,获得了由小到大的标签索引idxs。然后进行用户分片。

  • np.random.choice():从切片序号中选出两个序号,replace参数表示不放回取样
  • idxs[rand*num_imgs:(rand+1)*num_imgs]:取连续的300个排序后的索引号
  • np.concatenate():从哪个维度拼哪个维度就会增加,这里从200个索引号中随机选取了两个随机数,把这两个随机数对应位置的数据给连起来了

        最后函数返回了每个用户以及所对应的600个数据的字典。

3.mnist_noniid()

def mnist_noniid_unequal(dataset, num_users):
"""
    Sample non-I.I.D client data from MNIST dataset s.t clients
    have unequal amount of data
    :param dataset:
    :param num_users:
    :returns a dict of clients with each clients assigned certain
    number of training imgs
    """

        有点长,我分着说。把60000张数据分为1200份:

    # 60,000 training imgs --> 50 imgs/shard X 1200 shards
    num_shards, num_imgs = 1200, 50
    idx_shard = [i for i in range(num_shards)]
    dict_users = {i: np.array([]) for i in range(num_users)}
    idxs = np.arange(num_shards*num_imgs)
    labels = dataset.train_labels.numpy()

        获得排序后的索引号:

    # sort labels
    idxs_labels = np.vstack((idxs, labels))
    idxs_labels = idxs_labels[:, idxs_labels[1, :].argsort()]
    idxs = idxs_labels[0, :]

        设置每个用户所持有的数据份数范围:

    # Minimum and maximum shards assigned per client:
    min_shard = 1
    max_shard = 30

        也就是说,每个用户至少拥有1×50=50张图片,至多拥有30*50=1500张图片。

        接下来要把这1200份分给这些用户,并且保证每个用户至少被分到一个数据,且每个数据都要被分到。

    # Divide the shards into random chunks for every client
    # s.t the sum of these chunks = num_shards
    random_shard_size = np.random.randint(min_shard, max_shard+1,
                                          size=num_users)
    random_shard_size = np.around(random_shard_size /
                                  sum(random_shard_size) * num_shards)
    random_shard_size = random_shard_size.astype(int)
  • np.random.randint:返回为一个前闭后开的区间的列表,长度为用户数量
  • np.around:四舍六入,五归偶数

        经过这一步,所有的份数都被等比地调整,使其总和接近于为1200。(因为有小数被四舍六入,所以不严格等于1200)所以接下来就要针对这不严格的部分进行调整和分配。

    # Assign the shards randomly to each client
    if sum(random_shard_size) > num_shards:

        for i in range(num_users):
            # First assign each client 1 shard to ensure every client has
            # atleast one shard of data
            rand_set = set(np.random.choice(idx_shard, 1, replace=False))
            idx_shard = list(set(idx_shard) - rand_set)
            for rand in rand_set:
                dict_users[i] = np.concatenate(
                    (dict_users[i], idxs[rand*num_imgs:(rand+1)*num_imgs]),
                    axis=0)

        random_shard_size = random_shard_size-1

        # Next, randomly assign the remaining shards
        for i in range(num_users):
            if len(idx_shard) == 0:
                continue
            shard_size = random_shard_size[i]
            if shard_size > len(idx_shard):
                shard_size = len(idx_shard)
            rand_set = set(np.random.choice(idx_shard, shard_size,
                                            replace=False))
            idx_shard = list(set(idx_shard) - rand_set)
            for rand in rand_set:
                dict_users[i] = np.concatenate(
                    (dict_users[i], idxs[rand*num_imgs:(rand+1)*num_imgs]),
                    axis=0)
    else:

        for i in range(num_users):
            shard_size = random_shard_size[i]
            rand_set = set(np.random.choice(idx_shard, shard_size,
                                            replace=False))
            idx_shard = list(set(idx_shard) - rand_set)
            for rand in rand_set:
                dict_users[i] = np.concatenate(
                    (dict_users[i], idxs[rand*num_imgs:(rand+1)*num_imgs]),
                    axis=0)

        if len(idx_shard) > 0:
            # Add the leftover shards to the client with minimum images:
            shard_size = len(idx_shard)
            # Add the remaining shard to the client with lowest data
            k = min(dict_users, key=lambda x: len(dict_users.get(x)))
            rand_set = set(np.random.choice(idx_shard, shard_size,
                                            replace=False))
            idx_shard = list(set(idx_shard) - rand_set)
            for rand in rand_set:
                dict_users[k] = np.concatenate(
                    (dict_users[k], idxs[rand*num_imgs:(rand+1)*num_imgs]),
                    axis=0)

    return dict_users

        最后会获得随机分配的用户持有的非IID数据的索引字典.。

4.cifar_iid()、cifar_noniid()

        没有区别,不写了


三、本地模型参数更新——update.py

1.DatasetSplit(Dataset)

        先来看看Dataset类的官方解释:Dataset可以是任何东西,但它始终包含一个__len__函数(通过Python中的标准函数len调用)和一个用来索引到内容中的__getitem__函数。

class DatasetSplit(Dataset):
    """An abstract Dataset class wrapped around Pytorch Dataset class.
    """

    def __init__(self, dataset, idxs):
        self.dataset = dataset
        self.idxs = [int(i) for i in idxs]

    def __len__(self):
        return len(self.idxs)

    def __getitem__(self, item):
        image, label = self.dataset[self.idxs[item]]
        return torch.tensor(image), torch.tensor(label)

        这部分代码重写了Dataset类:

  • 重写了__len__(self)方法,返回数据列表长度,即数据集的样本数量
  • 重写了__getitem__(self,item)方法,获取image和label的张量

2.LocalUpdate(object)

        这是本地更新模型的代码,有点多我分着说:

class LocalUpdate(object):...

        首先是构造函数,首先定义了参数和日志,然后从train_val_test()函数获取了数据加载器,随后指定了运算设备。

        比较重要的是这里的损失函数是NLL损失函数,它跟交叉熵相似,唯一的区别在于NLL的log里面对结果进行了一次Softmax。

    def __init__(self, args, dataset, idxs, logger):
        self.args = args
        self.logger = logger
        self.trainloader, self.validloader, self.testloader = self.train_val_test(
            dataset, list(idxs))
        self.device = 'cuda' if args.gpu else 'cpu'
        # Default criterion set to NLL loss function
        self.criterion = nn.NLLLoss().to(self.device)

        接下来是train_val_test()函数,它用来分割数据集。输入数据集和索引,按照8:1:1来划分。注意到在指定batchsize的时候,除了训练集是从args参数里指定的,val和test都是取了总数的十分之一。

    def train_val_test(self, dataset, idxs):
        """
        Returns train, validation and test dataloaders for a given dataset
        and user indexes.
        """
        # split indexes for train, validation, and test (80, 10, 10)
        idxs_train = idxs[:int(0.8*len(idxs))]
        idxs_val = idxs[int(0.8*len(idxs)):int(0.9*len(idxs))]
        idxs_test = idxs[int(0.9*len(idxs)):]

        trainloader = DataLoader(DatasetSplit(dataset, idxs_train),
                                 batch_size=self.args.local_bs, shuffle=True)
        validloader = DataLoader(DatasetSplit(dataset, idxs_val),
                                 batch_size=int(len(idxs_val)/10), shuffle=False)
        testloader = DataLoader(DatasetSplit(dataset, idxs_test),
                                batch_size=int(len(idxs_test)/10), shuffle=False)
        return trainloader, validloader, testloader

        接下来是本地权重更新函数,输入模型全局更新的回合数,输出更新后的权重和损失平均值。首先选择了优化器,然后开始训练循环。

    def update_weights(self, model, global_round):
        # Set mode to train model
        model.train()
        epoch_loss = []

        # Set optimizer for the local updates
        if self.args.optimizer == 'sgd':
            optimizer = torch.optim.SGD(model.parameters(), lr=self.args.lr,
                                        momentum=0.5)
        elif self.args.optimizer == 'adam':
            optimizer = torch.optim.Adam(model.parameters(), lr=self.args.lr,
                                         weight_decay=1e-4)

        for iter in range(self.args.local_ep):
            batch_loss = []
            for batch_idx, (images, labels) in enumerate(self.trainloader):
                images, labels = images.to(self.device), labels.to(self.device)

                model.zero_grad()
                log_probs = model(images)
                loss = self.criterion(log_probs, labels)
                loss.backward()
                optimizer.step()

                if self.args.verbose and (batch_idx % 10 == 0):
                    print('| Global Round : {} | Local Epoch : {} | [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                        global_round, iter, batch_idx * len(images),
                        len(self.trainloader.dataset),
                        100. * batch_idx / len(self.trainloader), loss.item()))
                self.logger.add_scalar('loss', loss.item())
                batch_loss.append(loss.item())
            epoch_loss.append(sum(batch_loss)/len(batch_loss))

        return model.state_dict(), sum(epoch_loss) / len(epoch_loss)
  • self.logger.add_scalar('loss', loss.item()):这个函数是用来保存程序中的数据,然后利用tensorboard工具来进行可视化的
  • 每经过一次本地轮次,统计当前的loss,用于最后的平均损失统计
  • model.state_dict():是Pytorch中用于查看网络参数的方法,可以用torch.save()保存成pth文件

         接下来是评估函数:inference(self,model)。输入为模型,计算准确值、loss值,这里的代码很有参考意义:

    def inference(self, model):
        """ Returns the inference accuracy and loss.
        """

        model.eval()
        loss, total, correct = 0.0, 0.0, 0.0

        for batch_idx, (images, labels) in enumerate(self.testloader):
            images, labels = images.to(self.device), labels.to(self.device)

            # Inference
            outputs = model(images)
            batch_loss = self.criterion(outputs, labels)
            loss += batch_loss.item()

            # Prediction
            _, pred_labels = torch.max(outputs, 1)
            pred_labels = pred_labels.view(-1)
            correct += torch.sum(torch.eq(pred_labels, labels)).item()
            total += len(labels)

        accuracy = correct/total
        return accuracy, loss
  • model.eval():开启模型的评估模式
  • torch.max():第二个参数指维度,即返回第1维度(即行),这里返回了虽大数值的索引
  • pred_labels.view(-1):本意是根据另外一个数来自动调整维度,但是这里只有一个维度,因此就会将X里面的所有维度数据转化成一维的,并且按先后顺序排列。
  • torch.eq():对两个张量Tensor进行逐元素的比较,若相同位置的两个元素相同,则返回True;若不同,返回False。

        这里的函数通取测试集图像和标签,模型出结果后计算loss然后累加,

 3.test_inference(self,model)

        与LocalUpdate中的inference函数完全一致,只不过这里的输入参数除了args和model,还要指定test_dataset:

def test_inference(args, model, test_dataset):
    """ Returns the test accuracy and loss.
    """

    model.eval()
    loss, total, correct = 0.0, 0.0, 0.0

    device = 'cuda' if args.gpu else 'cpu'
    criterion = nn.NLLLoss().to(device)
    testloader = DataLoader(test_dataset, batch_size=128,
                            shuffle=False)

    for batch_idx, (images, labels) in enumerate(testloader):
        images, labels = images.to(device), labels.to(device)

        # Inference
        outputs = model(images)
        batch_loss = criterion(outputs, labels)
        loss += batch_loss.item()

        # Prediction
        _, pred_labels = torch.max(outputs, 1)
        pred_labels = pred_labels.view(-1)
        correct += torch.sum(torch.eq(pred_labels, labels)).item()
        total += len(labels)

    accuracy = correct/total
    return accuracy, loss

四、应用集——utils.py

        这里面封装了一些工具函数:get_dataset(),average_weights(),exp_details()

1.get_dataset(args)

        get_dataset(args)根据命令台参数获取相应的数据集和用户数据字典。就是个if else,有点简单就不说了。

2.average_weights(w)

        返回权重的平均值,即执行联邦平均算法:

def average_weights(w):
    """
    Returns the average of the weights.
    """
    w_avg = copy.deepcopy(w[0])
    for key in w_avg.keys():
        for i in range(1, len(w)):
            w_avg[key] += w[i][key]
        w_avg[key] = torch.div(w_avg[key], len(w))
    return w_avg
  • w:这个w是经过多轮本地训练后统计的权重list,在参数默认的情况下,是一个长度为10的列表,而每个元素都是一个字典,每个字典都包含了模型参数的名称(比如layer_input.weight或者layer_hidden.bias),以及其权重具体的值。
  • copy.deepcopy():深度复制,被复制的对象不会随着复制的对象的改变而改变。这里复制了第一个用户的权重字典。

        随后,对于每一类参数进行循环,累加每个用户模型里对应参数的值,最后取平均获得平均后的模型。

3.exp_details(args)

        可视化命令台参数args:

def exp_details(args):
    print('\nExperimental details:')
    print(f'    Model     : {args.model}')
    print(f'    Optimizer : {args.optimizer}')
    print(f'    Learning  : {args.lr}')
    print(f'    Global Rounds   : {args.epochs}\n')

    print('    Federated parameters:')
    if args.iid:
        print('    IID')
    else:
        print('    Non-IID')
    print(f'    Fraction of users  : {args.frac}')
    print(f'    Local Batch size   : {args.local_bs}')
    print(f'    Local Epochs       : {args.local_ep}\n')
    return

五、模型设置——models.py

        这个文件设置了一些比较常见的网络模型

1.MLP多层感知机模型

class MLP(nn.Module):
    def __init__(self, dim_in, dim_hidden, dim_out):
        super(MLP, self).__init__()
        self.layer_input = nn.Linear(dim_in, dim_hidden)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout()
        self.layer_hidden = nn.Linear(dim_hidden, dim_out)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = x.view(-1, x.shape[1]*x.shape[-2]*x.shape[-1])
        x = self.layer_input(x)
        x = self.dropout(x)
        x = self.relu(x)
        x = self.layer_hidden(x)
        return self.softmax(x)
  • nn.Dropout():你懂得,不懂就搜

2.CNN卷积神经网络

        太多了不予展示。

3.自创模型

        这里原代码是modelC,其构造函数下,super第一个参数是AllConvNet,在编译器中会报错。但是这里并非打错,而是让用户自定义。


六、主函数——federated_main.py

        (这里我贴的代码是我更改了注释的)

        首先是库的引用:

import os
import copy
import time
import pickle
import numpy as np
from tqdm import tqdm

import torch
from tensorboardX import SummaryWriter

from options import args_parser
from update import LocalUpdate, test_inference
from models import MLP, CNNMnist, CNNFashion_Mnist, CNNCifar
from utils import get_dataset, average_weights, exp_details

        随后直接开始主函数:

if __name__ == '__main__':
    start_time = time.time()

    # 定义路径
    path_project = os.path.abspath('..')  # 上级目录的绝对路径
    logger = SummaryWriter('../logs')  # python可视化工具

    args = args_parser()  # 输入命令行参数
    exp_details(args)  # 显示命令行参数情况

        由于是调试状态运行,所以没有更改参数,参数情况如下所示:

联邦学习开山之作代码解读与收获_第1张图片

         接下来加载数据集和用户数据字典:

    # 判断GPU是否可用:
    if args.gpu:
        torch.cuda.set_device(args.gpu)
    device = 'cuda' if args.gpu else 'cpu'

    # 加载数据集,用户本地数据字典
    train_dataset, test_dataset, user_groups = get_dataset(args)

        这里会返回60000的训练集,10000的测试集,以及长度为100的用户字典,用户字典是100个用户到各自600个IID训练数据的映射。

        然后开始建立模型,这里模型选择的是多层感知机:

    # 建立模型
    if args.model == 'cnn':
        # 卷积神经网络
        if args.dataset == 'mnist':
            global_model = CNNMnist(args=args)
        elif args.dataset == 'fmnist':
            global_model = CNNFashion_Mnist(args=args)
        elif args.dataset == 'cifar':
            global_model = CNNCifar(args=args)

    elif args.model == 'mlp':
        # 多层感知机
        img_size = train_dataset[0][0].shape
        len_in = 1
        for x in img_size:
            len_in *= x
            global_model = MLP(dim_in=len_in, dim_hidden=64,
                               dim_out=args.num_classes)
    else:
        exit('Error: unrecognized model')

        接下来就是设置模型进行第一轮训练,并复制权重:

    # 设置模型进行训练,并传输给计算设备
    global_model.to(device)
    global_model.train()
    print(global_model)

    # 复制权重
    global_weights = global_model.state_dict()

        模型如下所示:

联邦学习开山之作代码解读与收获_第2张图片

        这是一个输入层784个,隐藏层64个,输出层10个的多层感知机,且设置了0.5的Dropout。

        然后就开始正式训练:

    # 训练
    train_loss, train_accuracy = [], []
    val_acc_list, net_list = [], []
    cv_loss, cv_acc = [], []
    print_every = 2
    val_loss_pre, counter = 0, 0

    for epoch in tqdm(range(args.epochs)):
        local_weights, local_losses = [], []
        print(f'\n | Global Training Round : {epoch + 1} |\n')

        global_model.train()
        m = max(int(args.frac * args.num_users), 1)  # 随机选比例为frac的用户
        idxs_users = np.random.choice(range(args.num_users), m, replace=False)

        for idx in idxs_users:
            local_model = LocalUpdate(args=args, dataset=train_dataset,
                                      idxs=user_groups[idx], logger=logger)
            w, loss = local_model.update_weights(
                model=copy.deepcopy(global_model), global_round=epoch)
            local_weights.append(copy.deepcopy(w))
            local_losses.append(copy.deepcopy(loss))

        # 联邦平均,更新全局权重
        global_weights = average_weights(local_weights)

        # 将更新后的全局权重载入模型
        global_model.load_state_dict(global_weights)

        loss_avg = sum(local_losses) / len(local_losses)
        train_loss.append(loss_avg)

        # 每轮训练,都要计算所有用户的平均训练精度
        list_acc, list_loss = [], []
        global_model.eval()
        for c in range(args.num_users):
            local_model = LocalUpdate(args=args, dataset=train_dataset,
                                      idxs=user_groups[idx], logger=logger)
            acc, loss = local_model.inference(model=global_model)
            list_acc.append(acc)
            list_loss.append(loss)
        train_accuracy.append(sum(list_acc) / len(list_acc))

        # 每i轮打印全局Loss
        if (epoch + 1) % print_every == 0:
            print(f' \nAvg Training Stats after {epoch + 1} global rounds:')
            print(f'Training Loss : {np.mean(np.array(train_loss))}')
            print('Train Accuracy: {:.2f}% \n'.format(100 * train_accuracy[-1]))
  • 老实说除了train_loss,train_accuracy和print_every之外我都不知道其他的是干嘛的
  • tqdm是一个功能强大的进度条,支持在for循环中展示运行时间和进度
  • global_model.train():将模型设置为训练模式
  • idxs_users:随机选取用户的索引列表,这里来说,用户选取比例为0.1,用户总数100,那么就会随机抽取100×0.1=10个用户参与训练
  • 执行本地更新:对于选取的用户执行本地更新,数据集索引由user_groups[idx]获得,并记录更新后的本地参数和损失值
  • 联邦平均:把模型参数字典传入更新函数,返回平均后的模型参数字典,再载入到全局模型中

        每轮结束都统计所有100个用户的训练精度,每i轮都打印全局损失值。

        (注意,你跑模型不停滚动的什么Global Round,Local Epoch,那都是update.py里面的调用LocalUpdate类里的update_weights方法形成的,如果不想他这么频繁的滚动,到这个函数底下注释掉即可)

        全局训练后,模型在测试集的表现:

    # 训练后,测试模型在测试集的表现
    test_acc, test_loss = test_inference(args, global_model, test_dataset)

    print(f' \n Results after {args.epochs} global rounds of training:')
    print("|---- Avg Train Accuracy: {:.2f}%".format(100 * train_accuracy[-1]))
    print("|---- Test Accuracy: {:.2f}%".format(100 * test_acc))

        结果:

联邦学习开山之作代码解读与收获_第3张图片

         最后就是保存目标训练损失和训练精度了,最后输出时间。

    # 保存目标训练损失和训练精度
    file_name = '../save/objects/{}_{}_{}_C[{}]_iid[{}]_E[{}]_B[{}].pkl'. \
        format(args.dataset, args.model, args.epochs, args.frac, args.iid,
               args.local_ep, args.local_bs)

    with open(file_name, 'wb') as f:
        pickle.dump([train_loss, train_accuracy], f)

    print('\n Total Run Time: {0:0.4f}'.format(time.time() - start_time))
  • pkl文件:pickle.dump(数据,f)为写入,pickle.load(文件名)为读出,这里保存了Loss和Accuracy

七、作图

        在代码的最后,作者用注释写出的作图代码:

    # 画图
    import matplotlib
    import matplotlib.pyplot as plt
    matplotlib.use('Agg')

    # 绘制损失曲线
    plt.figure()
    plt.title('训练损失 vs 通信回合数')
    plt.plot(range(len(train_loss)), train_loss, color='r')
    plt.ylabel('训练损失')
    plt.xlabel('通信回合数')
    plt.savefig('../save/fed_{}_{}_{}_C[{}]_iid[{}]_E[{}]_B[{}]_loss.png'.
                format(args.dataset, args.model, args.epochs, args.frac,
                       args.iid, args.local_ep, args.local_bs))

    # 平均准度曲线
    plt.figure()
    plt.title('平均准度 vs 通信回合数')
    plt.plot(range(len(train_accuracy)), train_accuracy, color='k')
    plt.ylabel('平均准度')
    plt.xlabel('通信回合数')
    plt.savefig('../save/fed_{}_{}_{}_C[{}]_iid[{}]_E[{}]_B[{}]_acc.png'.
                format(args.dataset, args.model, args.epochs, args.frac,
                       args.iid, args.local_ep, args.local_bs))

        做图如下:

联邦学习开山之作代码解读与收获_第4张图片联邦学习开山之作代码解读与收获_第5张图片

 


八、个人总结

        这次细读代码让我收获良多,包括代码的组织,一些库的应用以及最重要的联邦学习的机理,作者用简单易懂的代码写出了一篇如此有意义的文章,是在敬佩。不但提高了我的码力,也让我正式跨进fl的大门。

你可能感兴趣的:(联邦学习,深度学习,人工智能,python,pytorch,机器学习)