联邦学习框架Federated-Learning-PyTorch 研读报告

框架来源 GitHub - AshwinRJ/Federated-Learning-PyTorch: Implementation of Communication-Efficient Learning of Deep Networks from Decentralized Data

baseline_main.py

从终端接收命令行参数

从终端接收命令行参数，封装后只有一行，如果想了解细节可以参考我的另一篇博客。

args = args_parser()

Python中是如何接收在终端中输入（自定义的）命令行参数的？_MikingG的博客-CSDN博客读取在终端中输入的命令行参数的背后原理及简单使用教程https://blog.csdn.net/weixin_64123373/article/details/132246130

决定运行设备

根据命令行传入参数觉得是在GPU还是在CPU上运行。

    # 决定代码将在 CPU 还是 GPU 上运行
    if args.gpu:
        torch.cuda.set_device(args.gpu)
    device = 'cuda' if args.gpu else 'cpu'

加载数据集

加载命令行中指定的数据集，封装后只有一行，如果想了解细节可以参考我的另一篇博客。

train_dataset, test_dataset, _ = get_dataset(args)

https://blog.csdn.net/weixin_64123373/article/details/132247146https://blog.csdn.net/weixin_64123373/article/details/132247146

建立模型

    # BUILD MODEL
    if args.model == 'cnn':
        # Convolutional neural netork
        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':
        # Multi-layer preceptron
        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')

作者在model.py中实现了CNNMnist、CNNFashion_Mnist、CNNCifar、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)


class CNNMnist(nn.Module):
    def __init__(self, args):
        super(CNNMnist, self).__init__()
        self.conv1 = nn.Conv2d(args.num_channels, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, args.num_classes)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, x.shape[1]*x.shape[2]*x.shape[3])
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)


class CNNFashion_Mnist(nn.Module):
    def __init__(self, args):
        super(CNNFashion_Mnist, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=5, padding=2),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(2))
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=5, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2))
        self.fc = nn.Linear(7*7*32, 10)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        return out


class CNNCifar(nn.Module):
    def __init__(self, args):
        super(CNNCifar, 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, args.num_classes)

    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 F.log_softmax(x, dim=1)

MLP讲解https://blog.csdn.net/weixin_64123373/article/details/132253020CNNMnist讲解https://blog.csdn.net/weixin_64123373/article/details/132253084CNNFashion_Mnist讲解https://blog.csdn.net/weixin_64123373/article/details/132253110CNNCifar讲解https://blog.csdn.net/weixin_64123373/article/details/132253132

设置优化器和标准

    # Training
    # Set optimizer and criterion
    if args.optimizer == 'sgd':
        optimizer = torch.optim.SGD(global_model.parameters(), lr=args.lr,
                                    momentum=0.5)
    elif args.optimizer == 'adam':
        optimizer = torch.optim.Adam(global_model.parameters(), lr=args.lr,
                                     weight_decay=1e-4)

随机梯度下降（Stochastic Gradient Descent, SGD）优化器可以用于训练神经网络模型global_model。通过适当选择学习率和动量，可以有效地调整训练过程，帮助模型快速和稳定地收敛到良好的解。

• torch.optim.SGD: 这是PyTorch库中实现SGD的类。SGD是最常用的优化算法之一，特别是在训练深度学习模型时。

• global_model.parameters(): 这里，global_model是你想要优化的模型。通过调用.parameters()方法，优化器可以知道哪些权重和偏置需要更新。

• lr=args.lr: 这里的lr是学习率（learning rate），它决定了模型在训练过程中参数更新的步长大小。较大的学习率可能会导致训练快速收敛，但也可能会导致模型在最优解附近震荡。较小的学习率可能会使收敛更稳定，但训练可能会较慢。args.lr是从函数外部传递的参数，用于设置学习率的具体值。

• momentum=0.5: 动量（momentum）是SGD的一个变体，它考虑了过去的梯度，以便更平滑地更新权重。这可以提高优化过程的速度和稳定性。动量值通常在0到1之间，本例中设置为0.5。

Adam优化器可以用于训练神经网络模型global_model。Adam是一种流行的优化算法，结合了Adagrad和RMSProp的优点，通常可以更快地收敛。

• torch.optim.Adam: 这是PyTorch库中实现Adam优化算法的类。

• global_model.parameters(): 这里的global_model是你想要优化的模型。通过调用.parameters()方法，优化器可以知道哪些权重和偏置需要更新。

• lr=args.lr: 这里的lr是学习率（learning rate），和SGD的学习率概念相同。它决定了参数更新的步长大小。args.lr是从函数外部传递的参数，用于设置学习率的具体值。

• weight_decay=1e-4: 这是权重衰减参数，用于L2正则化。正则化是一种防止模型过拟合的技术。通过在优化过程中对模型权重添加一些小的惩罚项，可以让模型变得更“简单”，从而有助于泛化到未见过的数据。1e-4是权重衰减的具体值，表示这一惩罚项的大小。

    trainloader = DataLoader(train_dataset, batch_size=64, shuffle=True)
    criterion = torch.nn.NLLLoss().to(device)
    epoch_loss = []

1. trainloader = DataLoader(train_dataset, batch_size=64, shuffle=True):

   • DataLoader: 这是PyTorch库中的一个类，用于封装数据集并提供批量加载的迭代器，以便在训练模型时按批次获取数据。
   • train_dataset: 这是要用于训练的数据集。
   • batch_size=64: 每个批次包括64个样本。批量训练是训练深度学习模型的标准做法，可以更有效地利用硬件并加快训练速度。
   • shuffle=True: 在每个训练周期开始时，数据会被随机重新排序。这有助于确保模型不会在训练数据的顺序上过拟合，从而提高泛化能力。

2. criterion = torch.nn.NLLLoss().to(device):

   • torch.nn.NLLLoss(): 这是PyTorch库中用于计算负对数似然损失（Negative Log-Likelihood Loss）的类。NLL损失通常与用于多分类问题的softmax激活函数结合使用。
   • .to(device): 这将损失函数转移到指定的设备（例如GPU或CPU）上。在进行训练之前，确保模型和损失函数都在同一个设备上计算是很重要的。

3. epoch_loss = []: 这一行初始化了一个空列表，用于存储每个训练周期（epoch）的损失值。在训练多个周期时，通过跟踪每个周期的损失，可以监视模型的学习进度，并在必要时调整超参数。

2.2补充如果模型和损失函数不在同一个设备上，那么在计算过程中可能会产生不必要的数据传输，例如在GPU和CPU之间来回移动数据。这种数据传输不仅可能导致代码错误，而且还可能显著降低计算效率。

通过使用例如.to(device)这样的方法可以明确地将模型、损失函数和数据移动到同一个设备上。这样做的目的是确保在执行前向传播、计算损失和进行反向传播时，所有的计算都在同一个设备上进行，从而实现高效的训练过程。

核心的训练

    for epoch in tqdm(range(args.epochs)):
        batch_loss = []

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

            optimizer.zero_grad()
            outputs = global_model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

1. for epoch in tqdm(range(args.epochs)):: 外部循环，将循环args.epochs次。每次循环代表一个训练周期(epoch)，其中整个训练集都会被遍历一次。tqdm是一个用于在控制台上显示进度条的库，它会显示训练进度。

2. batch_loss = []:: 初始化一个空列表，用于存储每个批次的损失值。这可以用于在每个训练周期结束时分析或可视化损失的下降趋势。

3. 内部循环，遍历训练数据的批次：

   • (images, labels) in enumerate(trainloader):: 从trainloader中迭代获取图像和标签的批次。
   • images, labels = images.to(device), labels.to(device):: 将图像和标签移动到执行计算的设备上（例如GPU或CPU）。

   • optimizer.zero_grad(): 清零梯度

     • 在每个训练步骤开始时，我们需要手动将梯度设置为零。如果不这样做，梯度会在每个训练步骤中累积，而不是替换，这可能导致意外的训练行为。

   • outputs = global_model(images): 前向传播

     • 这里，global_model是我们要训练的模型，而images是输入的一批图像数据。
     • 我们通过将输入数据传递给模型来进行前向传播，计算模型的预测输出。

   • loss = criterion(outputs, labels): 计算损失

     • 我们使用定义的损失函数criterion来计算模型的预测输出outputs与真实标签labels之间的损失。
     • 损失函数度量模型预测与实际目标之间的差距，是我们想要最小化的东西。

   • loss.backward(): 反向传播（强大的Pytorch自带）

     • 调用backward()函数将计算损失相对于模型参数的梯度。这些梯度存储在模型参数的.grad属性中，用于更新模型权重。
     • 反向传播是一种通过网络反向传播误差的过程，通过计算梯度来了解每个参数对最终误差的贡献大小。

   • optimizer.step(): 更新权重（强大的Pytorch自带）

     • 最后，我们调用优化器的step方法来更新模型的权重。
     • 这一步基于先前计算的梯度来调整模型的内部参数，如权重和偏置，以便在下一个迭代中减小损失。

监控训练过程

1. 条件检查：

   if batch_idx % 50 == 0:

   这部分检查当前的批次索引（batch_idx）是否是50的倍数。这意味着每50个批次，以下代码将执行一次。

2. 打印训练进度：

   print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
    epoch+1, batch_idx * len(images), len(trainloader.dataset),
    100. * batch_idx / len(trainloader), loss.item()))

   这部分使用字符串格式化来打印以下信息：

   • 当前的epoch（epoch + 1）。
   • 已处理的图像数量（batch_idx * len(images)）。
   • 总的图像数量（len(trainloader.dataset)）。
   • 训练进度的百分比。
   • 当前批次的损失（loss.item()）。

3. 收集批次损失：

   batch_loss.append(loss.item())

   这行代码将当前批次的损失添加到batch_loss列表中，该列表收集整个epoch的批次损失。

4. 计算和打印平均损失：

   loss_avg = sum(batch_loss)/len(batch_loss) 
   print('\nTrain loss:', loss_avg) 
   epoch_loss.append(loss_avg)

   这部分首先计算整个epoch的平均损失，然后将其打印出来。最后，将平均损失添加到epoch_loss列表中，该列表收集所有epoch的平均损失。

完整代码

            if batch_idx % 50 == 0:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                    epoch+1, batch_idx * len(images), len(trainloader.dataset),
                    100. * batch_idx / len(trainloader), loss.item()))
            batch_loss.append(loss.item())

        loss_avg = sum(batch_loss)/len(batch_loss)
        print('\nTrain loss:', loss_avg)
        epoch_loss.append(loss_avg)

画图

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

这段代码是用来绘制模型训练过程中的损失曲线，并将其保存为图像文件。具体来说：

1. plt.figure(): 这个函数调用创建了一个新的图像窗口，让你可以开始在其上绘制图形。

2. plt.plot(range(len(epoch_loss)), epoch_loss):

   • 这行代码绘制了损失曲线。range(len(epoch_loss)) 创建了一个序列，表示每个训练周期（epoch），而 epoch_loss 是一个包含每个周期损失值的列表。这行代码绘制了一个折线图，横坐标表示周期数，纵坐标表示对应的损失值。

3. plt.xlabel('epochs') 和 plt.ylabel('Train loss'):

   • 这两行代码分别设置了图像的x轴和y轴的标签。'epochs' 代表训练周期数，'Train loss' 代表训练损失。

4. plt.savefig('../save/nn_{}_{}_{}.png'.format(args.dataset, args.model, args.epochs)):

   • 这行代码将绘制的图像保存为一个PNG文件。
   • '../save/nn_{}_{}_{}.png' 是保存文件的路径和名称，其中包括三个占位符 {}，它们将被 .format(args.dataset, args.model, args.epochs) 中的相应值替换。这样，每个图像文件名都将包括所使用的数据集名称、模型名称和训练周期数，从而方便区分不同的训练实验。

测试

    # testing
    test_acc, test_loss = test_inference(args, global_model, test_dataset)
    print('Test on', len(test_dataset), 'samples')
    print("Test Accuracy: {:.2f}%".format(100*test_acc))

这段代码是执行模型在测试集上的推断，并打印测试准确率。详细解释如下：

1. test_acc, test_loss = test_inference(args, global_model, test_dataset):

   • 调用test_inference函数来进行模型的测试推断，在update.py中定义。
   • args可能包括了一些测试参数或模型的超参数。
   • global_model是要评估的训练过的模型。
   • test_dataset是用于测试模型的数据集。
   • 函数返回测试准确率test_acc和测试损失test_loss。

2. print('Test on', len(test_dataset), 'samples'):

   • 打印一条消息，显示测试集中的样本数量。len(test_dataset)计算测试集中的样本数。

3. print("Test Accuracy: {:.2f}%".format(100*test_acc)):

   • 打印测试准确率。{:.2f}是一个格式说明符，表示将准确率格式化为保留两位小数的浮点数。
   • 通过将准确率test_acc乘以100，将其从范围0到1转换为百分比形式。

update.py中的test_inference(args, model, test_dataset)

在这里可以解开loss为负的“谜团”。

完整代码

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

baseline.py完整代码

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


from tqdm import tqdm
import matplotlib.pyplot as plt

import torch
from torch.utils.data import DataLoader

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


if __name__ == '__main__':
    # 从命令行中接受参数
    args = args_parser()
    # 决定代码将在 CPU 还是 GPU 上运行
    if args.gpu:
        torch.cuda.set_device(args.gpu)
    device = 'cuda' if args.gpu else 'cpu'

    # load datasets
    train_dataset, test_dataset, _ = get_dataset(args)

    # BUILD MODEL
    if args.model == 'cnn':
        # Convolutional neural netork
        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':
        # Multi-layer preceptron
        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')

    # Set the model to train and send it to device.
    global_model.to(device)
    global_model.train()
    print(global_model)

    # Training
    # Set optimizer and criterion
    if args.optimizer == 'sgd':
        optimizer = torch.optim.SGD(global_model.parameters(), lr=args.lr,
                                    momentum=0.5)
    elif args.optimizer == 'adam':
        optimizer = torch.optim.Adam(global_model.parameters(), lr=args.lr,
                                     weight_decay=1e-4)

    trainloader = DataLoader(train_dataset, batch_size=64, shuffle=True)
    criterion = torch.nn.NLLLoss().to(device)
    epoch_loss = []

    for epoch in tqdm(range(args.epochs)):
        batch_loss = []

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

            optimizer.zero_grad()
            outputs = global_model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            if batch_idx % 50 == 0:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                    epoch+1, batch_idx * len(images), len(trainloader.dataset),
                    100. * batch_idx / len(trainloader), loss.item()))
            batch_loss.append(loss.item())

        loss_avg = sum(batch_loss)/len(batch_loss)
        print('\nTrain loss:', loss_avg)
        epoch_loss.append(loss_avg)

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

    # testing
    test_acc, test_loss = test_inference(args, global_model, test_dataset)
    print('Test on', len(test_dataset), 'samples')
    print("Test Accuracy: {:.2f}%".format(100*test_acc))

---

如果本篇研读报告对您有帮助，欢迎点赞，收藏，评论，转发~