-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathCryptoNets.cpp
195 lines (174 loc) · 11.2 KB
/
CryptoNets.cpp
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
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
//
// Created by 邓维维 on 2021/12/6.
//
#include "CryptoNets.h"
dww::CryptoNets::CryptoNets(const torch::Tensor &conv1_w, const torch::Tensor &conv1_b, int64_t conv1_in,
int64_t conv1_out, int64_t conv1_k, int64_t conv1_p, int64_t conv1_s, int64_t pool1_k,
int64_t pool1_p, int64_t pool1_s, const torch::Tensor &conv2_w,
const torch::Tensor &conv2_b, int64_t conv2_in, int64_t conv2_out, int64_t conv2_k,
int64_t conv2_p, int64_t conv2_s, int64_t pool2_k, int64_t pool2_p, int64_t pool2_s,
const torch::Tensor &l1_w, const torch::Tensor &l1_b, int64_t l1_in, int64_t l1_out,
const torch::Tensor &l2_w, const torch::Tensor &l2_b, int64_t l2_in, int64_t l2_out)
:
conv1(conv1_w,conv1_b,conv1_in,conv1_out,conv1_k,conv1_p,conv1_s),
pool1(pool1_k,pool1_p,pool1_s),
conv2(conv2_w,conv2_b,conv2_in,conv2_out,conv2_k,conv2_p,conv2_s),
pool2(pool2_k,pool2_p,pool2_s),
linear1(l1_w,l1_b,l1_in,l1_out),
linear2(l2_w,l2_b,l2_in,l2_out)
{
}
void dww::CryptoNets::forward(const torch::Tensor &input, dww::HEWrapper &tools, dww::Cipher_Tensor &output) {
assert(input.sizes().size() == 4 && "The input image is not a 2D image, its shape must be N * C * H * W!");
std::chrono::high_resolution_clock::time_point start,end;
start = std::chrono::high_resolution_clock::now();
Cipher_Tensor input_cipher(input,tools);
end = std::chrono::high_resolution_clock::now();
enc_time += std::chrono::duration_cast<std::chrono::duration<double>>(end - start).count();
// 获取当前同态运算的批处理的 image 图片数量，以用于后面的 Cipher_Tensor 中的 batch 的设置
int64_t batch = input_cipher.batch;
Cipher_Tensor conv1_output(conv1.out_numel(input_cipher.shape),conv1.out_shape(input_cipher.shape),batch);
Cipher_Tensor pool1_output(pool1.out_numel(conv1_output.shape),pool1.out_shape(conv1_output.shape),batch);
Cipher_Tensor conv2_output(conv2.out_numel(pool1_output.shape),conv2.out_shape(pool1_output.shape),batch);
Cipher_Tensor pool2_output(pool2.out_numel(conv2_output.shape),pool2.out_shape(conv2_output.shape),batch);
Cipher_Tensor linear1_output(linear1.out_,{linear1.out_},batch);
start = std::chrono::high_resolution_clock::now();
conv1.forward(input_cipher,tools,conv1_output);
end = std::chrono::high_resolution_clock::now();
conv_time += std::chrono::duration_cast<std::chrono::duration<double>>(end - start).count();
start = std::chrono::high_resolution_clock::now();
square.forward(conv1_output,tools);
end = std::chrono::high_resolution_clock::now();
square_time += std::chrono::duration_cast<std::chrono::duration<double>>(end - start).count();
start = std::chrono::high_resolution_clock::now();
pool1.forward(conv1_output,tools,pool1_output);
end = std::chrono::high_resolution_clock::now();
pool_time += std::chrono::duration_cast<std::chrono::duration<double>>(end - start).count();
start = std::chrono::high_resolution_clock::now();
conv2.forward(pool1_output,tools,conv2_output);
end = std::chrono::high_resolution_clock::now();
conv_time += std::chrono::duration_cast<std::chrono::duration<double>>(end - start).count();
start = std::chrono::high_resolution_clock::now();
pool2.forward(conv2_output,tools,pool2_output);
end = std::chrono::high_resolution_clock::now();
pool_time += std::chrono::duration_cast<std::chrono::duration<double>>(end - start).count();
start = std::chrono::high_resolution_clock::now();
pool2_output.shape = {pool2_output.numel()};
linear1.forward(pool2_output,tools,linear1_output);
end = std::chrono::high_resolution_clock::now();
linear_time += std::chrono::duration_cast<std::chrono::duration<double>>(end - start).count();
start = std::chrono::high_resolution_clock::now();
square.forward(linear1_output,tools);
end = std::chrono::high_resolution_clock::now();
square_time += std::chrono::duration_cast<std::chrono::duration<double>>(end - start).count();
start = std::chrono::high_resolution_clock::now();
linear2.forward(linear1_output,tools,output);
end = std::chrono::high_resolution_clock::now();
linear_time += std::chrono::duration_cast<std::chrono::duration<double>>(end - start).count();
}
void dww::CryptoNets_he_inference_test(const std::string& filename){
std::string path = "../model/cryptonets/" + filename + "/";
torch::Tensor conv1_weight,conv1_bias;
torch::Tensor conv2_weight,conv2_bias;
torch::Tensor linear1_weight,linear1_bias;
torch::Tensor linear2_weight,linear2_bias;
torch::load(conv1_weight,path + "conv1_weight.pt");
conv1_weight = conv1_weight.to(c10::kCPU);
torch::load(conv1_bias,path + "conv1_bias.pt");
conv1_bias = conv1_bias.to(c10::kCPU);
torch::load(conv2_weight,path + "conv2_weight.pt");
conv2_weight = conv2_weight.to(c10::kCPU);
torch::load(conv2_bias,path + "conv2_bias.pt");
conv2_bias = conv2_bias.to(c10::kCPU);
torch::load(linear1_weight,path + "linear1_weight.pt");
linear1_weight = linear1_weight.to(c10::kCPU);
torch::load(linear1_bias,path + "linear1_bias.pt");
linear1_bias = linear1_bias.to(c10::kCPU);
torch::load(linear2_weight,path + "linear2_weight.pt");
linear2_weight = linear2_weight.to(c10::kCPU);
torch::load(linear2_bias,path + "linear2_bias.pt");
linear2_bias = linear2_bias.to(c10::kCPU);
int64_t conv1_in = 1,conv1_out = 5, conv1_k = 5, conv1_s = 2,conv1_p = 1;
int64_t conv2_in = 5,conv2_out = 50, conv2_k = 5, conv2_s = 2,conv2_p = 0;
int64_t pool1_k = 3,pool1_s = 1,pool1_p = 1;
int64_t pool2_k = 3,pool2_s = 1,pool2_p = 1;
int64_t l1_in = 50*5*5,l1_out = 100;
int64_t l2_in = 100,l2_out = dww::label_num.at(filename);
CryptoNets model(conv1_weight,conv1_bias,conv1_in,conv1_out,conv1_k,conv1_p,conv1_s,
pool1_k,pool1_p,pool1_s,
conv2_weight,conv2_bias,conv2_in,conv2_out,conv2_k,conv2_p,conv2_s,
pool2_k,pool2_p,pool2_s,
linear1_weight,linear1_bias,
l1_in,l1_out,
linear2_weight,linear2_bias,
l2_in,l2_out
);
HEWrapper tools(16384,30);
dww::MedDataSet dataset(filename);
dww::MedDataSetLoader dataloader(dataset,DATA_CAT::TEST,tools.get_slots_num());
int64_t sz = dataloader.get_batch_num();
int64_t samples_num = dataloader.samples_num;
int64_t correct = 0;
std::ofstream test_log("../experiment/cryptonets",std::ios_base::app);
assert(test_log.is_open() && "File cryptonets can not open!");
test_log << "Dataset: " << filename << '\n';
test_log << "Model Information: \n";
test_log << model;
std::cout << "----> CryptoNets Homomorphic Convolution " << filename << " Datasets Starts <-----\n";
using std::chrono::high_resolution_clock;
high_resolution_clock::time_point start,end;
double time_consume = 0;
for(int64_t i = 0; i < sz; ++i){
torch::Tensor image = dataloader.images[i];
torch::Tensor label = dataloader.labels[i];
// 该 batch 中有的 image 图片个数
int64_t bt_sz = image.size(0);
// 保存最终预测结果的密文值
Cipher_Tensor output(model.linear2.out_,{model.linear2.out_},bt_sz);
start = high_resolution_clock::now();
// 同态运算一个卷积单元
model.forward(image,tools,output);
end = high_resolution_clock::now();
time_consume += std::chrono::duration_cast<std::chrono::duration<double>>(end - start).count();
start = high_resolution_clock::now();
// 对预测结果进行解密，获得每一个 image 的预测结果
std::vector<std::vector<double>> res(output.get_message_of_tensor(tools));
end = high_resolution_clock::now();
model.dec_time += std::chrono::duration_cast<std::chrono::duration<double>>(end - start).count();
// 记录每一个 image 的预测结果，使用下标的方式反映(某一个下标处值最大便是保存该下标值)
std::vector<int64_t> max_index(bt_sz);
for(int64_t j = 0; j < bt_sz; ++j)
max_index[j] = std::max_element(res[j].begin(),res[j].end()) - res[j].begin();
torch::Tensor temp = torch::tensor(max_index,torch::TensorOptions(torch::kInt64));
assert(temp.numel() == label.numel() && "The prediction number is not equal to labels!");
correct += (bt_sz - (temp - label).count_nonzero().item<int64_t>());
std::cout << "Dataset " << filename << ",batch " << i + 1 << " completed!\n";
test_log << "\t[" << i + 1 << "/" << sz << "], Acc = " << static_cast<double>(correct) / samples_num * 100 << "%\n";
}
test_log << "Total Inference Images : " << dataloader.samples_num << "\n";
test_log << "Total Acc : " << static_cast<double>(correct) / samples_num * 100 << "%\n";
test_log << "Total Encryption Operation Time Consume : " << model.enc_time << "(s)\n";
test_log << "Average Encryption Operation Time Consume : " << model.enc_time / samples_num << "(s)\n";
test_log << "Total Convolution Operation Time Consume : " << model.conv_time << "(s)\n";
test_log << "Average Convolution Operation Time Consume : " << model.conv_time / samples_num << "(s)\n";
test_log << "Total Pooling Operation Time Consume : " << model.pool_time << "(s)\n";
test_log << "Average Pooling Operation Time Consume : " << model.pool_time / samples_num << "(s)\n";
test_log << "Total Square Activation Operation Time Consume : " << model.square_time << "(s)\n";
test_log << "Average Square Activation Operation Time Consume : " << model.square_time / samples_num << "(s)\n";
test_log << "Total Linear Operation Time Consume : " << model.linear_time << "(s)\n";
test_log << "Average Square Activation Operation Time Consume : " << model.linear_time / samples_num << "(s)\n";
test_log << "Total Decryption Operation Time Consume : " << model.dec_time << "(s)\n";
test_log << "Average Decryption Operation Time Consume : " << model.dec_time / samples_num << "(s)\n";
test_log << "Total Time Consume : " << time_consume << "(s)\n";
test_log << "Average Time Consume Per Batch : " << time_consume / sz << "(s)\n";
test_log << "Average Time Consume Per Image : " << time_consume / samples_num << "(s)\n\n";
test_log.flush();
test_log.close();
std::cout << "----> CryptoNets Homomorphic Convolution " << filename << " Datasets End <-----\n";
}
std::ostream& dww::operator<<(std::ostream& out,const CryptoNets& self){
out << "conv1: " << self.conv1 << "pool1: " << self.pool1
<< "conv2: " << self.conv2 << "pool2: " << self.pool2
<< "linear1: " << self.linear1 << "linear2: " << self.linear2;
return out;
}