Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement im2col and update conv2d #134

Merged
merged 10 commits into from
Aug 10, 2020
4 changes: 4 additions & 0 deletions tenseal/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,9 @@
RelinKeys = _ts_cpp.RelinKeys
GaloisKeys = _ts_cpp.GaloisKeys

# utils
im2col_encoding = _ts_cpp.im2col_encoding


def context(
scheme, poly_modulus_degree, plain_modulus=None, coeff_mod_bit_sizes=None, n_threads=None
Expand Down Expand Up @@ -80,5 +83,6 @@ def context_from(buff, n_threads=None):
"ckks_vector_from",
"context",
"context_from",
"im2col_encoding",
"__version__",
]
15 changes: 15 additions & 0 deletions tenseal/binding.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
#include "tenseal/cpp/context/tensealcontext.h"
#include "tenseal/cpp/tensors/bfvvector.h"
#include "tenseal/cpp/tensors/ckksvector.h"
#include "tenseal/cpp/tensors/utils/utils.h"

using namespace tenseal;
using namespace seal;
Expand Down Expand Up @@ -81,6 +82,20 @@ PYBIND11_MODULE(_tenseal_cpp, m) {
.def("__deepcopy__",
[](const BFVVector &self, py::dict) { return self.deepcopy(); });

// CKKSVector utils
m.def("im2col_encoding",
[](shared_ptr<TenSEALContext> ctx, vector<vector<double>> &input,
const size_t kernel_n_rows, const size_t kernel_n_cols,
const size_t stride) {
vector<vector<double>> view_as_window;
vector<double> final_vector;
size_t windows_nb = im2col(input, view_as_window, kernel_n_rows,
kernel_n_cols, stride);
vertical_scan(view_as_window, final_vector);
CKKSVector ckks_vector = CKKSVector(ctx, final_vector);
return make_pair(ckks_vector, windows_nb);
});

Comment on lines +85 to +98
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe not in this PR, but the encoding shouldn't be doing encryption, just encoding it into a vector, then encryption should be something separate

py::class_<CKKSVector>(m, "CKKSVector")
// specifying scale
.def(py::init<shared_ptr<TenSEALContext> &, vector<double>, double>())
Expand Down
36 changes: 22 additions & 14 deletions tenseal/cpp/tensors/ckksvector.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -564,39 +564,47 @@ CKKSVector& CKKSVector::polyval_inplace(const vector<double>& coefficients) {
return *this;
}

CKKSVector CKKSVector::conv2d_im2col(const vector<double>& kernel,
size_t windows_nb) {
CKKSVector CKKSVector::conv2d_im2col(const vector<vector<double>>& kernel,
const size_t windows_nb) {
CKKSVector new_vec = *this;
new_vec.conv2d_im2col_inplace(kernel, windows_nb);
return new_vec;
}

CKKSVector& CKKSVector::conv2d_im2col_inplace(const vector<double>& kernel,
const size_t windows_nb) {
vector<double> plain_vec;
size_t chunks_nb = kernel.size();

CKKSVector& CKKSVector::conv2d_im2col_inplace(
const vector<vector<double>>& kernel, const size_t windows_nb) {
if (windows_nb == 0) {
throw invalid_argument("Windows number can't be zero");
}

if (kernel.empty()) {
throw invalid_argument("Kernel vector can't be empty");
if (kernel.empty() ||
(any_of(kernel.begin(), kernel.end(),
[](const vector<double>& i) { return i.empty(); }))) {
throw invalid_argument("Kernel matrix can't be empty");
Comment on lines +580 to +583
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What about something like this? to check whether the vector is a well organized matrix?

Suggested change
if (kernel.empty() ||
(any_of(kernel.begin(), kernel.end(),
[](const vector<double>& i) { return i.empty(); }))) {
throw invalid_argument("Kernel matrix can't be empty");
if (kernel.empty() ||
(any_of(kernel.begin(), kernel.end(),
[](const vector<double>& i) { return i.size() == kernel[0].size(); }))) {
throw invalid_argument("Kernel matrix can't be empty");

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The test will be wrong in the case where kernel[0].size() == 0.
in other hand horizontal_scan always check if the rows have the same size.

}

// check if vector size is not a power of 2
if (!(chunks_nb && (!(chunks_nb & (chunks_nb - 1))))) {
throw invalid_argument("Kernel size should be a power of 2");
}
// flat the kernel
vector<double> flatten_kernel;
horizontal_scan(kernel, flatten_kernel);

// calculate the next power of 2
size_t kernel_size = kernel.size() * kernel[0].size();
kernel_size = 1 << (static_cast<size_t>(ceil(log2(kernel_size))));

// pad the kernel with zeros to the next power of 2
flatten_kernel.resize(kernel_size, 0);

size_t chunks_nb = flatten_kernel.size();

if (this->_size / windows_nb != chunks_nb) {
throw invalid_argument("Matrix shape doesn't match with vector size");
}

vector<double> plain_vec;
plain_vec.reserve(this->_size);

for (size_t i = 0; i < chunks_nb; i++) {
vector<double> tmp(windows_nb, kernel[i]);
vector<double> tmp(windows_nb, flatten_kernel[i]);
plain_vec.insert(plain_vec.end(), tmp.begin(), tmp.end());
}

Expand Down
7 changes: 4 additions & 3 deletions tenseal/cpp/tensors/ckksvector.h
Original file line number Diff line number Diff line change
Expand Up @@ -128,9 +128,10 @@ class CKKSVector {
* The input matrix should be encoded in a vertical scan (column-major).
* The kernel vector should be padded with zeros to the next power of 2
*/
CKKSVector conv2d_im2col(const vector<double>& kernel, size_t windows_nb);
CKKSVector& conv2d_im2col_inplace(const vector<double>& kernel,
size_t windows_nb);
CKKSVector conv2d_im2col(const vector<vector<double>>& kernel,
const size_t windows_nb);
CKKSVector& conv2d_im2col_inplace(const vector<vector<double>>& kernel,
const size_t windows_nb);

/**
* Load/Save the vector from/to a serialized protobuffer.
Expand Down
103 changes: 103 additions & 0 deletions tenseal/cpp/tensors/utils/utils.h
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
#ifndef TENSEAL_UTILS_UTILS_H
#define TENSEAL_UTILS_UTILS_H

#include <algorithm>
#include <iterator>
#include <memory>

#include "seal/seal.h"
Expand All @@ -10,6 +12,107 @@ namespace tenseal {
using namespace seal;
using namespace std;

/**
* horizontally scan matrix (vector of vectors)
**/
template <typename T>
void horizontal_scan(const vector<vector<T>>& src, vector<T>& dst) {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These operations might be easier to execute using a gsl::multi_span structure. For sure not for this PR, I don't know if they actually help. But they should prevent extra-copies

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I will check this out, maybe they will be helpful.

size_t in_height = src.size();
size_t in_width = src[0].size();

dst.resize(in_height * in_width);

// check if each row size is equals to in_width
if (any_of(src.begin(), src.end(), [in_width](const vector<T>& i) {
return i.size() != in_width;
})) {
throw invalid_argument("rows sizes are different");
}

auto start = src.begin();
auto end = src.end();
auto iter = dst.begin();
while (start != end) {
iter = copy(start->begin(), start->end(), iter);
start++;
}
}

/**
* vertically scan matrix (vector of vectors)
**/
template <typename T>
void vertical_scan(const vector<vector<T>>& src, vector<T>& dst) {
size_t in_height = src.size();
size_t in_width = src[0].size();

dst.resize(in_height * in_width);

// check if each row size is equals to in_width
if (any_of(src.begin(), src.end(), [in_width](const vector<T>& i) {
return i.size() != in_width;
})) {
throw invalid_argument("rows sizes are different");
}

for (size_t i = 0; i < in_height; i++) {
for (size_t j = 0; j < in_width; j++) {
dst[i + j * in_height] = src[i][j];
}
}
}

/**
* Image Block to Columns implementation
**/
template <typename T>
size_t im2col(const vector<vector<T>>& src, vector<vector<T>>& dst,
const size_t window_height, const size_t window_width,
const size_t stride) {
// input shape
size_t in_height = src.size();
size_t in_width = src[0].size();

if (src.empty()) {
throw invalid_argument("empty matrix");
}

// check if each row size is equals to in_width
if (any_of(src.begin(), src.end(), [in_width](const vector<T>& i) {
return i.size() != in_width;
})) {
throw invalid_argument("rows sizes are different");
}

// output shape
size_t out_height = (in_height - window_height) / stride + 1;
size_t out_width = (in_width - window_width) / stride + 1;
dst.reserve(out_height);

// windows number
size_t windows_nb = out_height * out_width;

// kernel_size
size_t kernel_size = window_width * window_height;
// calculate the next power of 2
kernel_size = 1 << (static_cast<size_t>(ceil(log2(kernel_size))));

for (size_t j = 0; j < in_height - window_height + 1; j += stride) {
for (size_t i = 0; i < in_width - window_width + 1; i += stride) {
// pad the window vector to the next power of 2 of kernel size
vector<T> window_vec(kernel_size, 0);
auto iter = window_vec.begin();
for (size_t k = 0; k < window_height; k++) {
iter = copy(src[j + k].begin() + i,
src[j + k].begin() + i + window_width, iter);
}
dst.push_back(window_vec);
}
}

return windows_nb;
}

/*
Replicate the current vector as many times to fill `final_size` elements.
*/
Expand Down
25 changes: 13 additions & 12 deletions tests/python/tenseal/tensors/test_ckks_vector.py
Original file line number Diff line number Diff line change
Expand Up @@ -1125,8 +1125,9 @@ def test_polynomial_rescale_off(context, data, polynom):
@pytest.mark.parametrize(
"input_size, kernel_size", [(2, 2), (3, 2), (4, 2), (4, 3), (7, 3), (12, 5)]
)
def test_conv2d_im2col(context, input_size, kernel_size):
def generate_input(input_size, kernel_size, stride=1):
@pytest.mark.parametrize("stride", [1, 2, 3])
def test_conv2d_im2col(context, input_size, kernel_size, stride):
def generate_input(input_size, kernel_size, stride):
# generated random values and prepare the inputs
x = np.random.randn(input_size, input_size)
kernel = np.random.randn(kernel_size, kernel_size)
Expand All @@ -1136,27 +1137,27 @@ def generate_input(input_size, kernel_size, stride=1):
(x.shape[1] - kernel.shape[1]) // stride + 1,
)

new_x = view_as_windows(x, kernel.shape, step=stride)
new_x = new_x.reshape(out_h * out_w, kernel.shape[0] * kernel.shape[1])
padded_im2col_x = view_as_windows(x, kernel.shape, step=stride)
padded_im2col_x = padded_im2col_x.reshape(out_h * out_w, kernel.shape[0] * kernel.shape[1])

next_power2 = pow(2, math.ceil(math.log2(kernel.size)))
pad_width = next_power2 - kernel.size
new_x = np.pad(new_x, ((0, 0), (0, pad_width)))
padded_im2col_x = np.pad(padded_im2col_x, ((0, 0), (0, pad_width)))

kernel = np.pad(kernel.flatten(), (0, pad_width))
return new_x, kernel
padded_kernel = np.pad(kernel.flatten(), (0, pad_width))
return x, padded_im2col_x, kernel, padded_kernel

# generated galois keys in order to do rotation on ciphertext vectors
context.generate_galois_keys()

x, kernel = generate_input(input_size, kernel_size)
windows_nb = x.shape[0]
x, padded_im2col_x, kernel, padded_kernel = generate_input(input_size, kernel_size, stride)
# windows_nb = padded_im2col_x.shape[0]

x_enc, windows_nb = ts.im2col_encoding(context, x, kernel.shape[0], kernel.shape[1], stride)

x_enc = ts.ckks_vector(context, x.flatten(order="F").tolist())
y_enc = x_enc.conv2d_im2col(kernel.tolist(), windows_nb)
decrypted_result = y_enc.decrypt()

expected = (x @ kernel).tolist()
expected = (padded_im2col_x @ padded_kernel).tolist()
assert _almost_equal(decrypted_result, expected, 0)


Expand Down
Loading