FedProx代码详解

yued0997

已于 2023-11-07 17:30:39 修改

阅读量1k

点赞数 3

文章标签： python

于 2023-09-27 23:05:38 首次发布

本文链接：https://blog.csdn.net/qq_56483157/article/details/133364495

版权

文献：https://arxiv.org/abs/1812.06127

文献笔记：[论文阅读]Federated Optimization In Heterogeneous Networks_荒诞主义的博客-CSDN博客

代码来源：https://github.com/ki-ljl/FedProx-PyTorch

逐行分析代码记录所学

一、args.py

--用于加载参数

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

    parser.add_argument('--E', type=int, default=30, help='number of rounds of training')
    parser.add_argument('--r', type=int, default=30, help='number of communication rounds')
    parser.add_argument('--K', type=int, default=10, help='number of total clients')
    parser.add_argument('--input_dim', type=int, default=28, help='input dimension')
    parser.add_argument('--lr', type=float, default=0.01, help='learning rate')
    parser.add_argument('--C', type=float, default=0.5, help='sampling rate')
    parser.add_argument('--B', type=int, default=50, help='local batch size')
    parser.add_argument('--mu', type=float, default=0.01, help='proximal term constant')
    parser.add_argument('--optimizer', type=str, default='adam', help='type of optimizer')
    parser.add_argument('--device', default=torch.device("cuda" if torch.cuda.is_available() else "cpu"))
    parser.add_argument('--weight_decay', type=float, default=1e-4, help='weight decay')
    parser.add_argument('--step_size', type=int, default=10, help='step size')
    parser.add_argument('--gamma', type=float, default=0.1, help='learning rate decay per global round')
    clients = ['Task1_W_Zone' + str(i) for i in range(1, 11)]
    parser.add_argument('--clients', default=clients)

    args = parser.parse_args()

    return args

argparse的用法详见上篇帖子：FedAvg代码详解_荒诞主义的博客-CSDN博客

–E: 30, 训练轮数
–r: 30, 通信轮数
–K: 10, 客户端总数
–input_dim: 28, 输入维度
–lr: 0.01, 学习率
–C: 0.5, 采样率
–B: 50, 本地批大小
–mu: 0.01, 近端项常数
–optimizer: ‘adam’, 优化器类型
–device: torch.device(“cuda” if torch.cuda.is_available() else “cpu”), 指定设备
–weight_decay: 1e-4, 权重衰减
–step_size: 10, 步长
–gamma: 0.1, 每个全局轮次的学习率衰减
–clients: [‘Task1_W_Zone1’, ‘Task1_W_Zone2’, ‘Task1_W_Zone3’, ‘Task1_W_Zone4’, ‘Task1_W_Zone5’, ‘Task1_W_Zone6’, ‘Task1_W_Zone7’, ‘Task1_W_Zone8’, ‘Task1_W_Zone9’, ‘Task1_W_Zone10’], 客户端列表

二、get_data.py

args = args_parser()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
clients_wind = ['Task1_W_Zone' + str(i) for i in range(1, 11)]


def load_data(file_name):
    df = pd.read_csv('data/Wind/Task 1/Task1_W_Zone1_10/' + file_name + '.csv', encoding='gbk')
    columns = df.columns
    df.fillna(df.mean(), inplace=True)
    for i in range(3, 7):
        MAX = np.max(df[columns[i]])
        MIN = np.min(df[columns[i]])
        df[columns[i]] = (df[columns[i]] - MIN) / (MAX - MIN)

    return df


class MyDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __getitem__(self, item):
        return self.data[item]

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


def nn_seq_wind(file_name, B):
    print('data processing...')
    dataset = load_data(file_name)
    # split
    train = dataset[:int(len(dataset) * 0.6)]
    val = dataset[int(len(dataset) * 0.6):int(len(dataset) * 0.8)]
    test = dataset[int(len(dataset) * 0.8):len(dataset)]

    def process(data):
        columns = data.columns
        wind = data[columns[2]]
        wind = wind.tolist()
        data = data.values.tolist()
        seq = []
        for i in range(len(data) - 30):
            train_seq = []
            train_label = []
            for j in range(i, i + 24):
                train_seq.append(wind[j])
            for c in range(3, 7):
                train_seq.append(data[i + 24][c])
            train_label.append(wind[i + 24])
            train_seq = torch.FloatTensor(train_seq).view(-1)
            train_label = torch.FloatTensor(train_label).view(-1)
            seq.append((train_seq, train_label))

        seq = MyDataset(seq)

        seq = DataLoader(dataset=seq, batch_size=B, shuffle=False, num_workers=0)

        return seq

    Dtr = process(train)
    Val = process(val)
    Dte = process(test)

    return Dtr, Val, Dte


def get_mape(x, y):
    """
    :param x:true
    :param y:pred
    :return:MAPE
    """
    return np.mean(np.abs((x - y) / x))

逐行分析：

args = args_parser()

用于解析输入的参数

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

根据系统可用的GPU决定训练使用的设备类型：

如果有可用的GPU，就使用CUDA设备；否则使用CPU。

clients_wind = ['Task1_W_Zone' + str(i) for i in range(1, 11)]

创建一个包含10个任务的客户端列表：

这里使用range(1, 11)生成数字1到10，并将其转换为字符串并与Task1_W_Zone组合，最终得到一个包含10个任务名称的列表。

def load_data(file_name):
    df = pd.read_csv('data/Wind/Task 1/Task1_W_Zone1_10/' + file_name + '.csv', encoding='gbk')
    columns = df.columns
    df.fillna(df.mean(), inplace=True)
    for i in range(3, 7):
        MAX = np.max(df[columns[i]])
        MIN = np.min(df[columns[i]])
        df[columns[i]] = (df[columns[i]] - MIN) / (MAX - MIN)

    return df

load_data的函数，用于加载数据：

在load_data(file_name)函数中，通过pd.read_csv()函数读取CSV文件的数据并创建了一个DataFrame对象df。文件路径为data/Wind/Task 1/Task1_W_Zone1_10/加上参数file_name再加上.csv。函数中的encoding='gbk'指定了文件的编码格式为GBK。
获取DataFrame对象df的所有列名，并将其赋值给变量columns。
用DataFrame对象df的均值填充其中的缺失值。fillna()函数用于填充缺失值，df.mean()返回每列的均值，inplace=True表示在原地进行填充操作，即修改原始的DataFrame对象df。
用于对数据的指定列进行归一化处理。range(3, 7)生成一个从3到7（不包括7）的整数序列，循环变量为i。在每次循环中，首先使用np.max()和np.min()分别获取指定列的最大值和最小值，然后将其保存到变量MAX和MIN中。接下来，通过计算(df[columns[i]] - MIN) / (MAX - MIN)将指定列的每个元素进行归一化处理，结果保存回原始的DataFrame对象df的相应列中。
最后，函数返回处理后的DataFrame对象df。

整个函数的目的是加载CSV文件，并对其中的数据进行缺失值填充和归一化处理。

class MyDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __getitem__(self, item):
        return self.data[item]

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

定义了一个名为MyDataset的类，它继承自torch.utils.data.Dataset类。通过继承Dataset类，我们可以创建自定义的数据集类，用于加载和处理数据。
__init__()方法，是MyDataset类的构造函数。它接收一个名为data的参数，并将其保存为类的属性self.data。构造函数用于初始化类的实例，并设置属性的初始值。来初始化数据，
__getitem__()方法。它定义了如何通过索引获取数据集中的样本。方法接收一个名为item的参数，该参数表示样本的索引。方法通过索引访问self.data属性中的数据，并将其返回。这样，我们可以使用dataset[i]的方式获取数据集中索引为i的样本。
__len__()方法。它定义了数据集的长度，也就是数据集中样本的数量。方法返回self.data属性的长度，即数据集中样本的个数。

def nn_seq_wind(file_name, B):
    print('data processing...')
    dataset = load_data(file_name)
    train = dataset[:int(len(dataset) * 0.6)]
    val = dataset[int(len(dataset) * 0.6):int(len(dataset) * 0.8)]
    test = dataset[int(len(dataset) * 0.8):len(dataset)]

    def process(data):
        columns = data.columns
        wind = data[columns[2]]
        wind = wind.tolist()
        data = data.values.tolist()
        seq = []
        for i in range(len(data) - 30):
            train_seq = []
            train_label = []
            for j in range(i, i + 24):
                train_seq.append(wind[j])
            for c in range(3, 7):
                train_seq.append(data[i + 24][c])
            train_label.append(wind[i + 24])
            train_seq = torch.FloatTensor(train_seq).view(-1)
            train_label = torch.FloatTensor(train_label).view(-1)
            seq.append((train_seq, train_label))

        seq = MyDataset(seq)
        seq = DataLoader(dataset=seq, batch_size=B, shuffle=False, num_workers=0)

        return seq

    Dtr = process(train)
    Val = process(val)
    Dte = process(test)

    return Dtr, Val, Dte

逐行分析：

    print('data processing...')
    dataset = load_data(file_name)

输出一个提示信息，并调用load_data(file_name)函数加载数据。加载的数据存储在dataset变量中。

    train = dataset[:int(len(dataset) * 0.6)]
    val = dataset[int(len(dataset) * 0.6):int(len(dataset) * 0.8)]
    test = dataset[int(len(dataset) * 0.8):len(dataset)]

这几行代码将数据集dataset划分为训练集（train）、验证集（val）和测试集（test）。通过切片操作，根据数据集的长度将数据按照60%、20%和20%的比例进行划分。

    def process(data):
        columns = data.columns
        wind = data[columns[2]]
        wind = wind.tolist()
        data = data.values.tolist()
        seq = []
        for i in range(len(data) - 30):
            train_seq = []
            train_label = []
            for j in range(i, i + 24):
                train_seq.append(wind[j])
            for c in range(3, 7):
                train_seq.append(data[i + 24][c])
            train_label.append(wind[i + 24])
            train_seq = torch.FloatTensor(train_seq).view(-1)
            train_label = torch.FloatTensor(train_label).view(-1)
            seq.append((train_seq, train_label))

        seq = MyDataset(seq)
        seq = DataLoader(dataset=seq, batch_size=B, shuffle=False, num_workers=0)

        return seq

这段代码是一个函数 process(data)，该函数接受一个参数 data，并返回一个 DataLoader 对象。

首先，代码提取了 data 的列名，并将第三列的数据赋值给变量 wind。然后将 wind 转换为列表，并重新赋值给 wind 变量。接下来，将 data 转换为列表格式。

代码创建了一个空列表 seq 用于存储训练序列和标签。

然后通过循环遍历数据的长度减去30的范围，创建 train_seq 和 train_label 列表。

在内部的循环中，将 wind 列表从索引 i 到 i + 24 的值添加到 train_seq 中，然后将 data 列表从索引 i + 24 的第3到第6个元素添加到 train_seq 中。

同时，将 wind 列表中索引为 i + 24 的值添加到 train_label 中。

接下来，使用 torch.FloatTensor 将 train_seq 和 train_label 转换为 PyTorch 的 FloatTensor 类型，并进行形状变换为一维。

然后将 (train_seq, train_label) 元组添加到 seq 列表中。

接下来，将 seq 列表传递给 MyDataset 类创建一个数据集对象，并重新将其赋值给 seq 变量。

最后，根据给定的参数值创建一个 DataLoader 对象，将数据集对象 seq 作为数据集，B 作为批大小，shuffle 设置为 False，num_workers 设置为 0。

最后，返回 seq 对象作为函数的输出。

train_seq = []
            train_label = []
            for j in range(i, i + 24):
                train_seq.append(wind[j])
            for c in range(3, 7):
                train_seq.append(data[i + 24][c])
            train_label.append(wind[i + 24])

对于每个训练序列，将前24个数据作为训练特征（客户端的wind数据），后续4个数据作为训练标签。将特征和标签分别存储在train_seq和train_label中。

train_seq = torch.FloatTensor(train_seq).view(-1)
            train_label = torch.FloatTensor(train_label).view(-1)

将train_seq和train_label转换为FloatTensor格式，并调整其维度为一维。这样做是为了适应后续的模型训练。

seq.append((train_seq, train_label))：将训练序列和对应的标签作为元组，添加到seq列表中。
seq = MyDataset(seq)：创建一个自定义的数据集MyDataset，并传入seq列表作为参数，用于后续的数据加载。
seq = DataLoader(dataset=seq, batch_size=B, shuffle=False, num_workers=0)：使用DataLoader对数据集进行批量加载，设置批量大小为B，并且在联邦学习中通常会进行数据洗牌以保证随机性，设置shuffle=True。
return seq：返回加载好的包含训练序列和标签的数据集seq。

    Dtr = process(train)
    Val = process(val)
    Dte = process(test)

    return Dtr, Val, Dte

这段代码的作用是调用process()函数分别处理训练集、验证集和测试集，并将处理后的数据加载器保存在Dtr、Val和Dte变量中。最后，通过return语句将这三个数据加载器作为结果返回。

def get_mape(x, y):
    return np.mean(np.abs((x - y) / x))

这几行代码分别通过调用process()函数对训练集、验证集和测试集进行处理，将处理后的数据加载器保存在Dtr、Val和Dte变量中。最后，将这三个数据加载器作为结果返回。

三、model.py

class ANN(nn.Module):
    def __init__(self, args, name):
        super(ANN, self).__init__()
        self.name = name
        self.len = 0
        self.loss = 0
        self.fc1 = nn.Linear(args.input_dim, 20)
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()
        self.dropout = nn.Dropout()
        self.fc2 = nn.Linear(20, 20)
        self.fc3 = nn.Linear(20, 20)
        self.fc4 = nn.Linear(20, 1)

    def forward(self, data):
        x = self.fc1(data)
        x = self.sigmoid(x)
        x = self.fc2(x)
        x = self.sigmoid(x)
        x = self.fc3(x)
        x = self.sigmoid(x)
        x = self.fc4(x)
        x = self.sigmoid(x)

        return x

这个ANN类实现了一个简单的前向神经网络模型，具有四个全连接层和三个Sigmoid激活函数层。每个线性层之后都应用了Sigmoid激活函数。模型的输入维度由args.input_dim确定，输出维度为1。在实际使用中，可以根据具体任务的要求进行修改和调整。

逐行分析：

class ANN(nn.Module):
    def __init__(self, args, name):
        super(ANN, self).__init__()
        self.name = name
        self.len = 0
        self.loss = 0

这是ANN类的构造函数。它接收两个参数：args和name。args是一个参数对象，name是一个名称。在构造函数中，首先调用nn.Module类的构造函数进行初始化，然后将name保存为类的属性self.name，并初始化self.len和self.loss为0。构造函数用于初始化类的实例，并设置属性的初始值。

        self.fc1 = nn.Linear(args.input_dim, 20)
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()
        self.dropout = nn.Dropout()
        self.fc2 = nn.Linear(20, 20)
        self.fc3 = nn.Linear(20, 20)
        self.fc4 = nn.Linear(20, 1)

这几行代码定义了神经网络的各层。nn.Linear()用于定义一个线性层，第一个参数是输入特征的维度，第二个参数是输出特征的维度。nn.ReLU()、nn.Sigmoid()和nn.Dropout()分别定义了激活函数ReLU、Sigmoid和Dropout层。通过定义这些层，我们可以构建神经网络模型的结构。

    def forward(self, data):
        x = self.fc1(data)
        x = self.sigmoid(x)
        x = self.fc2(x)
        x = self.sigmoid(x)
        x = self.fc3(x)
        x = self.sigmoid(x)
        x = self.fc4(x)
        x = self.sigmoid(x)

        return x

这是ANN类的forward()方法，它定义了数据在前向传播过程中的运算流程。给定输入数据data，首先将其传入第一个全连接层self.fc1，然后对输出进行Sigmoid激活函数操作。接着将结果传入第二个全连接层self.fc2，再次进行Sigmoid激活函数操作。依次类推，将输入依次传入self.fc3和self.fc4层，并进行相应的Sigmoid激活函数操作。最后，将结果返回。

通过定义这个前向传播方法，我们可以根据输入数据的流经神经网络的顺序计算输出结果。

四、客户端与服务器

4.1 server.py

class FedProx:
    def __init__(self, args):
        self.args = args
        self.nn = ANN(args=self.args, name='server').to(args.device)
        self.nns = []
        for i in range(self.args.K):
            temp = copy.deepcopy(self.nn)
            temp.name = self.args.clients[i]
            self.nns.append(temp)

    def server(self):
        for t in tqdm(range(self.args.r)):
            print('round', t + 1, ':')
            # sampling
            m = np.max([int(self.args.C * self.args.K), 1])
            index = random.sample(range(0, self.args.K), m)  # st
            # dispatch
            self.dispatch(index)
            # local updating
            self.client_update(index)
            # aggregation
            self.aggregation(index)

        return self.nn

    def aggregation(self, index):
        s = 0
        for j in index:
            # normal
            s += self.nns[j].len

        params = {}
        for k, v in self.nns[0].named_parameters():
            params[k] = torch.zeros_like(v.data)

        for j in index:
            for k, v in self.nns[j].named_parameters():
                params[k] += v.data * (self.nns[j].len / s)

        for k, v in self.nn.named_parameters():
            v.data = params[k].data.clone()

    def dispatch(self, index):
        for j in index:
            for old_params, new_params in zip(self.nns[j].parameters(), self.nn.parameters()):
                old_params.data = new_params.data.clone()

    def client_update(self, index):  # update nn
        for k in index:
            self.nns[k] = train(self.args, self.nns[k], self.nn)

    def global_test(self):
        model = self.nn
        model.eval()
        for client in self.args.clients:
            model.name = client
            test(self.args, model)

该FedProx类实现了FedProx算法的逻辑，通过在服务器和客户端之间交替进行参数分发、本地训练和参数聚合，实现模型的联邦学习。

逐行分析：

class FedProx:
    def __init__(self, args):
        self.args = args
        self.nn = ANN(args=self.args, name='server').to(args.device)
        self.nns = []
        for i in range(self.args.K):
            temp = copy.deepcopy(self.nn)
            temp.name = self.args.clients[i]
            self.nns.append(temp)

FedProx 类的 __init__ 方法用于初始化联邦学习的相关参数和模型，其中包括了初始化全局模型self.nn和self.nns列表；

args：模型的参数配置和超参数的集合
nn：全局模型，即服务器端的模型
nns：每个客户端模型的集合
K：客户端的数量
clients：客户端的标识集合
device：模型运行的设备（比如 CPU 或 GP

    def server(self):
        for t in tqdm(range(self.args.r)):
            print('round', t + 1, ':')
            # sampling
            m = np.max([int(self.args.C * self.args.K), 1])
            index = random.sample(range(0, self.args.K), m)  # st
            # dispatch
            self.dispatch(index)
            # local updating
            self.client_update(index)
            # aggregation
            self.aggregation(index)

        return self.nn

for t in tqdm(range(self.args.r)):这是一个迭代循环，循环次数由参数 self.args.r 指定，表示联邦学习的轮数。

print('round', t + 1, ':')在每轮训练开始时打印当前轮数。

m = np.max([int(self.args.C * self.args.K), 1])计算每轮随机选择参与训练的客户端数量 m，这里使用了参数 C 控制每轮参与训练的比例，参数 K 表示客户端的总数。

index = random.sample(range(0, self.args.K), m)从客户端列表中随机选择 m 个客户端，这些客户端将会参与当前轮的训练。这里使用了 Python 的 random.sample 方法来实现随机选择。

self.dispatch(index)调用 dispatch 方法，将全局模型参数分发给被选择的客户端模型。

self.client_update(index)调用 client_update 方法，客户端使用收到的全局模型参数进行本地模型训练，并更新自己的模型参数。

self.aggregation(index)调用 aggregation 方法，根据客户端的更新情况对全局模型参数进行聚合，以得到新的全局模型参数。

return self.nn返回经过多轮训练后的全局模型。

def aggregation(self, index):
        s = 0
        for j in index:
            # normal
            s += self.nns[j].len

        params = {}
        for k, v in self.nns[0].named_parameters():
            params[k] = torch.zeros_like(v.data)

        for j in index:
            for k, v in self.nns[j].named_parameters():
                params[k] += v.data * (self.nns[j].len / s)

        for k, v in self.nn.named_parameters():
            v.data = params[k].data.clone()

逐行解释：

def aggregation(self, index):

这是 aggregation 方法的定义，其中 self 是指代 FedProx 类的实例本身，index 是被选择的客户端的索引列表。

    s = 0 
    for j in index: 
        s += self.nns[j].len

这段代码计算被选中的客户端模型的总样本数量，并将其存储在变量 s 中，以便后续加权平均使用。

    params = {} 
    for k, v in self.nns[0].named_parameters(): 
        params[k] = torch.zeros_like(v.data)

在这里，代码初始化了一个空的字典 params，用于存储聚合后的参数，初始值为全局模型参数的零张量，以便稍后进行累加计算。

    for j in index: 
        for k, v in self.nns[j].named_parameters(): 
            params[k] += v.data * (self.nns[j].len / s)

上述代码对选中的客户端模型的参数进行加权累加。使用了一个加权平均的方法，其中权重是选中客户端的样本数量与总样本数量的比值。这意味着训练样本更多的客户端对全局模型参数的影响权重更大。

    for k, v in self.nn.named_parameters(): 
        v.data = params[k].data.clone()

最后，这段代码将计算得到的聚合参数复制到全局模型中，使全局模型得以更新。

    def dispatch(self, index):
        for j in index:
            for old_params, new_params in zip(self.nns[j].parameters(), self.nn.parameters()):
                old_params.data = new_params.data.clone()

这段代码的作用是将全局模型的参数分发给选中的客户端模型，以确保每个客户端在开始训练之前都具有最新的全局模型参数。

这里使用一个循环来遍历被选中的客户端模型的索引列表。
在这个循环中，代码使用 zip 函数同时迭代选中的客户端模型 self.nns[j] 的参数和全局模型 self.nn 的参数。这意味着对于每个客户端模型和全局模型都会迭代对应的参数。
在迭代过程中，代码将全局模型的参数值（new_params.data）复制到对应的客户端模型参数（old_params.data）上。这样做可以确保选中的客户端模型接收到了最新的全局模型参数。

    def client_update(self, index):  # update nn
        for k in index:
            self.nns[k] = train(self.args, self.nns[k], self.nn)

这段代码的作用是对选中的客户端模型进行训练，并将其更新为全局模型的最新参数。

在这个循环中，代码遍历了被选中的客户端模型的索引列表。
对于每个被选中的客户端模型，代码调用了名为 train 的函数，该函数接收客户端模型 self.nns[k]、全局模型 self.nn 和参数 self.args 作为输入。被选中的客户端模型会使用全局模型来进行训练，并将训练后的模型更新为最新参数。

    def global_test(self):
        model = self.nn
        model.eval()
        for client in self.args.clients:
            model.name = client
            test(self.args, model)

这段代码的作用是在全局模型上进行测试，针对每个客户端模型进行测试，并记录测试结果。

将全局模型 self.nn 赋值给变量 model。这样做是为了在测试阶段使用全局模型进行测试。
将模型切换到评估（evaluation）模式。在评估模式下，模型会关闭一些特定于训练时的操作，例如 dropout 等，以确保测试结果的稳定性。
在这个循环中，代码遍历了客户端模型的列表 self.args.clients。
这行代码试图给模型对象添加一个名为 name 的属性，并赋值为当前客户端的名称 client。然而，一般来说，PyTorch 模型对象是不支持动态添加属性的。可能这里是想要记录当前测试的是哪个客户端模型，但是这种方式并不是标准的做法。
调用名为 test 的函数，该函数接收客户端参数 self.args 和模型 model 作为输入。这个函数的作用是在给定的客户端上对模型进行测试。

4.2 client.py

def get_val_loss(args, model, Val):
    model.eval()
    loss_function = nn.MSELoss().to(args.device)
    val_loss = []
    for (seq, label) in Val:
        with torch.no_grad():
            seq = seq.to(args.device)
            label = label.to(args.device)
            y_pred = model(seq)
            loss = loss_function(y_pred, label)
            val_loss.append(loss.item())

    return np.mean(val_loss)


def train(args, model, server):
    model.train()
    Dtr, Val, Dte = nn_seq_wind(model.name, args.B)
    model.len = len(Dtr)
    global_model = copy.deepcopy(server)
    lr = args.lr
    if args.optimizer == 'adam':
        optimizer = torch.optim.Adam(model.parameters(), lr=lr,
                                     weight_decay=args.weight_decay)
    else:
        optimizer = torch.optim.SGD(model.parameters(), lr=lr,
                                    momentum=0.9, weight_decay=args.weight_decay)
    stepLR = StepLR(optimizer, step_size=args.step_size, gamma=args.gamma)
    # training
    min_epochs = 10
    best_model = None
    min_val_loss = 5
    print('training...')
    loss_function = nn.MSELoss().to(args.device)
    for epoch in tqdm(range(args.E)):
        train_loss = []
        for (seq, label) in Dtr:
            seq = seq.to(args.device)
            label = label.to(args.device)
            y_pred = model(seq)
            optimizer.zero_grad()
            # compute proximal_term
            proximal_term = 0.0
            for w, w_t in zip(model.parameters(), global_model.parameters()):
                proximal_term += (w - w_t).norm(2)

            loss = loss_function(y_pred, label) + (args.mu / 2) * proximal_term
            train_loss.append(loss.item())
            loss.backward()
            optimizer.step()
        stepLR.step()
        # validation
        val_loss = get_val_loss(args, model, Val)
        if epoch + 1 >= min_epochs and val_loss < min_val_loss:
            min_val_loss = val_loss
            best_model = copy.deepcopy(model)

        print('epoch {:03d} train_loss {:.8f} val_loss {:.8f}'.format(epoch, np.mean(train_loss), val_loss))
        model.train()

    return best_model


def test(args, ann):
    ann.eval()
    Dtr, Val, Dte = nn_seq_wind(ann.name, args.B)
    pred = []
    y = []
    for (seq, target) in tqdm(Dte):
        with torch.no_grad():
            seq = seq.to(args.device)
            y_pred = ann(seq)
            pred.extend(list(chain.from_iterable(y_pred.data.tolist())))
            y.extend(list(chain.from_iterable(target.data.tolist())))

    pred = np.array(pred)
    y = np.array(y)
    print('mae:', mean_absolute_error(y, pred), 'rmse:',
          np.sqrt(mean_squared_error(y, pred)))

这段函数的主要目标是用于训练、验证和测试神经网络模型。它们在训练过程中使用损失函数来优化模型的权重，并在验证和测试过程中评估模型的性能。

逐行分析：

def get_val_loss(args, model, Val):
    model.eval()
    loss_function = nn.MSELoss().to(args.device)
    val_loss = []
    for (seq, label) in Val:
        with torch.no_grad():
            seq = seq.to(args.device)
            label = label.to(args.device)
            y_pred = model(seq)
            loss = loss_function(y_pred, label)
            val_loss.append(loss.item())

    return np.mean(val_loss)

这段代码定义了一个函数`get_val_loss(args, model, Val)`，用于计算模型在验证集上的损失函数值。
首先，代码将模型设置为评估模式(`model.eval()`)，这是为了在验证过程中禁用一些特定层（如Dropout层）的随机行为。
然后，代码定义了损失函数(`loss_function = nn.MSELoss().to(args.device)`)。在这里，使用均方误差（MSE）作为损失函数。
接下来，代码通过遍历验证集中的样本 `(seq, label)`，并将它们移动到指定的设备上(`seq.to(args.device)`和`label.to(args.device)`)。这是为了确保在计算损失函数时使用与模型相同的设备。
在进入循环之前，使用`torch.no_grad()`上下文管理器来关闭梯度计算，因为验证阶段不需要进行梯度更新。
在循环中，将输入数据 seq 和 label 移动到指定的设备上（例如 GPU）。使用模型 model 进行推断，得到预测值 y_pred，而后计算预测值与真实标签之间的损失值 loss，并将将损失值保存到列表 val_loss 中。

这个函数的目的是在每个训练周期的末尾计算模型在验证集上的性能，以便对模型进行评估和选择。

def train(args, model, server):
    model.train()
    Dtr, Val, Dte = nn_seq_wind(model.name, args.B)
    model.len = len(Dtr)
    global_model = copy.deepcopy(server)
    lr = args.lr
    if args.optimizer == 'adam':
        optimizer = torch.optim.Adam(model.parameters(), lr=lr,
                                     weight_decay=args.weight_decay)
    else:
        optimizer = torch.optim.SGD(model.parameters(), lr=lr,
                                    momentum=0.9, weight_decay=args.weight_decay)
    stepLR = StepLR(optimizer, step_size=args.step_size, gamma=args.gamma)
    # training
    min_epochs = 10
    best_model = None
    min_val_loss = 5
    print('training...')
    loss_function = nn.MSELoss().to(args.device)
    for epoch in tqdm(range(args.E)):
        train_loss = []
        for (seq, label) in Dtr:
            seq = seq.to(args.device)
            label = label.to(args.device)
            y_pred = model(seq)
            optimizer.zero_grad()
            # compute proximal_term
            proximal_term = 0.0
            for w, w_t in zip(model.parameters(), global_model.parameters()):
                proximal_term += (w - w_t).norm(2)

            loss = loss_function(y_pred, label) + (args.mu / 2) * proximal_term
            train_loss.append(loss.item())
            loss.backward()
            optimizer.step()
        stepLR.step()
        # validation
        val_loss = get_val_loss(args, model, Val)
        if epoch + 1 >= min_epochs and val_loss < min_val_loss:
            min_val_loss = val_loss
            best_model = copy.deepcopy(model)

        print('epoch {:03d} train_loss {:.8f} val_loss {:.8f}'.format(epoch, np.mean(train_loss), val_loss))
        model.train()

    return best_model

min_epochs = 10: 定义最小训练轮数为10次。
best_model = None: 定义存储最佳模型的变量。
min_val_loss = 5: 定义最小验证集损失为5。
loss_function = nn.MSELoss().to(args.device): 定义使用均方误差损失函数。
for epoch in tqdm(range(args.E)):: 对于每个训练轮数进行循环。
for (seq, label) in Dtr:: 对于每个训练样本进行循环。
seq = seq.to(args.device): 将输入数据移到设备上进行计算。
label = label.to(args.device): 将标签数据移到设备上进行计算。
y_pred = model(seq): 使用模型进行预测。
optimizer.zero_grad(): 清除之前的梯度。
proximal_term = 0.0: 初始化近端项为0。
for w, w_t in zip(model.parameters(), global_model.parameters()):: 对于每个模型参数和全局模型参数进行循环。
proximal_term += (w - w_t).norm(2): 计算模型参数和全局模型参数之间的差异，并使用L2范数来度量。
loss = loss_function(y_pred, label) + (args.mu / 2) * proximal_term: 计算损失函数，包括预测值和标签之间的均方误差损失以及近端项。
train_loss.append(loss.item()): 将训练损失添加到列表中。
loss.backward(): 反向传播，计算梯度。
optimizer.step(): 更新模型参数。
stepLR.step(): 调整学习率。
val_loss = get_val_loss(args, model, Val): 计算验证集的损失。
if epoch + 1 >= min_epochs and val_loss < min_val_loss:: 如果达到最小训练轮数并且验证集损失小于最小验证集损失，则更新最小验证集损失和最佳模型。
model.train(): 将模型设置为训练模式。

def test(args, ann):
    ann.eval()
    Dtr, Val, Dte = nn_seq_wind(ann.name, args.B)
    pred = []
    y = []
    for (seq, target) in tqdm(Dte):
        with torch.no_grad():
            seq = seq.to(args.device)
            y_pred = ann(seq)
            pred.extend(list(chain.from_iterable(y_pred.data.tolist())))
            y.extend(list(chain.from_iterable(target.data.tolist())))

    pred = np.array(pred)
    y = np.array(y)
    print('mae:', mean_absolute_error(y, pred), 'rmse:',
          np.sqrt(mean_squared_error(y, pred)))

首先，代码将模型设置为评估模式 (ann.eval())，以禁用一些特定层（如Dropout层）的随机行为。
然后，代码调用nn_seq_wind(ann.name, args.B)函数以获取训练集、验证集和测试集。这个函数根据模型的名称和批次大小返回相应的数据集。
接下来，代码初始化了两个空列表pred和y，用于保存模型的预测值和目标值。
在测试集上进行循环，遍历每个样本(seq, target)。在这个循环中，代码进行了以下操作：

seq = seq.to(args.device): 将输入数据移到设备上进行计算。
y_pred = ann(seq): 使用训练好的模型进行预测。
pred.extend(list(chain.from_iterable(y_pred.data.tolist()))):: 将预测值添加到列表中。
y.extend(list(chain.from_iterable(target.data.tolist()))):: 将真实值添加到列表中。
pred = np.array(pred)：将预测值转换为NumPy数组。
y = np.array(y)：将真实值转换为NumPy数组。
print('mae:', mean_absolute_error(y, pred), 'rmse:', np.sqrt(mean_squared_error(y, pred))): 打印平均绝对误差和均方根误差。

五、main.py

def main():
    args = args_parser()
    fedProx = FedProx(args)
    fedProx.server()
    fedProx.global_test()


if __name__ == '__main__':
    main()

args = args_parser(): 调用函数来获取命令行参数，并将其赋值给变量。
fedProx = FedProx(args): 创建一个对象，传入参数。
fedProx.server(): 调用对象的方法，启动服务器。
fedProx.global_test(): 调用对象的方法，进行全局测试。

这段代码是将联邦学习的训练和测试过程封装到一个函数中，并通过main()函数的调用来执行这些操作。这种设计方式可以提高代码的整洁性和可维护性。在其他的模块中，你可以直接导入这个模块，并通过调用main()函数来执行联邦学习的训练和测试。

大致就这样，欢迎各位大佬纠错或指导！