1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
|
import torch
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
import logging
import syft as sy
import torchvision
hook = sy.TorchHook(torch)
qin = sy.VirtualWorker(hook=hook, id="qin")
zheng = sy.VirtualWorker(hook=hook, id="zheng")
#设定参数
args = {
'use_cuda': True,
'batch_size': 64,
'test_batch_size': 1000,
'lr': 0.01,
'log_interval': 10,
'epochs': 10,
'save_model': False
}
use_cuda = args['use_cuda'] and torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
# 我们首先加载数据，然后使用。联邦方法。
# 此联邦数据集现在提供给 Federated DataLoader。测试数据集保持不变。
# 下面是训练数据，需要分发给远处打工人
federated_train_loader = sy.FederatedDataLoader(
datasets.MNIST('./mnist_data', train=True, download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
]))
.federate((qin, zheng)),
batch_size=args['batch_size'], shuffle=True
)
# federate()函数已经实现了分发（省去了我们一个一个send()的麻烦）
# 下面是测试数据，在我们本地
test_loader = DataLoader(
datasets.MNIST('./mnist_data', train=False,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])),
batch_size=args['batch_size'], shuffle=True
)
# CNN网络
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv = nn.Sequential(
nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, stride=1),
#输出26*26*32
nn.ReLU(),
nn.Conv2d(in_channels=32, out_channels=64,kernel_size=3, stride=1),
#输出24*24*64
nn.ReLU()
)
self.fc = nn.Sequential(
nn.Linear(in_features=64*12*12, out_features=128),
nn.ReLU(),
nn.Linear(in_features=128, out_features=10),
)
self.dropout = nn.Dropout2d(0.25) # 随机丢弃
def forward(self, x):
x = self.conv(x)#输入的时候是28*28*1,输出应该是24*24*64
x = F.max_pool2d(x, 2)#用步长为2的池化,输出12*12*64
x = x.view(-1, 64*12*12)#此时将其拉成一条直线进入全连接
x = self.fc(x)
x = F.log_softmax(x, dim=1)
return x
def train(args, model, device, train_loader, optimizer, epoch):
model.train()
# 远程迭代
# for batch_idx, (data, target) in enumerate(federated_train_loader):
# if batch_idx < 5:
# print(batch_idx, type(data), data.location)
# else:
# break
# 0 <class 'torch.Tensor'> < VirtualWorker id:qin # objects:4>
# 1 <class 'torch.Tensor'> < VirtualWorker id:qin # objects:4>
# 2 <class 'torch.Tensor'> < VirtualWorker id:qin # objects:4>
# 3 <class 'torch.Tensor'> < VirtualWorker id:qin # objects:4>
# 4 <class 'torch.Tensor'> < VirtualWorker id:qin # objects:4>
for batch_idx, (data, target) in enumerate(train_loader): # enumrate用来编序号
# model = model.send(data.location) # 发送模型到远程
model.send(data.location)
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
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:{}[{}/{}({:.06f}%)]\tLoss:{:.06f}'.format(
epoch,
batch_idx * args['batch_size'],
len(train_loader) * args['batch_size'],
100. * batch_idx / len(train_loader),
loss.item()
))
def test(model, device, test_loader):
model.eval()
'''返回model的返回值以字符串显示,使用PyTorch进行训练和测试时
一定注意要把实例化的model指定train/eval，eval（）时，
框架会自动把BN和DropOut固定住，不会取平均，而是用训练好的值，不然的话，
一旦test的batch_size过小，很容易就会被BN层导致生成图片颜色失真极大！！'''
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()
'''nll_loss的解释请看
https://blog.csdn.net/qq_22210253/article/details/85229988
和https://www.cnblogs.com/ranjiewen/p/10059490.html'''
#进行预测最可能的分类
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)))
model = Net().to(device)
optimizer = optim.SGD(model.parameters(), lr=args['lr'])
logging.info("开始训练!!\n")
for epoch in range(1, args['epochs']+1):
train(args, model, device, federated_train_loader, optimizer, epoch)
test(model, device, test_loader)
# if (args["save_model"]):
torch.save(model.state_dict(), "mnist_cnn.pt")
|
