联邦学习FedAvg自编写代码

联邦学习中，联邦平均算法获得了很大的使用空间，因此常常被用于进行同步训练操作
不多废话了，以下为Fedavg代码
由于使用场景为NonIID场景，因此我使用了别人的一个MNIST数据集自定义的代码（见附录）
FedAvg代码如下，功能具体看注释
工作环境：python3.8.5 + pytorch(无cuda)
divergence模块可直接删除

# coding: utf-8

# In[1]:


import argparse
import torch
import os
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torch.autograd import Variable
from PIL import Image
import torch
import copy
import pandas as pd
import random
import time
import sys
import re
import matplotlib.pyplot as plt
#import divergence

name = str(sys.argv[0])

# In[2]:

home_path = "./"
class MyDataset(torch.utils.data.Dataset): #创建自己的类：MyDataset,这个类是继承的torch.utils.data.Dataset
    def __init__(self,root,data,label,transform=None, target_transform=None): #初始化一些需要传入的参数
        super(MyDataset,self).__init__()
        imgs = []                      #创建一个名为img的空列表，一会儿用来装东西
        self.img_route = root
        for i in range(len(data)):
            imgs.append((data[i],int(label[i])))
        self.imgs = imgs
        self.transform = transform
        self.target_transform = target_transform
 
    def __getitem__(self, index):    #这个方法是必须要有的，用于按照索引读取每个元素的具体内容
        fn, label = self.imgs[index] #fn是图片path #fn和label分别获得imgs[index]也即是刚才每行中word[0]和word[1]的信息
        route = self.img_route + str(label) + "/" + fn
        img = Image.open(route) #按照path读入图片from PIL import Image # 按照路径读取图片
        if self.transform is not None:
            img = self.transform(img) #是否进行transform
        return img,label  #return很关键，return回哪些内容，那么我们在训练时循环读取每个batch时，就能获得哪些内容
 
    def __len__(self): #这个函数也必须要写，它返回的是数据集的长度，也就是多少张图片，要和loader的长度作区分
        return len(self.imgs)


# In[3]:


filePath = home_path + 'data/MNIST/image_turn/'
train_data = []
train_label = []
for i in range(10):
    train_data.append(os.listdir(filePath+str(i)))
    train_label.append([i]*len(train_data[i]))
filePath = home_path + 'data/MNIST/image_test_turn/'
test_data = []
test_label = []
for i in range(10):
    test_data.append(os.listdir(filePath+str(i)))
    test_label.append([i]*len(test_data[i]))
test_ori = []
test_label_ori = []
for x in range(10):
    test_ori += test_data[x]
    test_label_ori += test_label[x]
test_data=MyDataset(home_path + "data/MNIST/image_test_turn/",test_ori,test_label_ori, transform=transforms.ToTensor())
test_loader = DataLoader(dataset=test_data, batch_size=64)


# In[4]:


class MyConvNet(nn.Module):
    def __init__(self):
        super(MyConvNet,self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(
                in_channels = 1,
                out_channels = 16,
                kernel_size = 3,
                stride = 1,
                padding = 1,
            ),
            nn.ReLU(),
            nn.AvgPool2d(
                kernel_size = 2,
                stride = 2)
        )

        self.conv2 = nn.Sequential(
            nn.Conv2d(16,32,3,1,0),
            nn.ReLU(),
            nn.MaxPool2d(2,2),
        )
        self.classifier = nn.Sequential(
            nn.Linear(32*6*6,256),
            nn.ReLU(),
            nn.Linear(256,128),
            nn.ReLU(),
            nn.Linear(128,10)
        )

    def forward(self,x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0),-1)
        output = self.classifier(x)
        return  output


# In[5]:


def train_model(model,traindataloader,criterion,optimizer,batch_max,num_epochs):
    train_loss_all = []
    train_acc_all = []
    for epoch in range(num_epochs):
        train_loss = 0.0
        train_corrects = 0
        train_num = 0
        temp = random.sample(traindataloader, batch_max)
        for (b_x,b_y) in temp:
            model.train()
            if(torch.cuda.is_available()):
                b_x = b_x.cuda()
                b_y = b_y.cuda()
            output = model(b_x)
            pre_lab = torch.argmax(output,1)
            loss =  criterion(output, b_y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            train_loss += loss.item()
            train_corrects += torch.sum(pre_lab == b_y.data)
            train_num += b_x.size(0)
        train_loss_all.append(train_loss / train_num)
        train_acc_all.append(train_corrects.double().item()/train_num)
        print("Train Loss:{:.4f}  Train Acc: {:.4f}" .format(train_loss_all[-1],train_acc_all[-1]))
    return model


# In[6]:


def local_train(local_convnet_dict,traindataloader,epochs,batch_max):
    if(torch.cuda.is_available()):
        local_convnet = MyConvNet().cuda()
    else:
        local_convnet = MyConvNet()
    local_convnet.load_state_dict(local_convnet_dict)
    optimizer = optim.Adam(local_convnet.parameters(), lr=0.01, weight_decay=5e-4)
    criterion = nn.CrossEntropyLoss()
    local_convnet = train_model(local_convnet,traindataloader, criterion, optimizer, batch_max, epochs)
    minus_convnet_dict = MyConvNet().state_dict()
    for name in local_convnet.state_dict():
        minus_convnet_dict[name] = local_convnet_dict[name] - local_convnet.state_dict()[name]
    return minus_convnet_dict


# In[7]:


def Central_model_update(Central_model, minus_convnet_client):
    weight = 1

    model_dict = Central_model.state_dict()
    for name in Central_model.state_dict():
        for local_dict in minus_convnet_client:
            model_dict[name] = model_dict[name] - weight * local_dict[name]/len(minus_convnet_client)
    Central_model.load_state_dict(model_dict)
    return Central_model


# In[12]:


def train_data_loader(client_num, ClientSort1):
    global train_data
    global train_label
    train_loaders = []
    for i in range(client_num):
        train_ori = []
        label_ori = []
        for j in range(10):
            train_ori += train_data[j]
            label_ori += train_label[j]
        train_datas=MyDataset(home_path + "data/MNIST/image_turn/",train_ori,label_ori, transform=transforms.ToTensor())
        train_loader = DataLoader(dataset=train_datas, batch_size=100, shuffle=True)
        train_list = []
        for step,(b_x,b_y) in enumerate(train_loader):
            train_list.append((b_x,b_y))
        train_loaders.append(train_list)
    return train_loaders
'''
def train_data_loader(client_num, ClientSort1):
    global train_data
    global train_label
    train_loaders = []
    for i in range(client_num):
        if i < ClientSort1: #同步部分
            train_ori = []
            label_ori = []
            for j in range(10):
                if(i!=j):
                    train_ori += train_data[j]
                    label_ori += train_label[j]
        else: #异步部分
            train_ori = []
            label_ori = []
            for j in range(10):
                if(j == i-10 ):
                    train_ori += train_data[j]
                    label_ori += train_label[j]
        train_datas = MyDataset(home_path + "data/MNIST/image_turn/",train_ori,label_ori, transform=transforms.ToTensor())
        train_loader = DataLoader(dataset=train_datas, batch_size=128, shuffle=True)
        train_list = []
        for step,(b_x,b_y) in enumerate(train_loader):
            train_list.append((b_x,b_y))
        train_loaders.append(train_list)
    return train_loaders
'''
# In[13]:

def test_accuracy(Central_model):
    global test_loader
    test_correct = 0
    for data in test_loader:
        Central_model.eval()
        inputs, lables = data
        if(torch.cuda.is_available()):
            inputs = inputs.cuda()
        inputs, lables = Variable(inputs), Variable(lables)
        outputs = Central_model(inputs)
        if(torch.cuda.is_available()):
            outputs = outputs.cpu()
        _, id = torch.max(outputs.data, 1)
        test_correct += torch.sum(id == lables.data)
        test_correct = test_correct
    print("correct:%.3f%%" % (100 * test_correct / len(test_ori)))
    return 100 * test_correct / len(test_ori)

# In[14]:


############################
##      中央共享模型      ##
############################
if(torch.cuda.is_available()):
    Central_model = MyConvNet().cuda()
else:
    Central_model = MyConvNet()
local_client_num = 10 #局部客户端数量
ClientSort1 = 10
#Central_model.load_state_dict(torch.load('F:/params.pkl'))
global_epoch = 1000
#print(test_accuracy(Central_model))


# In[ ]:


train_loaders = train_data_loader(local_client_num, ClientSort1)
result = []
count = 0
for i in range(global_epoch):
    count += 1
    minus_model = []
    for j in range(local_client_num):
        minus_model.append(local_train(Central_model.state_dict(),train_loaders[j],1,1))
    Central_model = Central_model_update(Central_model, minus_model)
    print("epoch: ",count, "\naccuracy:")
    result.append(float(test_accuracy(Central_model)))
    #divergence.cal_divergence(Central_model.state_dict(), i, "F:/同步Fig/")


# In[ ]:


plt.xlabel('round')
plt.ylabel('accuracy')
plt.plot(range(0,len(result)), result, color='r',linewidth='1.0', label='同步FedAvg')
plt.savefig(home_path + name + ".jpg")
filename = open(home_path + name + ".txt", mode='w')
for namet in result:
    filename.write(str(namet))
    filename.write('\n')
filename.close()
torch.save(Central_model.state_dict(),filePath + name + '.pkl')

附录：

import numpy as np
import struct
 
from PIL import Image
import os
 
data_file = './data/MNIST/raw/train-images-idx3-ubyte' #需要修改的路径

data_file_size = 47040016
data_file_size = str(data_file_size - 16) + 'B'
 
data_buf = open(data_file, 'rb').read()
 
magic, numImages, numRows, numColumns = struct.unpack_from(
    '>IIII', data_buf, 0)
datas = struct.unpack_from(
    '>' + data_file_size, data_buf, struct.calcsize('>IIII'))
datas = np.array(datas).astype(np.uint8).reshape(
    numImages, 1, numRows, numColumns)
 
label_file = './data/MNIST/raw/train-labels-idx1-ubyte' #需要修改的路径
 
# It's 60008B, but we should set to 60000B
label_file_size = 60008
label_file_size = str(label_file_size - 8) + 'B'
 
label_buf = open(label_file, 'rb').read()
 
magic, numLabels = struct.unpack_from('>II', label_buf, 0)
labels = struct.unpack_from(
    '>' + label_file_size, label_buf, struct.calcsize('>II'))
labels = np.array(labels).astype(np.int64)
 
datas_root = './data/MNIST/image_turn/' #需要修改的路径
if not os.path.exists(datas_root):
    os.mkdir(datas_root)
 
for i in range(10):
    file_name = datas_root + os.sep + str(i)
    if not os.path.exists(file_name):
        os.mkdir(file_name)
 
for ii in range(numLabels):
    img = Image.fromarray(datas[ii, 0, 0:28, 0:28])
    label = labels[ii]
    file_name = datas_root + os.sep + str(label) + os.sep + \
        'mnist_train_' + str(ii) + '.png'
    img.save(file_name)

import numpy as np
import struct
 
from PIL import Image
import os
 
data_file = './data/MNIST/raw/t10k-images-idx3-ubyte' #需要修改的路径
 

data_file_size = 7840016
data_file_size = str(data_file_size - 16) + 'B'
 
data_buf = open(data_file, 'rb').read()
 
magic, numImages, numRows, numColumns = struct.unpack_from(
    '>IIII', data_buf, 0)
datas = struct.unpack_from(
    '>' + data_file_size, data_buf, struct.calcsize('>IIII'))
datas = np.array(datas).astype(np.uint8).reshape(
    numImages, 1, numRows, numColumns)
 
label_file = './data/MNIST/raw/t10k-labels-idx1-ubyte'#需要修改的路径
 
# It's 10008B, but we should set to 10000B
label_file_size = 10008
label_file_size = str(label_file_size - 8) + 'B'
 
label_buf = open(label_file, 'rb').read()
 
magic, numLabels = struct.unpack_from('>II', label_buf, 0)
labels = struct.unpack_from(
    '>' + label_file_size, label_buf, struct.calcsize('>II'))
labels = np.array(labels).astype(np.int64)
 
datas_root = './data/MNIST/image_test_turn/' #需要修改的路径
 
if not os.path.exists(datas_root):
    os.mkdir(datas_root)
 
for i in range(10):
    file_name = datas_root + os.sep + str(i)
    if not os.path.exists(file_name):
        os.mkdir(file_name)
 
for ii in range(numLabels):
    img = Image.fromarray(datas[ii, 0, 0:28, 0:28])
    label = labels[ii]
    file_name = datas_root + os.sep + str(label) + os.sep + \
        'mnist_test_' + str(ii) + '.png'
    img.save(file_name)