FedAvg代码详解

FedAvg代码详解

代码位置：https://github.com/shaoxiongji/federated-learning

README

首先查看作者给的README文档：

这一部分给出了论文的地址，如果有兴趣的话可以读一下论文，讲的就是FedAvg的思想以及对独立同分布数据和非独立同分布数据的一些研究。

下来是requirements，就是代码的依赖库，需要安装

如果有不会安装pytorch的同学可以看我之前的博客，从安装CUDA和CUDNN开始的。

运行：

给出了运行的命令，如果使用MLP和CNN模型单独训练就执行main_nn.py文件，使用联邦学习的训练就执行main_fed.py文件。并且给出了命令行参数：–dataset设置训练集；–iid设置数据是否为独立同分布；以及–num_channels设置数据的通道数量，如果是MNIST数据集就是1，CIFAR-10是3；–epochs设置训练的轮数；–gpu设置是否使用GPU；–all_client设置平局所有客户端模型。

作者给出的结果：

项目目录

• data文件夹下存放数据集，有MNIST和CIFAR-10
• models文件夹下存放和模型相关的文件
• save存放训练的结果
• utils存放一些工具

分析代码

所有_init_.py文件都不需要分析

Models目录下代码

Fed.py

import copy
import torch
from torch import nn


def FedAvg(w):
    w_avg = copy.deepcopy(w[0])  # 创建一个深拷贝，用于平均参数的累加
    for k in w_avg.keys():  # 遍历参数字典的键
        for i in range(1, len(w)):  # 遍历参与平均的参数列表
            w_avg[k] += w[i][k]  # 将每个参数按键累加到平均参数中
        w_avg[k] = torch.div(w_avg[k], len(w))  # 将累加的参数值除以参与平均的参数个数，得到平均值
    return w_avg  # 返回平均参数

Fed文件主要做的事情就是对权重求平均

这段代码可能难以理解的就是这个深拷贝，我这里给出一个测试代码，大家可以自行理解，这里就不展开了，有问题可以写在评论区或者私信我。

import copy

a = 1   # =赋值 不可变元素如字符串、数值等
b = a
print('a和b的值:', a, b)
print('a和b的id:', id(a), id(b))
a = 2
print('修改a后a和b的值:', a, b)
print('修改a后a和b的id:', id(a), id(b))
# a和b的值: 1 1
# a和b的id: 2279821764912 2279821764912
# 修改a后a和b的值: 2 1
# 修改a后a和b的id: 2279821764944 2279821764912

c = [1,2,3] # =赋值 可变元素如列表、字典等
d = c
print('c和d的值:', c, d)
print('c和d的id:', id(c), id(d))
c.append(4)
print('修改c后c和d的值:', c, d)
print('修改c后c和d的id:', id(c), id(d))
# c和d的值: [1, 2, 3] [1, 2, 3]
# c和d的id: 2279827005184 2279827005184
# 修改c后c和d的值: [1, 2, 3, 4] [1, 2, 3, 4]
# 修改c后c和d的id: 2279827005184 2279827005184

orignal_list = [1, 2, 3, [4]]   # 定义一个有嵌套层次的列表
copy_list = copy.copy(orignal_list) # 浅拷贝
deepcopy_list = copy.deepcopy(orignal_list) # 深拷贝
print(copy_list, deepcopy_list)
print(id(orignal_list), id(copy_list), id(deepcopy_list))
orignal_list.append(5)
orignal_list[-2].append(6)  # 修改嵌套内的可变对象
print(orignal_list, copy_list, deepcopy_list)
# [1, 2, 3, [4]] [1, 2, 3, [4]]
# 1985362690688 1985365369664 1985365370112
# [1, 2, 3, [4, 6], 5] [1, 2, 3, [4, 6]] [1, 2, 3, [4]]	# copy和deepcopy的差别就是:copy只复制了第一层 而更深的层次还是一种引用;deepcopy是递归复制 所有层次完全复制到新的内存空间

Nets.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Python version: 3.6

import torch
from torch import nn
import torch.nn.functional as F


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()  # ReLU激活函数
        self.dropout = nn.Dropout()  # Dropout层
        self.layer_hidden = nn.Linear(dim_hidden, dim_out)  # 隐藏层

    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)  # Dropout层，用于防止过拟合
        x = self.relu(x)  # ReLU激活函数
        x = self.layer_hidden(x)  # 隐藏层
        return x


class CNNMnist(nn.Module):
    def __init__(self, args):
        super(CNNMnist, self).__init__()
        self.conv1 = nn.Conv2d(args.num_channels, 10, kernel_size=5)  # 第一个卷积层，输入通道数为args.num_channels，输出通道数为10
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)  # 第二个卷积层，输入通道数为10，输出通道数为20
        self.conv2_drop = nn.Dropout2d()  # 二维Dropout层
        self.fc1 = nn.Linear(320, 50)  # 全连接层1，输入大小为320，输出大小为50
        self.fc2 = nn.Linear(50, args.num_classes)  # 全连接层2，输入大小为50，输出大小为args.num_classes

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))  # 第一个卷积层后接ReLU激活函数和最大池化层
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))  # 第二个卷积层后接ReLU激活函数、Dropout层和最大池化层
        x = x.view(-1, x.shape[1]*x.shape[2]*x.shape[3])  # 将输入展平
        x = F.relu(self.fc1(x))  # 全连接层1后接ReLU激活函数
        x = F.dropout(x, training=self.training)  # Dropout层，用于防止过拟合
        x = self.fc2(x)  # 全连接层2，输出最终结果
        return x


class CNNCifar(nn.Module):
    def __init__(self, args):
        super(CNNCifar, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)  # 第一个卷积层，输入通道数为3，输出通道数为6
        self.pool = nn.MaxPool2d(2, 2)  # 最大池化层
        self.conv2 = nn.Conv2d(6, 16, 5)  # 第二个卷积层，输入通道数为6，输出通道数为16
        self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 全连接层1，输入大小为16*5*5，输出大小为120
        self.fc2 = nn.Linear(120, 84)  # 全连接层2，输入大小为120，输出大小为84
        self.fc3 = nn.Linear(84, args.num_classes)  # 全连接层3，输入大小为84，输出大小为args.num_classes

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))  # 第一个卷积层后接ReLU激活函数和最大池化层
        x = self.pool(F.relu(self.conv2(x)))  # 第二个卷积层后接ReLU激活函数和最大池化层
        x = x.view(-1, 16 * 5 * 5)  # 将输入展平
        x = F.relu(self.fc1(x))  # 全连接层1后接ReLU激活函数
        x = F.relu(self.fc2(x))  # 全连接层2后接ReLU激活函数
        x = self.fc3(x)  # 全连接层3，输出最终结果
        return x

Nets.py是定义模型的文件代码，我们在这里实现我们需要使用的模型就可以

test.py

def test_img(net_g, datatest, args):
    net_g.eval()
    # 将模型设置为评估模式 eval函数的作用就是不启用dropout和BN 否则测试的时候不会是训练好的权重

    test_loss = 0
    correct = 0
    data_loader = DataLoader(datatest, batch_size=args.bs)
    # 创建测试集的数据加载器

    for idx, (data, target) in enumerate(data_loader):
        if args.gpu != -1:
            data, target = data.cuda(), target.cuda()
        # 将数据和标签移到GPU上（如果可用）

        log_probs = net_g(data)	# 模型会自动调用forward函数前向传播 对于分类任务 返回值就是每一类的概率 也就是权重w

        test_loss += F.cross_entropy(log_probs, target, reduction='sum').item()
        # 计算批次损失的总和

        y_pred = log_probs.data.max(1, keepdim=True)[1] 
        # 在第一个维度中寻找最大值并不改变维度

        correct += y_pred.eq(target.data.view_as(y_pred)).long().cpu().sum()
        # 计算正确预测的数量

    test_loss /= len(data_loader.dataset)
    # 计算平均损失

    accuracy = 100.00 * correct / len(data_loader.dataset)
    # 计算准确率

    if args.verbose:
        print('\nTest set: Average loss: {:.4f} \nAccuracy: {}/{} ({:.2f}%)\n'.format(
            test_loss, correct, len(data_loader.dataset), accuracy))
    # 打印测试结果（平均损失和准确率）

    return accuracy, test_loss
    # 返回准确率和损失

test.py中只有一个test_img函数，输入模型和数据集，返回准确率和测试的损失

Update.py

class DatasetSplit(Dataset):	
    # Dataset的子类 可以创建一个数据子集对象 只包含原始数据集中特定索引的样本 在分割数据集用于训练和测试时使用
    def __init__(self, dataset, idxs):
        self.dataset = dataset
        self.idxs = list(idxs)

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

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

这个类在实例化的时候，接收参数dataset和idxs，数据集和要分割的索引，将数据集按照索引进行分割

class LocalUpdate(object):	# 本地模型更新 根据本地数据集进行训练并返回
    def __init__(self, args, dataset=None, idxs=None):
        self.args = args
        self.loss_func = nn.CrossEntropyLoss()	# 交叉熵损失
        self.selected_clients = []	# 选择的客户端节点
        self.ldr_train = DataLoader(DatasetSplit(dataset, idxs), batch_size=self.args.local_bs, shuffle=True)	
        # 加载数据集的子集 通过idxs分割 本地的batch_size由命令行参数给出

    def train(self, net):	# 本地训练
        net.train()		# 设置为训练模式 会启用dropout函数和BN函数
        # train and update
        optimizer = torch.optim.SGD(net.parameters(), lr=self.args.lr, momentum=self.args.momentum)	# 梯度下降算法使用随机梯度下降

        epoch_loss = []
        for iter in range(self.args.local_ep):	# 根据local_ep确定本地训练轮数
            batch_loss = []
            for batch_idx, (images, labels) in enumerate(self.ldr_train):
                images, labels = images.to(self.args.device), labels.to(self.args.device)
                net.zero_grad()
                log_probs = net(images)
                loss = self.loss_func(log_probs, labels)
                loss.backward()		
                # 反向传播 用于计算损失函数对模型参数的梯度 从而更新模型参数以最小化损失函数
                optimizer.step()	# 根据
                if self.args.verbose and batch_idx % 10 == 0:
                    print('Update Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                        iter, batch_idx * len(images), len(self.ldr_train.dataset),
                               100. * batch_idx / len(self.ldr_train), loss.item()))
                batch_loss.append(loss.item())
            epoch_loss.append(sum(batch_loss)/len(batch_loss))
        return net.state_dict(), sum(epoch_loss) / len(epoch_loss)

LocalUpdate类主要定义了联邦学习中本地更新的过程，每个客户端使用自己的本地模型进行训练，并将更新后的模型参数传回到中央服务器进行聚合，代码中的形式就是return，net.state_dict()就是模型参数，返回值的第二项就是所有轮次的平均损失。

Utils目录下代码

options.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Python version: 3.6

import argparse


def args_parser():
    parser = argparse.ArgumentParser()
    # federated arguments
    parser.add_argument('--epochs', type=int, default=10, help="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=5, help="the number of local epochs: E")
    parser.add_argument('--local_bs', type=int, default=10, help="local batch size: B")
    parser.add_argument('--bs', type=int, default=128, help="test batch size")
    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)")
    parser.add_argument('--split', type=str, default='user', help="train-test split type, user or sample")

    # 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('--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")
    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('--iid', action='store_true', help='whether i.i.d or not')
    parser.add_argument('--num_classes', type=int, default=10, help="number of classes")
    parser.add_argument('--num_channels', type=int, default=3, help="number of channels of imges")
    parser.add_argument('--gpu', type=int, default=0, help="GPU ID, -1 for CPU")
    parser.add_argument('--stopping_rounds', type=int, default=10, help='rounds of early stopping')
    parser.add_argument('--verbose', action='store_true', help='verbose print')
    parser.add_argument('--seed', type=int, default=1, help='random seed (default: 1)')
    parser.add_argument('--all_clients', action='store_true', help='aggregation over all clients')
    args = parser.parse_args()
    return args

这个代码就没什么好说的，命令行的参数含义

sampling.py

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))]	
    # dict_users是一个空字典 all_idxs是一个从0到len(dataset)-1的连续整数的列表
    for i in range(num_users):
        dict_users[i] = set(np.random.choice(all_idxs, num_items, replace=False))	
        # 为每一个客户端随机选择num_items个索引 并且设置不重复选择 然后将这些索引作为一个集合set存在字典中 使用i作为键 这样每个客户端在在 dict_users 字典中都有一个对应的图像索引集合，包含了随机选择的 num_items 个不重复的图像索引。通过这种方式，确保了每个客户端的数据集合起来能够覆盖整个数据集，并且每个客户端的数据量相等。
        all_idxs = list(set(all_idxs) - dict_users[i])	# 每给一个用户分配数据后 将分配过的数据删去
    return dict_users	# 返回一个用户序号-数据集的字典

从MNIST数据集中抽取独立同分布IID的客户端数据，将数据集划分为多个子数据集给每个客户端使用（但是根据这个代码，我发现其实这个函数并没有考虑数据的同分布问题，只是单纯的数量相同）

import numpy as np

def mnist_noniid(dataset, num_users):
    """
    从MNIST数据集中采样非独立同分布（non-I.I.D.）的客户端数据
    :param dataset: MNIST数据集
    :param num_users: 客户端数量
    :return: 包含客户端数据索引的字典
    """
    num_shards, num_imgs = 200, 300  # 数据集划分的分片数和每个分片的图像数量
    idx_shard = [i for i in range(num_shards)]  # 分片索引列表
    dict_users = {i: np.array([], dtype='int64') for i in range(num_users)}  # 存储客户端数据索引的字典
    idxs = np.arange(num_shards * num_imgs)  # 所有图像的索引
    labels = dataset.train_labels.numpy()  # 所有图像的标签

    # 根据标签对索引进行排序
    idxs_labels = np.vstack((idxs, labels))	# 将两个数据idxs和labels按垂直方向堆叠
    idxs_labels = idxs_labels[:, idxs_labels[1, :].argsort()]	# 按照图像标签进行排序
    idxs = idxs_labels[0, :]	# 取出第一行 即排序后的所有图像索引
	
    # 划分并分配数据
    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

作用就是采取非独立同分布的客户端数据，不像iid那样均分数据，因为取数据的时候先将数据按照类别排好序之后再进行分配，将数据排好序后，根据随机的分片索引在随机的位置上取连续的数据，这样取到的数据很可能是同一个标签的。

def cifar_iid(dataset, num_users):
    """
    Sample I.I.D. client data from CIFAR10 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

和mnist_iid函数完全一致，我也不知道为什么要放两个一模一样的函数。

核心代码

main_fed.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Python version: 3.6

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import copy
import numpy as np
from torchvision import datasets, transforms
import torch

from utils.sampling import mnist_iid, mnist_noniid, cifar_iid
from utils.options import args_parser
from models.Update import LocalUpdate
from models.Nets import MLP, CNNMnist, CNNCifar
from models.Fed import FedAvg
from models.test import test_img


if __name__ == '__main__':
    # parse args
    args = args_parser()	# 解析命令行参数
    args.device = torch.device('cuda:{}'.format(args.gpu) if torch.cuda.is_available() and args.gpu != -1 else 'cpu')

    # load dataset and split users
    if args.dataset == 'mnist':	# 加载MNIST数据集
        trans_mnist = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
        dataset_train = datasets.MNIST('../data/mnist/', train=True, download=True, transform=trans_mnist)
        dataset_test = datasets.MNIST('../data/mnist/', train=False, download=True, transform=trans_mnist)
        # sample users	# 划分数据集给客户端
        if args.iid:
            dict_users = mnist_iid(dataset_train, args.num_users)
        else:
            dict_users = mnist_noniid(dataset_train, args.num_users)
    elif args.dataset == 'cifar':	# 加载CIFAR-10数据集
        trans_cifar = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])	# 因为CIFAR数据集是3通道
        dataset_train = datasets.CIFAR10('../data/cifar', train=True, download=True, transform=trans_cifar)
        dataset_test = datasets.CIFAR10('../data/cifar', train=False, download=True, transform=trans_cifar)
        if args.iid:
            dict_users = cifar_iid(dataset_train, args.num_users)
        else:
            exit('Error: only consider IID setting in CIFAR10')	
            # 代码没有实现CIFAR的非独立同分布
    else:
        exit('Error: unrecognized dataset')
    img_size = dataset_train[0][0].shape	# 输出数据集的图像格式

    # build model	# 调用模型
    if args.model == 'cnn' and args.dataset == 'cifar':
        net_glob = CNNCifar(args=args).to(args.device)
    elif args.model == 'cnn' and args.dataset == 'mnist':
        net_glob = CNNMnist(args=args).to(args.device)
    elif args.model == 'mlp':
        len_in = 1
        for x in img_size:
            len_in *= x
        net_glob = MLP(dim_in=len_in, dim_hidden=200, dim_out=args.num_classes).to(args.device)
    else:
        exit('Error: unrecognized model')
    print(net_glob)	# 输出网络的前向函数的返回值，也就是一个张量
    net_glob.train()	# 设置为训练模式

    # copy weights
    w_glob = net_glob.state_dict()	# 取出模型参数即权重

    # 训练
    loss_train = []
    cv_loss, cv_acc = [], []
    val_loss_pre, counter = 0, 0
    net_best = None
    best_loss = None
    val_acc_list, net_list = [], []

    if args.all_clients: 	# 如果命令行参数指定了全部客户端
        print("Aggregation over all clients")
        w_locals = [w_glob for i in range(args.num_users)]	# 聚合所有的本地模型参数
    for iter in range(args.epochs):
        loss_locals = []
        if not args.all_clients:
            w_locals = []
            
        m = max(int(args.frac * args.num_users), 1)	
        # 计算每轮训练中参与训练的用户的数量
        idxs_users = np.random.choice(range(args.num_users), m, replace=False)
        # 随机选择m个不重复的用户索引
        for idx in idxs_users:
            local = LocalUpdate(args=args, dataset=dataset_train, idxs=dict_users[idx])
            # 实例化LocalUpdate类
            w, loss = local.train(net=copy.deepcopy(net_glob).to(args.device))
            # 调用LocalUpdate类的函数train 使用深拷贝取出刚才的模型 这里就因为每个客户端都需要一个单独的模型来训练 所以必须使用深拷贝
            if args.all_clients:
                w_locals[idx] = copy.deepcopy(w)
            else:
                w_locals.append(copy.deepcopy(w))
            loss_locals.append(copy.deepcopy(loss))	# 计算本地损失
        # update global weights
        w_glob = FedAvg(w_locals)	# 使用FedAvg算法聚合模型参数

        # copy weight to net_glob	
        net_glob.load_state_dict(w_glob)	# 加载全局模型

        # print loss
        loss_avg = sum(loss_locals) / len(loss_locals)	# 输出平均损失
        print('Round {:3d}, Average loss {:.3f}'.format(iter, loss_avg))
        loss_train.append(loss_avg)

    # plot loss curve
    plt.figure()	#使用plt可视化损失变化
    plt.plot(range(len(loss_train)), loss_train)
    plt.ylabel('train_loss')
    plt.savefig('./save/fed_{}_{}_{}_C{}_iid{}.png'.format(args.dataset, args.model, args.epochs, args.frac, args.iid))

    # testing
    net_glob.eval()	# 测试模型 设置为eval评估模式
    acc_train, loss_train = test_img(net_glob, dataset_train, args) # 调用test_img测试函数
    acc_test, loss_test = test_img(net_glob, dataset_test, args)
    print("Training accuracy: {:.2f}".format(acc_train))
    print("Testing accuracy: {:.2f}".format(acc_test))

main_nn.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Python version: 3.6

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
import torch.optim as optim
from torchvision import datasets, transforms

from utils.options import args_parser
from models.Nets import MLP, CNNMnist, CNNCifar


def test(net_g, data_loader):	
    # 计算测试时的平均损失和准确率 实际上和test.py中的test_img代码基本相同
    # testing
    net_g.eval()
    test_loss = 0
    correct = 0
    l = len(data_loader)
    for idx, (data, target) in enumerate(data_loader):
        data, target = data.to(args.device), target.to(args.device)
        log_probs = net_g(data)
        test_loss += F.cross_entropy(log_probs, target).item()
        y_pred = log_probs.data.max(1, keepdim=True)[1]
        correct += y_pred.eq(target.data.view_as(y_pred)).long().cpu().sum()

    test_loss /= len(data_loader.dataset)
    print('\nTest set: Average loss: {:.4f} \nAccuracy: {}/{} ({:.2f}%)\n'.format(
        test_loss, correct, len(data_loader.dataset),
        100. * correct / len(data_loader.dataset)))

    return correct, test_loss


if __name__ == '__main__':
    # parse args
    args = args_parser()	# 解析命令行参数
    args.device = torch.device('cuda:{}'.format(args.gpu) if torch.cuda.is_available() and args.gpu != -1 else 'cpu')

    torch.manual_seed(args.seed)

    # load dataset and split users	# 这一部分加载数据和main_fed一致
    if args.dataset == 'mnist':
        dataset_train = datasets.MNIST('./data/mnist/', train=True, download=True,
                   transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,))
                   ]))
        img_size = dataset_train[0][0].shape
    elif args.dataset == 'cifar':
        transform = transforms.Compose(
            [transforms.ToTensor(),
             transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
        dataset_train = datasets.CIFAR10('./data/cifar', train=True, transform=transform, target_transform=None, download=True)
        img_size = dataset_train[0][0].shape
    else:
        exit('Error: unrecognized dataset')

    # build model
    if args.model == 'cnn' and args.dataset == 'cifar':
        net_glob = CNNCifar(args=args).to(args.device)
    elif args.model == 'cnn' and args.dataset == 'mnist':
        net_glob = CNNMnist(args=args).to(args.device)
    elif args.model == 'mlp':
        len_in = 1
        for x in img_size:
            len_in *= x
        net_glob = MLP(dim_in=len_in, dim_hidden=64, dim_out=args.num_classes).to(args.device)
    else:
        exit('Error: unrecognized model')
    print(net_glob)

    # training	# 训练过程 直接用整个数据集进行训练
    optimizer = optim.SGD(net_glob.parameters(), lr=args.lr, momentum=args.momentum)
    train_loader = DataLoader(dataset_train, batch_size=64, shuffle=True)

    list_loss = []
    net_glob.train()
    for epoch in range(args.epochs):
        batch_loss = []
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(args.device), target.to(args.device)
            optimizer.zero_grad()
            output = net_glob(data)
            loss = F.cross_entropy(output, target)
            loss.backward()
            optimizer.step()
            if batch_idx % 50 == 0:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                    epoch, batch_idx * len(data), len(train_loader.dataset),
                           100. * batch_idx / len(train_loader), loss.item()))
            batch_loss.append(loss.item())
        loss_avg = sum(batch_loss)/len(batch_loss)
        print('\nTrain loss:', loss_avg)
        list_loss.append(loss_avg)

    # plot loss
    plt.figure()
    plt.plot(range(len(list_loss)), list_loss)
    plt.xlabel('epochs')
    plt.ylabel('train loss')
    plt.savefig('./save/nn_{}_{}_{}.png'.format(args.dataset, args.model, args.epochs))

    # testing	# 调用上面的test函数进行测试
    if args.dataset == 'mnist':
        dataset_test = datasets.MNIST('./data/mnist/', train=False, download=True,
                   transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,))
                   ]))
        test_loader = DataLoader(dataset_test, batch_size=1000, shuffle=False)
    elif args.dataset == 'cifar':
        transform = transforms.Compose(
            [transforms.ToTensor(),
             transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
        dataset_test = datasets.CIFAR10('./data/cifar', train=False, transform=transform, target_transform=None, download=True)
        test_loader = DataLoader(dataset_test, batch_size=1000, shuffle=False)
    else:
        exit('Error: unrecognized dataset')

    print('test on', len(dataset_test), 'samples')
    test_acc, test_loss = test(net_glob, test_loader)

我们可以看到main_nn就是正常的一个训练过程，main_fed是联邦学习的训练测试过程。代码部分讲解就到这里，我们给出一个自己画的main_fed.py代码流程图。