pysyft对MNIST数据集进行加载训练_sy.federateddataloader

利用pysyft加载MNIST数据集，通过神经网络进行训练，参考官方文档，自己改了些数值，每一行代码基本都注释了对应的含义

导入pytorch
import logging
 
import torch.utils.data
# 用于构建NN
import torch.nn as nn
 
# 需要用到这个库里面的激活函数
import torch.nn.functional as F
 
# 用于构建优化器
import torch.optim as optim
 
# 用于初始化数据
from torchvision import datasets, transforms
 
# 用于分布式训练
import syft as sy
 
 
hook = sy.TorchHook(torch)
Bob = sy.VirtualWorker(hook,id='Bob')
Alice = sy.VirtualWorker(hook,id='Alice')
class Arguments(object):
    def __init__(self):
        self.batch_size = 1
        self.test_batch_size = 1000
        self.epochs = 3
        self.lr = 0.01
        self.momentum = 0.5
        self.no_cuda = False
        self.seed = 1
        self.log_interval = 30
        self.save_model = False
#实例化参数类
args = Arguments()
#判断是否使用GPu
use_cuda = not args.no_cuda and torch.cuda.is_available()
#固定化随机数种子，使得每次训练的随机数都是固定的
torch.manual_seed(args.seed)
#判断是否使用GPU
device = torch.device('cuda' if use_cuda else 'cpu')
#tramsform将图片转化为Tensor，Normalize用来正则化，降低模型复杂度
fed_dataset = datasets.MNIST('data',download=True,train=True,
                            transform=transforms.Compose([transforms.ToTensor()]))
 
#定义数据加载器，shuffle是采用随机的方式抽取数据,顺便也把数据集定义在了客户端上
 
fed_loader = sy.FederatedDataLoader(federated_dataset=fed_dataset.federate((Alice,Bob)),batch_size=args.batch_size,shuffle=True)
# print('Bob\'object: {}'.format(Bob._objects))   #检验分发的数据
# print('Alice\'object: {}'.format(Alice._objects))  #检验分发的数据
#定义测试集
test_dataset = datasets.MNIST('data',download=True,train=False,
                            transform=transforms.Compose([transforms.ToTensor()]))
 
#定义测试集加载器
test_loader = torch.utils.data.DataLoader(test_dataset,batch_size=args.test_batch_size,shuffle=True)
#构建神经网络模型
class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        #输入维度为1，输出维度为20，卷积核大小为：5*5，步幅为1
        self.conv1 = nn.Conv2d(1,20,5,1)   #灰度图片的通道
        self.conv2 = nn.Conv2d(20,50,5,1)   #输入为20，输出为50，卷积核5*5，步幅为1
        self.fc1 = nn.Linear(4*4*50,500)
        #最后映射到10维上
        self.fc2 = nn.Linear(500,10)
 
    def forward(self,x):
        #print(x.shape)
        x = F.relu(self.conv1(x))#28*28*1 -> 24*24*20    （28-5+1=24）
        #print(x.shape)
        #卷机核：2*2 步幅：2
        x = F.max_pool2d(x,2,2)#24*24*20 -> 12*12*20  池化，通道不变，维度降半
        #print(x.shape)
        x = F.relu(self.conv2(x))#12*12*20 -> 8*8*50      （12-5+1=8）
        #   #print(x.shape)
        x = F.max_pool2d(x,2,2)#8*8*50 -> 4*4*50
        #print(x.shape)
        x = x.view(-1,4*4*50)     #拉伸  -1自动计算维度
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        #使用logistic函数作为softmax进行激活吗就
        return F.log_softmax(x, dim = 1)   #计算分类够每个数字的概率值
 
#定义数据训练:
model = Net().to(device)
optimizer = optim.SGD(model.parameters(), lr=args.lr)
def train(model:model, device:device, fed_loader:sy.FederatedDataLoader, optimizer:optimizer, epoch:args.epochs):
    model.train()
    # 远程迭代
    for batch_idx, (data,target) in enumerate(fed_loader):  # enumerate用来编序号
        model.send(data.location)  # 发送模型到远程
        data, target = data.to(device), target.to(device)    #部署到device    #此行有问题
        optimizer.zero_grad()  #梯度清0
        output = model(data)
        loss = F.nll_loss(output, target)   #计算损失
        loss.backward()   #反向传播
        optimizer.step()
        model.get()   # 以上都是发送命令给远程，下面是取回更新的模型
        if batch_idx % args.log_interval== 0:  # 打印间隔时间
          # 由于损失也是在远处产生的，因此我们需要把它取回来
          loss=loss.get()
          print('Train Epoch : {} [ {} / {} ({:.0f}%)] \tLoss: {:.6f}'.format(
            epoch,
            batch_idx * args.batch_size,
            len(fed_loader)*args.batch_size,
            100.*batch_idx/len(fed_loader),
            loss.item())
            )
def test(model:model, device:device, test_loader:torch.utils.data.DataLoader):
    model.eval()
    test_loss = 0 #测试损失
    correct=0 #正确率
    with torch.no_grad():
        for data,target in test_loader:
            data,target = data.to(device),target.to(device)
            output=model(data)
            #将损失加起来
            test_loss+=F.nll_loss(output,target,reduction='sum').item()
 
            #进行预测最可能的分类
            pred =output.argmax(dim=1,keepdim=True)
            correct+=pred.eq(target.view_as(pred)).sum().item()
            test_loss/=len(test_loader.dataset)
            print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
                test_loss, correct, len(test_loader.dataset),
                100. * correct / len(test_loader.dataset)))
 
logging.info("开始训练!!\n")
for epoch in range(1, args.epochs + 1):
    train(model, device, fed_loader, optimizer, epoch)
    test(model, device, test_loader)

训练效果还不错