目录顺序为函数调用顺序!!
main_fed_snr.py :主要模型的训练。
packages
models
utils
options.py
sampling.py
Nets.py
options.py
shufflenetv2.py
test.py
Update.py
Fed.py
BPDecoding.py
options.py:建立命令行参数解析对象,添加实例属性。
sampling.py:引入了数据集分配方式。
Nets.py:定义了几种模型。如果args.dataset == 'mnist' ,就调用Nets.py;否则,选择ShuffleNetV2。
shufflenetv2.py:定义了shufflenetv2网络模型。
test.py:测试过程,首先根据数据集对测试集数据进行加载。对网络模型进行测试。
Update.py:客户端本地训练。
Fed.py:进行模型参数中心聚合。包含 完美信道 和 异步。
BPDecoding.py:计算接收信号功率,添加噪声并利用两个估计器(ML+对齐样本),进行
估计求和。
一、Main函数
def Main(args):
setup_seed(args.seed)
# load dataset and split users---加载数据集和拆分客户端
if args.dataset == 'mnist':
trans_mnist = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
#train参数表示是训练集还是测试集,如果是true,则从训练集创建,否则从测试集创建。transform参数是为了对数据集图片进行处理。
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)#num_users参数定义在options文件中,40
else:
dict_users = mnist_noniid(dataset_train, args.num_users)
elif args.dataset == 'cifar':
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
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)),
])
trans_cifar = 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, download=True, transform=transform_train)
dataset_test = datasets.CIFAR10('../data/cifar', train=False, download=True, transform=transform_test)
if args.iid:
dict_users = cifar_iid(dataset_train, args.num_users)#cifar_iid函数在sampling文件
else:
dict_users = cifar_noniid(dataset_train, args.num_users)
else:
exit('Error: unrecognized dataset')
# build model---构建模型选择-ShuffleNetV2
if args.model == 'cnn' and args.dataset == 'cifar':
# net_glob = CNNCifar(args=args).to(args.device)
net_glob = ShuffleNetV2(1).to(args.device)
elif args.model == 'cnn' and args.dataset == 'mnist':
net_glob = CNNMnist(args=args).to(args.device)
else:
exit('Error: unrecognized model')
net_glob.train()
print("============================== Federated Learning ... ...")
# training---训练过程
loss_train = []
acc_store = np.array([])
#torch.nn.Module模块中的state_dict变量存放训练过程中需要学习的权重和偏置系数 state_dict本质上Python字典对象
w_glob = net_glob.state_dict() # initial global weights 复制权重
for iter in range(args.epochs):
# record the running time of an iteration
startTime = time.time()
# Every 5 itertaions, evaluate the learning performance in terms of "test accuracy"---每5次迭代,根据“测试准确性”评估学习性能
if np.mod(iter,5) == 0:
net_glob.eval()
acc_test, _ = test_img(net_glob, dataset_test, args)
acc_store = np.append(acc_store, acc_test.numpy())
print("Test accuracies every 5 itertaions =", acc_store)
net_glob.train()
# ----------------------------------------------------------------------- Preparation
# global model at the beginning of the iteration---迭代开始时的全局模型
history_dict = net_glob.state_dict()
# random choose M devices out of the args.num_users devices---从参数中随机选择M个设备。num_users设备
M = max(int(args.frac * args.num_users), 1)
idxs_users = np.random.choice(range(args.num_users), M, replace=False)
args.lr = args.lr * 0.95 # learning rate adjustment
# ----------------------------------------------------------------------- Local Training
w_locals = [] # store the local "updates" (the difference) of M devices---存储M个设备的本地“更新”(差异)
loss_locals = [] # store the local training loss--存储本地培训损失
for idx in idxs_users:
#LocalUpdate在Update文件中
local = LocalUpdate(args=args, dataset=dataset_train, idxs=dict_users[idx]) #对每个worker进行本地更新
#本地训练,返回net.state_dict()和平均loss
w, loss = local.train(net=copy.deepcopy(net_glob).to(args.device), history_dict=history_dict)
if args.all_clients: ## 选择利用所有客户进行聚合
w_locals[idx] = copy.deepcopy(w)
else:
w_locals.append(copy.deepcopy(w))
loss_locals.append(copy.deepcopy(loss))
# ----------------------------------------------------------------------- Federated Averaging
if args.Aligned == 0: # perfect channel ->->-> no misalignments, no noise
current_dict = FedAvg(w_locals, args) #对全局w进行聚合更新
elif args.Aligned == 1: # channel misalignment, symbol misalignment, noise
current_dict = FedAvg_ML(w_locals, args)
else:
exit("unknown args.Aligned")
# ----------------------------------------------------------------------- Reconstruct the new model at the PS
#联邦平均,输出为更新后的全局模型
for k in current_dict.keys():
w_glob[k] = history_dict[k] + current_dict[k]
# load new model
net_glob.load_state_dict(w_glob)
# print training loss
loss_avg = sum(loss_locals) / len(loss_locals)
print('Round {:3d}, Average loss {:.3f}, Time Cosumed {:.3f}'.format(iter, loss_avg, time.time()-startTime))
loss_train.append(loss_avg)
# testing
net_glob.eval()
acc_train, loss_train = test_img(net_glob, dataset_train, args)
acc_test, loss_test = test_img(net_glob, dataset_test, args)
acc_store = np.append(acc_store, acc_test.numpy())
return np.array([acc_store])
这是一个用于联邦学习(Federated Learning)的 Python 脚本,其中的 `Main` 函数是整个联邦学习过程的主要执行部分。我将逐步解释这个函数的主要功能:
1. `setup_seed(args.seed)`: 设置随机数种子,以确保实验的可重复性。这通常是为了在不同运行中得到相同的随机数。
2. 数据集加载和用户拆分:
- 针对不同的数据集(MNIST或CIFAR-10),通过 `if args.dataset == 'mnist':` 或 `elif args.dataset == 'cifar':` 进行选择。对于MNIST数据集,采用`mnist_iid` 或 `mnist_noniid` 函数拆分用户;对于CIFAR-10数据集,采用 `cifar_iid` 或 `cifar_noniid` 函数拆分用户。
3. 模型构建:
- 根据数据集和模型选择,构建对应的神经网络模型。在这里,根据 `args.model` 和 `args.dataset` 的值选择使用 ShuffleNetV2 或 CNNMnist 模型。
4. 迭代训练过程:
- 使用循环迭代进行训练,每次迭代被称为一个 "Round"。
- 每 5 次迭代,评估模型在测试集上的性能,并打印测试准确性。
- 在每次迭代开始时,记录全局模型的当前状态 (`history_dict`)。
- 随机选择一部分用户进行训练,设置学习率,并执行本地训练 (`local.train`)。本地训练返回本地模型更新和损失。
- 根据本地模型更新,使用联邦平均(`FedAvg` 或 `FedAvg_ML`)将本地模型的更新聚合到全局模型。
- 重构全局模型,加载新的模型参数,用于下一轮的训练。
- 打印每轮的平均损失和运行时间。
5. 训练完成后进行测试:
- 在整个训练完成后,将模型设置为评估模式 (`net_glob.eval()`)。
- 在训练集和测试集上进行最终的评估,并记录准确性和损失。
- 将测试准确性存储在 `acc_store` 数组中,最后将其作为输出返回。
总体来说,这段代码实现了一个基本的联邦学习框架,其中涉及模型的初始化、数据集的加载和拆分、本地训练、全局模型的更新和重构,以及最终的性能评估。
1.1 数据集加载和用户拆分
if args.dataset == 'mnist':
trans_mnist = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
#train参数表示是训练集还是测试集,如果是true,则从训练集创建,否则从测试集创建。transform参数是为了对数据集图片进行处理。
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)#num_users参数定义在options文件中,40
else:
dict_users = mnist_noniid(dataset_train, args.num_users)
elif args.dataset == 'cifar':
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
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)),
])
trans_cifar = 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, download=True, transform=transform_train)
dataset_test = datasets.CIFAR10('../data/cifar', train=False, download=True, transform=transform_test)
if args.iid:
dict_users = cifar_iid(dataset_train, args.num_users)#cifar_iid函数在sampling文件
else:
dict_users = cifar_noniid(dataset_train, args.num_users)
else:
exit('Error: unrecognized dataset')
这部分代码主要负责加载数据集和准备用于联邦学习的数据。以下是对每个部分的解释:
1. **MNIST数据集:**
- 如果指定的数据集是'MNIST',则首先创建一个用于处理图像的变换管道 `trans_mnist`,该管道包括将图像转换为张量和进行归一化操作。
- 然后使用`datasets.MNIST`从本地下载或加载MNIST数据集。对于训练集,`train=True`;对于测试集,`train=False`。
- 如果 `args.iid`(独立同分布)为True,则调用 `mnist_iid` 函数,该函数将训练集分为 `args.num_users` 个独立同分布的用户。否则,调用 `mnist_noniid` 函数,该函数将训练集分为 `args.num_users` 个非独立同分布的用户。
2. **CIFAR数据集:**
- 如果指定的数据集是'CIFAR',则定义了两个转换管道 `transform_train` 和 `transform_test`,分别用于训练和测试数据。这些转换包括随机裁剪、水平翻转、转换为张量和归一化。
- 同样,使用 `datasets.CIFAR10` 从本地下载或加载CIFAR-10数据集。对于训练集,`train=True`;对于测试集,`train=False`。
- 如果 `args.iid` 为True,则调用 `cifar_iid` 函数【在sampling.py中】,该函数将训练集分为 `args.num_users` 个独立同分布的用户。否则,调用 `cifar_noniid` 函数,该函数将训练集分为 `args.num_users` 个非独立同分布的用户。
3. **错误处理:**
- 如果指定的数据集不是'MNIST'或'CIFAR',则输出错误消息并退出程序。
总体而言,这段代码的目标是根据指定的数据集和联邦学习参数,准备好用于联邦学习的训练集和测试集,并根据是否独立同分布的要求划分用户。
1.2 模型构建
if args.model == 'cnn' and args.dataset == 'cifar':
# net_glob = CNNCifar(args=args).to(args.device)
net_glob = ShuffleNetV2(1).to(args.device)
elif args.model == 'cnn' and args.dataset == 'mnist':
net_glob = CNNMnist(args=args).to(args.device)
else:
exit('Error: unrecognized model')
net_glob.train()
这段代码负责选择和初始化用于联邦学习的模型。以下是对每个部分的解释:
1. **选择CNN或ShuffleNetV2模型:**
- 如果命令行参数 `args.model` 的值是'cnn' 并且 `args.dataset` 的值是'cifar',则选择使用ShuffleNetV2模型。
- 如果 `args.model` 的值是'cnn' 并且 `args.dataset` 的值是'mnist',则选择使用CNNMnist模型。
- 如果上述条件都不满足,则输出错误消息并退出程序。
2. **模型初始化:**
- 如果选择使用ShuffleNetV2模型,使用`ShuffleNetV2【在shuffletNetV2.py中】(1)`创建一个ShuffleNetV2实例,并将其移动到指定的设备 `args.device` 上。
- 如果选择使用CNNMnist模型,使用`CNNMnist(args=args)`创建一个CNNMnist实例,并将其移动到指定的设备 `args.device` 上。
- 调用 `net_glob.train()` 将模型设置为训练模式。
总体而言,这段代码的作用是根据命令行参数选择合适的模型,并将该模型实例移动到指定的设备上。这是联邦学习中全局模型的初始化步骤。
1.3 迭代训练过程
loss_train = []
acc_store = np.array([])
#torch.nn.Module模块中的state_dict变量存放训练过程中需要学习的权重和偏置系数 state_dict本质上Python字典对象
w_glob = net_glob.state_dict() # initial global weights 复制权重
for iter in range(args.epochs):
# record the running time of an iteration
startTime = time.time()
# Every 5 itertaions, evaluate the learning performance in terms of "test accuracy"---每5次迭代,根据“测试准确性”评估学习性能
if np.mod(iter,5) == 0:
net_glob.eval()
acc_test, _ = test_img(net_glob, dataset_test, args)
acc_store = np.append(acc_store, acc_test.numpy())
print("Test accuracies every 5 itertaions =", acc_store)
net_glob.train()
# ----------------------------------------------------------------------- Preparation
# global model at the beginning of the iteration---迭代开始时的全局模型
history_dict = net_glob.state_dict()
# random choose M devices out of the args.num_users devices---从参数中随机选择M个设备。num_users设备
M = max(int(args.frac * args.num_users), 1)
idxs_users = np.random.choice(range(args.num_users), M, replace=False)
args.lr = args.lr * 0.95 # learning rate adjustment
# ----------------------------------------------------------------------- Local Training
w_locals = [] # store the local "updates" (the difference) of M devices---存储M个设备的本地“更新”(差异)
loss_locals = [] # store the local training loss--存储本地培训损失
for idx in idxs_users:
#LocalUpdate在Update文件中
local = LocalUpdate(args=args, dataset=dataset_train, idxs=dict_users[idx]) #对每个worker进行本地更新
#本地训练,返回net.state_dict()和平均loss
w, loss = local.train(net=copy.deepcopy(net_glob).to(args.device), history_dict=history_dict)
if args.all_clients: ## 选择利用所有客户进行聚合
w_locals[idx] = copy.deepcopy(w)
else:
w_locals.append(copy.deepcopy(w))
loss_locals.append(copy.deepcopy(loss))
这部分代码实现了联邦学习的主要训练循环。以下是对每个部分的解释:
1. **初始化变量:**
- `loss_train`是一个列表,用于存储每轮训练的平均损失。
- `acc_store`是一个NumPy数组,用于存储每5轮测试后的测试准确度。
- `w_glob`是全局模型的权重字典,初始时由 `net_glob` 的状态字典复制而来。
2. **训练循环:**
- `for iter in range(args.epochs):` 开始了训练的主循环,共进行 `args.epochs` 次迭代。
- `startTime = time.time()` 记录每轮迭代的开始时间。
3. **每5轮进行一次测试:**
- `if np.mod(iter, 5) == 0:` 每5轮测试一次模型性能。将全局模型设置为评估模式 (`net_glob.eval()`),然后使用测试数据集进行测试,记录测试准确度,并将准确度追加到 `acc_store` 中。
- 恢复模型为训练模式 (`net_glob.train()`)。
4. **全局模型准备:**
- `history_dict = net_glob.state_dict()` 在每轮迭代开始时,将全局模型的状态字典保存为 `history_dict`。
5. **选择一部分本地设备:**
- `M = max(int(args.frac * args.num_users), 1)` 计算每轮选择的本地设备数目。
- `idxs_users = np.random.choice(range(args.num_users), M, replace=False)` 随机选择M个本地设备的索引。
6. **本地训练:**
- 对于每个选定的本地设备,使用 `LocalUpdate` 类【在update.py中】进行本地模型训练,得到本地模型的权重和损失。将本地更新和本地损失分别保存在 `w_locals` 和 `loss_locals` 中。
7. **联邦平均:**
- `current_dict` 通过调用 `FedAvg[在Fed.py中]` 或 `FedAvg_ML` 方法,对本地更新进行联邦平均,得到当前轮次的全局更新。
8. **更新全局模型:**
- 将全局模型的权重更新为 `history_dict` 和 `current_dict` 的组合。
- `net_glob.load_state_dict(w_glob)` 将更新后的权重加载到 `net_glob` 中。
9. **打印训练损失:**
- 计算并打印当前轮次的平均本地训练损失。
- 将平均损失追加到 `loss_train` 中。
总体而言,这段代码实现了一个联邦学习的迭代训练过程,其中包括周期性的模型评估、本地训练、联邦平均和全局模型更新。
二、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))]
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
def mnist_noniid(dataset, num_users):
"""
Sample non-I.I.D client data from MNIST dataset
:param dataset:
:param num_users:
:return:
"""
num_shards, num_imgs = 120, 500
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.targets.numpy()
# sort labels
idxs_labels = np.vstack((idxs, labels))
idxs_labels = idxs_labels[:,idxs_labels[1,:].argsort()]
idxs = idxs_labels[0,:]
# divide and assign
for i in range(num_users):
rand_set = set(np.random.choice(idx_shard, 3, 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
'''从 CIFAR 数据集中创建非独立同分布(Non-I.I.D.)的客户端数据集的函数.
帮助创建数据分布,以便在联邦学习等任务中模拟客户端之间的不同数据分布情况。
'''
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
def cifar_noniid_shard(dataset, num_users):
"""
Sample non-I.I.D client data from cifar dataset
:param dataset:
:param num_users:
:return:
"""
classPeruser = 5
num_shards, num_imgs = num_users * classPeruser, int(len(dataset)/num_users / classPeruser)
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 = np.array(dataset.targets)
# sort labels
idxs_labels = np.vstack((idxs, labels))
idxs_labels = idxs_labels[:,idxs_labels[1,:].argsort()]
idxs = idxs_labels[0,:]
# divide and assign
for i in range(num_users):
rand_set = set(np.random.choice(idx_shard, classPeruser, 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
def cifar_noniid(dataset, num_users):
"""
Sample non-I.I.D client data from cifar dataset
:param dataset:
:param num_users:
:return:
Each device randomly sample
"""
lenRandom = 40000
num_items = int(lenRandom/num_users)
dict_users, all_idxs = {}, [i for i in range(len(dataset))]
for ii in range(num_users):
dict_users[ii] = set(np.random.choice(all_idxs, num_items, replace=False))
all_idxs = list(set(all_idxs) - dict_users[ii])
labels = np.array(dataset.targets)
labels = labels[all_idxs]
# sort labels
idxs = np.arange(len(labels))
idxs_labels = np.vstack((idxs, labels))
idxs_labels = idxs_labels[:,idxs_labels[1,:].argsort()]
idxs = idxs_labels[0,:]
# divide and assign
numImage = int(len(idxs)/num_users)
for ii in range(num_users):
temp = idxs[ii*numImage:(ii+1)*numImage]
dict_users[ii] = np.concatenate((list(dict_users[ii]), temp), axis=0)
return dict_users
if __name__ == '__main__':
dataset_train = datasets.MNIST('../data/mnist/', train=True, download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
]))
num = 100
d = mnist_noniid(dataset_train, num)
这个 `sampling.py` 文件包含了用于创建非独立同分布(Non-I.I.D.)数据分布的函数,特别适用于联邦学习中的客户端数据。以下是对每个函数的解释:
1. **`mnist_iid` 函数:**
- 从MNIST数据集中按照独立同分布(I.I.D.)的方式为客户端抽样数据。
- 将MNIST数据集分为 `num_users` 个用户,每个用户分配相等数量的图像。
2. **`mnist_noniid` 函数:**
- 从MNIST数据集中按照非独立同分布(Non-I.I.D.)的方式为客户端抽样数据。
- 将MNIST数据集分为 `num_users` 个用户,每个用户从120个虚构的"shard"(数据分片) 中随机选择3个,每个shard包含500个图像。
3. **`cifar_iid` 函数:**
- 从CIFAR-10数据集中按照独立同分布(I.I.D.)的方式为客户端抽样数据。
- 将CIFAR-10数据集分为 `num_users` 个用户,每个用户分配相等数量的图像。
4. **`cifar_noniid_shard` 函数:**
- 从CIFAR-10数据集中按照非独立同分布(Non-I.I.D.)的方式为客户端抽样数据。
- 将CIFAR-10数据集分为 `num_users * classPeruser` 个用户,每个用户从总共的120个shard中随机选择 `classPeruser` 个,每个shard包含500个图像。
5. **`cifar_noniid` 函数:**
- 从CIFAR-10数据集中按照非独立同分布(Non-I.I.D.)的方式为客户端抽样数据。
- 每个用户从整个CIFAR-10数据集中随机选择相等数量的图像,但不同用户之间的选择是独立的。
在 `__main__` 部分,展示了如何使用这些函数。具体地,使用 `mnist_noniid` 函数从MNIST数据集中为100个用户创建了非独立同分布的客户端数据。
三、ShuffleNetV2.py
class ShuffleBlock(nn.Module):
def __init__(self, groups=2):
super(ShuffleBlock, self).__init__()
self.groups = groups
def forward(self, x):
'''Channel shuffle: [N,C,H,W] -> [N,g,C/g,H,W] -> [N,C/g,g,H,w] -> [N,C,H,W]'''
N, C, H, W = x.size()
g = self.groups
return x.view(N, g, C//g, H, W).permute(0, 2, 1, 3, 4).reshape(N, C, H, W)
class SplitBlock(nn.Module):
def __init__(self, ratio):
super(SplitBlock, self).__init__()
self.ratio = ratio
def forward(self, x):
c = int(x.size(1) * self.ratio)
return x[:, :c, :, :], x[:, c:, :, :]
class BasicBlock(nn.Module):
def __init__(self, in_channels, split_ratio=0.5):
super(BasicBlock, self).__init__()
self.split = SplitBlock(split_ratio)
in_channels = int(in_channels * split_ratio)
self.conv1 = nn.Conv2d(in_channels, in_channels,
kernel_size=1, bias=False)
self.bn1 = nn.BatchNorm2d(in_channels)
self.conv2 = nn.Conv2d(in_channels, in_channels,
kernel_size=3, stride=1, padding=1, groups=in_channels, bias=False)
self.bn2 = nn.BatchNorm2d(in_channels)
self.conv3 = nn.Conv2d(in_channels, in_channels,
kernel_size=1, bias=False)
self.bn3 = nn.BatchNorm2d(in_channels)
self.shuffle = ShuffleBlock()
def forward(self, x):
x1, x2 = self.split(x)
out = F.relu(self.bn1(self.conv1(x2)))
out = self.bn2(self.conv2(out))
out = F.relu(self.bn3(self.conv3(out)))
out = torch.cat([x1, out], 1)
out = self.shuffle(out)
return out
class DownBlock(nn.Module):
def __init__(self, in_channels, out_channels):
super(DownBlock, self).__init__()
mid_channels = out_channels // 2
# left
self.conv1 = nn.Conv2d(in_channels, in_channels,
kernel_size=3, stride=2, padding=1, groups=in_channels, bias=False)
self.bn1 = nn.BatchNorm2d(in_channels)
self.conv2 = nn.Conv2d(in_channels, mid_channels,
kernel_size=1, bias=False)
self.bn2 = nn.BatchNorm2d(mid_channels)
# right
self.conv3 = nn.Conv2d(in_channels, mid_channels,
kernel_size=1, bias=False)
self.bn3 = nn.BatchNorm2d(mid_channels)
self.conv4 = nn.Conv2d(mid_channels, mid_channels,
kernel_size=3, stride=2, padding=1, groups=mid_channels, bias=False)
self.bn4 = nn.BatchNorm2d(mid_channels)
self.conv5 = nn.Conv2d(mid_channels, mid_channels,
kernel_size=1, bias=False)
self.bn5 = nn.BatchNorm2d(mid_channels)
self.shuffle = ShuffleBlock()
def forward(self, x):
# left
out1 = self.bn1(self.conv1(x))
out1 = F.relu(self.bn2(self.conv2(out1)))
# right
out2 = F.relu(self.bn3(self.conv3(x)))
out2 = self.bn4(self.conv4(out2))
out2 = F.relu(self.bn5(self.conv5(out2)))
# concat
out = torch.cat([out1, out2], 1)
out = self.shuffle(out)
return out
class ShuffleNetV2(nn.Module):
def __init__(self, net_size):
super(ShuffleNetV2, self).__init__()
out_channels = configs[net_size]['out_channels']
num_blocks = configs[net_size]['num_blocks']
self.conv1 = nn.Conv2d(3, 24, kernel_size=3,
stride=1, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(24)
self.in_channels = 24
self.layer1 = self._make_layer(out_channels[0], num_blocks[0])
self.layer2 = self._make_layer(out_channels[1], num_blocks[1])
self.layer3 = self._make_layer(out_channels[2], num_blocks[2])
self.conv2 = nn.Conv2d(out_channels[2], out_channels[3],
kernel_size=1, stride=1, padding=0, bias=False)
self.bn2 = nn.BatchNorm2d(out_channels[3])
self.linear = nn.Linear(out_channels[3], 10)
def _make_layer(self, out_channels, num_blocks):
layers = [DownBlock(self.in_channels, out_channels)]
for i in range(num_blocks):
layers.append(BasicBlock(out_channels))
self.in_channels = out_channels
return nn.Sequential(*layers)
def forward(self, x):
out = F.relu(self.bn1(self.conv1(x)))
# out = F.max_pool2d(out, 3, stride=2, padding=1)
out = self.layer1(out)
out = self.layer2(out)
out = self.layer3(out)
out = F.relu(self.bn2(self.conv2(out)))
out = F.avg_pool2d(out, 4)
out = out.view(out.size(0), -1)
out = self.linear(out)
return out
configs = {
0.5: {
'out_channels': (48, 96, 192, 1024),
'num_blocks': (3, 7, 3)
},
1: {
'out_channels': (116, 232, 464, 1024),
'num_blocks': (3, 7, 3)
},
1.5: {
'out_channels': (176, 352, 704, 1024),
'num_blocks': (3, 7, 3)
},
2: {
'out_channels': (224, 488, 976, 2048),
'num_blocks': (3, 7, 3)
}
}
def test():
net = ShuffleNetV2(net_size=0.5)
x = torch.randn(3, 3, 32, 32)
y = net(x)
这是一个使用 PyTorch 实现的 ShuffleNetV2 模型,ShuffleNetV2 是一种轻量级的卷积神经网络结构。让我解释代码的各个部分:
1. `ShuffleBlock` 类:
- `ShuffleBlock` 定义了一个通道洗牌(Channel Shuffle)操作。这个操作将输入张量的通道分成若干组(由 `groups` 参数指定),然后对这些组进行洗牌,并最终将它们合并成一个输出张量。这通常用于提高模型的表达能力。
2. `SplitBlock` 类:
- `SplitBlock` 定义了一个张量分割操作。根据给定的比率 `ratio`,将输入张量在通道维度上分割成两个部分。
3. `BasicBlock` 类:
- `BasicBlock` 是 ShuffleNetV2 中的基本块。它包含了一系列卷积、批归一化和激活函数操作。其中,`SplitBlock` 用于将输入分割,然后通过一系列卷积和批归一化操作,最终通过通道洗牌操作将分割后的张量与原始输入合并。
4. `DownBlock` 类:
- `DownBlock` 定义了 ShuffleNetV2 中的下采样块。它包括两个分支,一个分支通过深度可分离卷积实现降采样,另一个分支通过标准卷积实现。最后,两个分支的结果通过通道洗牌操作合并。
5. `ShuffleNetV2` 类:
- `ShuffleNetV2` 是整个模型的主要定义。它包含了卷积层、批归一化层、多个 `DownBlock` 和多个 `BasicBlock`,以及全局平均池化和线性层用于分类。
6. `configs` 字典:
- `configs` 字典包含了不同网络大小的配置信息,其中定义了不同网络大小下的输出通道数和块的数量。
7. `test` 函数:
- `test` 函数创建了一个 ShuffleNetV2 模型(使用 `net_size=0.5`),并对随机输入进行前向传播。
总体而言,这段代码实现了 ShuffleNetV2 模型,并通过 `configs` 字典定义了不同网络大小的配置。这种网络结构设计旨在在保持模型轻量级的同时提供良好的性能,特别适用于移动设备和边缘计算环境。
四、update.py
#客户端本地训练
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, num_workers=4)
def train(self, net, history_dict):
net.train()
# train and update
optimizer = torch.optim.SGD(net.parameters(), lr=self.args.lr, momentum=self.args.momentum)
# optimizer = torch.optim.SGD(net.parameters(), lr=self.args.lr, momentum=0.9, weight_decay=5e-4)
# scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=200)
epoch_loss = []
for iter in range(self.args.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))
current_dict = net.state_dict()
for k in current_dict.keys():
current_dict[k] -= history_dict[k]
return current_dict, sum(epoch_loss) / len(epoch_loss)
这个 `update.py` 文件定义了两个类:`DatasetSplit` 和 `LocalUpdate`。
1. `DatasetSplit` 类:
- 该类继承自 PyTorch 的 `Dataset` 类,用于创建一个分割后的数据集。`DatasetSplit` 接受原始数据集和一个索引列表 `idxs`,并返回仅包含指定索引的数据集。在这个文件中,它的主要作用是根据客户端的索引创建相应的本地训练数据集。
2. `LocalUpdate` 类:
- 该类是客户端的本地更新类。在初始化时,它接收一些参数(`args`,`dataset`,`idxs`),其中 `args` 包含了一些超参数(如学习率、本地训练轮数等),`dataset` 是整个数据集,`idxs` 是本地客户端所拥有的样本的索引列表。
- `__init__` 方法初始化了一些参数,包括损失函数 (`nn.CrossEntropyLoss()`),选择的客户端 (`self.selected_clients`) 和一个 PyTorch 数据加载器 (`self.ldr_train`),用于加载本地训练数据集。
- `train` 方法执行本地训练。在每个本地训练轮次中,它遍历数据加载器 (`self.ldr_train`),在每个小批次上执行前向传播、损失计算、反向传播和参数更新。每个小批次的损失值被记录,最后计算每个本地训练轮次的平均损失。在训练结束后,该方法返回更新后的模型参数和平均损失。
总体而言,这个文件中的类主要负责为每个客户端执行本地训练,并返回本地更新后的模型参数和训练损失。这是联邦学习中客户端本地训练的核心逻辑。
五、Fed.py
def FedAvg_ML(w, args):
M = len(w) # 设备的数量
# 从每个设备提取符号(权重更新)并转换为numpy数组(复杂序列)
StoreRecover = np.array([]) # 记录如何将numpy数组转换回字典
for m in np.arange(M):
wEach = w[m]
eachWeightNumpy = np.array([])
for k in wEach.keys():
if k in ToIgnore:
continue
temp = wEach[k].cpu().numpy()
temp, unflatten = flatten(temp)
if m == 0:
StoreRecover = np.append(StoreRecover, unflatten)
eachWeightNumpy = np.append(eachWeightNumpy, temp)
if m == 0:
complexSymbols = eachWeightNumpy[0:int(len(eachWeightNumpy)/2)] + 1j * eachWeightNumpy[int(len(eachWeightNumpy)/2):]
TransmittedSymbols = np.array([complexSymbols])
else:
complexSymbols = eachWeightNumpy[0:int(len(eachWeightNumpy)/2)] + 1j * eachWeightNumpy[int(len(eachWeightNumpy)/2):]
TransmittedSymbols = np.r_[TransmittedSymbols, np.array([complexSymbols])]
# 每个设备的复杂符号的数量
L = len(TransmittedSymbols[0]) # 631927
# ---------------------------------------------------------------------------------- 逐个数据包传输
# 添加 1 个全零列 => 631927 + 1 = 631928(11 * 4 * 43 * 167 * 2)= 57448 * 11 或 28724 * 22
TransmittedSymbols = np.c_[TransmittedSymbols, np.zeros([M, 1])] # 4 * 631928
# -------------------------------------------------------- 长数据包
if args.short == 0:
numPkt = 44
lenPkt = int((L+1)/numPkt)
pool = Pool(processes=numPkt) # 创建 11 个进程
results = []
for idx in range(numPkt):
# 在一个数据包中传输的复杂符号
onePkt = TransmittedSymbols[:, (idx*lenPkt):((idx+1)*lenPkt)]
# 每个数据包中接收到的复杂符号(经过估计和平均化后)
results.append(pool.apply_async(per_pkt_transmission, (args, M, onePkt, )))
pool.close() # 关闭进程池,不再添加进程
pool.join() # 等待所有进程完成
for idx in range(len(results)):
try:
output = results[idx].get()
except:
pdb.set_trace()
if idx == 0:
ReceivedComplexPkt = output
else:
ReceivedComplexPkt = np.append(ReceivedComplexPkt, output)
else:
# -------------------------------------------------------- 短数据包
numPkt = 946 # 22 * 43
lenPkt = int((L+1)/numPkt)
# 多进程
numCPU = 22
for loop in range(43):
pktBatch = TransmittedSymbols[:, (loop*lenPkt*numCPU):((loop+1)*lenPkt*numCPU)]
pool = Pool(processes=numCPU) # 创建 11 个进程
results = []
for idx in range(numCPU):
# 在一个数据包中传输的复杂符号
onePkt = pktBatch[:, (idx*lenPkt):((idx+1)*lenPkt)]
# 每个数据包中接收到的复杂符号(经过估计和平均化后)
results.append(pool.apply_async(per_pkt_transmission, (args, M, onePkt, )))
pool.close() # 关闭进程池,不再添加进程
pool.join() # 等待所有进程完成
for idx in range(len(results)):
try:
output = results[idx].get()
except:
pdb.set_trace()
if idx == 0:
ReceivedBatch = output
else:
ReceivedBatch = np.append(ReceivedBatch, output)
if loop == 0:
ReceivedComplexPkt = ReceivedBatch
else:
ReceivedComplexPkt = np.append(ReceivedComplexPkt, ReceivedBatch)
# 恢复真实权重 ->->-> numpy数组
ReceivedPkt = np.append(np.real(ReceivedComplexPkt[:-1]), np.imag(ReceivedComplexPkt[:-1])) # 必须删除最后一个元素(0)
## =========================================================================================
## =========================================================================================
## ========================================== 从numpy数组重构字典
# 先运行联邦平均化(以解决批标准化层的问题)
w_avg = FedAvg(w, args, 1)
startIndex = 0
idx = 0
for k in w_avg.keys():
# 只更新w_avg中的非批标准化层
if k not in ToIgnore:
lenLayer = w_avg[k].numel()
# 获取数据
ParamsLayer = ReceivedPkt[startIndex:(startIndex+lenLayer)]
# 重塑
ParamsLayer_reshaped = StoreRecover[idx](ParamsLayer)
# 转换为torch并移至cuda()
w_avg[k] = torch.from_numpy(ParamsLayer_reshaped).cuda()
startIndex += lenLayer
idx += 1
return w_avg
这段代码主要涉及联邦学习中的异步通信和模型参数的传输。在这里,符号是指模型参数的一种表示形式。这段代码首先将本地模型的权重转换为复杂数符号,然后通过异步的方式传输和聚合这些符号,最后将聚合后的符号还原为全局模型的权重。
-
FedAvg
函数:- 这个函数是用于同步联邦学习的,其中所有设备在每轮迭代后将其本地模型参数传输到中央服务器进行平均化。
- 函数参数中的
flag
默认为 0,表示采用同步方式,即等待所有设备传输完毕后再进行模型参数的平均化。 - 函数接收一个包含所有本地模型参数的列表
w
,并返回这些参数的平均值,用于更新中央服务器上的全局模型。
-
FedAvg_ML
函数:- 这个函数是用于异步联邦学习的,其中每个设备在准备好时可以选择进行模型参数的传输。
- 函数中通过
TransmittedSymbols
将每个设备的模型参数转换为符号序列,并模拟了符号的传输过程。 - 通过异步传输,不同设备的更新可以在不同的时间进行,而不是等待所有设备都准备好。
- 最后,通过联邦平均化(
FedAvg
函数)将接收到的符号转换为全局模型的权重。
六、BPDecoding.py
这是一个基于概率传播(Probability Propagation,简称BP)算法的通信系统仿真代码。以下是代码中各部分的解释:
per_pkt_transmission 函数
这段代码实现了一个仿真模型,模拟了在一个复杂通信通道上进行数据包传输的过程,并对接收到的信号进行了处理,包括通道效应、噪声的添加,以及对齐样本估计器和最大似然估计器的使用。
1、生成通道延迟和相位:
taus
是随机生成的通道延迟。
hh
是复杂通道的相位,根据参数 args.phaseOffset
的不同值生成。
args.phaseOffset
是一个参数,其作用是确定复杂通道的相位偏移.
当 args.phaseOffset
为 0 时,通道相位保持不变,即相位偏移为0。
当 args.phaseOffset
为 1 时,通道相位在每个设备上随机选择 [0, 2π/4)
范围值。
当 args.phaseOffset
为 2 时,通道相位在每个设备上随机选择 [0, 3π/4)
范围值。
当 args.phaseOffset
为 3 时,通道相位在每个设备上随机选择 [0, 4π/4)
范围值。
2、复杂通道传输:
对每个设备的符号序列 TransmittedSymbols
应用复杂通道效应,即将每个符号乘以对应的相位。
通过在发送端将每个符号乘以对应的相位,模拟了信号在传输过程中的相位变化。在接收端,对于每个设备的信号,可以通过除以对应的相位来进行相位补偿,从而尽可能地还原发送端的相位信息。这是为了在模拟中更真实地反映通信系统中的相位失真效应。
3、接收信号处理:
计算接收信号功率并添加高斯噪声。加入高斯噪声是为了模拟真实通信系统的工作环境。
对接收信号进行过采样,以便进行对齐样本估计器和最大似然估计器的操作。
4、对齐样本估计器:
对齐样本估计器(aligned sample estimator)是用于从接收到的信号中提取特定位置的样本,并基于这些样本进行参数估计或信号检测的一种技术。在数字通信系统中,由于信号经过传输和接收过程中可能存在多径效应、时延扩展等问题,导致接收到的信号在时间上不对齐,因此需要对信号进行处理以提取有效的信息。对齐样本估计器针对特定时刻或位置的信号样本进行分析和处理。
自相关函数表达了同一过程不同时刻的相互依赖关系,而互相关函数表示不同过程的某一时刻的相互依赖关系.互相关函数是描述随机信号X(t),Y(t)在任意两个不同时刻t1,t2,的取值之间的相关程度.自相关函数是描述随机信号X(t)在任意两个不同时刻t1,t2,的取值之间的相关程度。
假设4个设备的原始信号中各包含了5个子序列,则原始信号序列和也有5个子序列:
最终生成的样本序列samples是经过重复、添0、列相加等处理后包含噪声的信号,所以最终的samples样本序列是(包含了23个子序列):
匹配滤波的过程可以分为三个步骤:模板生成、相关运算和判决首先,需要根据已知的信号特征生成一个模板或者滤波器。模板可以是一个固定的函数,也可以是一个由已知信号样本计算得到的函数。然后,将待检测信号与模板进行相关运算,得到相关输出。最后,对相关输出进行判决,根据预先设定的闯值确定信号的存在与否。
yk(t)互相关函数用来寻找输入信号中和模板信号最相似的部分【图像中的峰值】。峰值的位置对应着输入信号中与模板信号对齐的延迟.
WMFS方案中,用M个匹配滤波器对输入信号进行匹配,对齐样本估计器选择第M个滤波器
如图,第M个滤波器从标红处对齐。
当 args.Estimator
为 1 时,使用对齐样本估计器。这个估计器通过选择合适的滤波器索引来进行匹配滤波。??
在数字通信中,匹配滤波是一种常见的信号处理技术,用于检测接收到的信号中的特定信号模式。对齐样本计算器(aligned sample estimator)使用了这一原理,通过选择合适的滤波器索引,进行匹配滤波以提高信号的检测性能。
理解对齐样本计算器通过匹配滤波进行操作的关键点如下:
1. 信号模式的对齐: 在通信系统中,由于多路径传播等原因,接收到的信号可能会发生时延(相位偏移)。对齐样本计算器的目标是将接收到的信号的样本与发送信号的样本进行对齐,以便更准确地检测信号。
2. 滤波器索引的选择:对齐样本计算器使用滤波器来对信号进行处理,滤波器的作用是匹配信号的形状,从而增强信号的特征。选择合适的滤波器索引意味着选择一个与接收到的信号的时延相匹配的滤波器,以便最大化信号的相关性。
3. 匹配滤波:选择的滤波器索引用于对接收到的信号进行匹配滤波。匹配滤波的目的是最大化信号和滤波器的卷积结果,从而在时域上对信号进行最佳匹配。
4. 样本估计:通过匹配滤波,对齐样本计算器得到了每个滤波器索引下的滤波输出。通常,选择具有最大滤波输出的索引,该输出对应于最佳的对齐位置。这个索引即为所谓的“对齐样本”,它表示了在匹配滤波过程中信号的最佳对齐位置。
通过这种方式,对齐样本计算器能够通过选择适当的滤波器索引,对信号进行匹配滤波并估计最佳对齐样本的位置,从而提高信号检测的准确性。这对于在通信系统中处理时延和相位失真等问题非常有用。
Q:与接收到的信号的时延相匹配的滤波器,匹配是什么意思?什么怎么就算相匹配了?
A:
匹配滤波是一种信号处理技术,其目标是在接收到的信号中找到与事先知道的信号模式最相似的部分。在对齐样本计算器的上下文中,匹配滤波用于找到与发送信号模式相匹配的时延位置。
在匹配滤波中,匹配的度量通常使用卷积操作来计算。给定信号序列 x[n] 和滤波器序列 h[n],它们的离散卷积 表示为:
在对齐样本计算器中,我们希望找到一个时延 d,使得卷积 y[n]在时延 d 处达到最大值。这个时延 d 就是信号的对齐样本。
相匹配的意思是,通过调整时延d,我们尽可能地使滤波器 h[n] 与接收到的信号 x[n] 的一部分(位于时延 d 的位置)相似。匹配程度的度量通常使用卷积的峰值来衡量,即找到卷积结果 y[n]的最大值。
因此,当在某个时延 d 处进行卷积时,卷积的结果达到最大值,说明此时的滤波器 h[n] 与接收到的信号 x[n] 在这个时刻相匹配。在对齐样本计算器中,选择使卷积结果最大的时延位置,即为相匹配的对齐样本。
综合来说,匹配滤波的目标是找到一个时延,使得滤波器与接收到的信号在这个时延位置相匹配,使得卷积的峰值最大。这个时延即为对齐样本,表示最佳的对齐位置。
5、最大似然估计器:
当 args.Estimator
为 2 时,使用最大似然估计器。这个估计器尝试在已知噪声和通道条件下找到最可能的信号。计算了相关的矩阵和向量,然后通过求解线性方程组得到估计值。
6、SP_ML 估计器:
当 args.Estimator
为 3 时,使用了 SP_ML 估计器。这个估计器可能是一种基于置信传播算法的估计方法。
BP_Decoding 函数
这段代码实现了一种称为 Belief Propagation(信念传播,BP)的解码算法,用于处理通过信道传输的符号。主要包括右传播和左传播两个过程。
1. 准备观测节点的高斯消息:
通过计算观测节点的 Lambda 矩阵和 Eta 矩阵准备高斯消息。
2. 右向传播:
通过右向传播计算每个变量节点的消息,这涉及到消息 m1、m2 和 m3。
m1 是来自底部的消息,m2 是右侧消息的乘积,m3 是这两者的和。
3. 左向传播:
通过左向传播计算每个变量节点的消息,同样涉及到消息 m1、m2 和 m3。
4. 边缘化和 BP 解码:
通过边缘化计算每个变量节点的最终解码值。
test 函数
1. 生成传输的符号:
对每个设备生成随机的符号,模拟传输的符号。
2. 符号估计与均方误差计算:
使用不同的符号估计器(对齐样本估计器、最大似然估计器、SP-ML 估计器),计算估计的符号与真实符号之间的均方误差(MSE)。
这个代码文件主要涉及到通信系统中的信道建模、符号传输和接收、概率传播解码等方面的模拟。