diff --git a/tenseal/__init__.py b/tenseal/__init__.py index 59ce42e3..ca553898 100644 --- a/tenseal/__init__.py +++ b/tenseal/__init__.py @@ -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 @@ -80,5 +83,6 @@ def context_from(buff, n_threads=None): "ckks_vector_from", "context", "context_from", + "im2col_encoding", "__version__", ] diff --git a/tenseal/binding.cpp b/tenseal/binding.cpp index 10d9d03a..218f2757 100644 --- a/tenseal/binding.cpp +++ b/tenseal/binding.cpp @@ -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; @@ -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 ctx, vector> &input, + const size_t kernel_n_rows, const size_t kernel_n_cols, + const size_t stride) { + vector> view_as_window; + vector 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); + }); + py::class_(m, "CKKSVector") // specifying scale .def(py::init &, vector, double>()) diff --git a/tenseal/cpp/tensors/ckksvector.cpp b/tenseal/cpp/tensors/ckksvector.cpp index e93eb580..4830f242 100644 --- a/tenseal/cpp/tensors/ckksvector.cpp +++ b/tenseal/cpp/tensors/ckksvector.cpp @@ -564,39 +564,47 @@ CKKSVector& CKKSVector::polyval_inplace(const vector& coefficients) { return *this; } -CKKSVector CKKSVector::conv2d_im2col(const vector& kernel, - size_t windows_nb) { +CKKSVector CKKSVector::conv2d_im2col(const vector>& 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& kernel, - const size_t windows_nb) { - vector plain_vec; - size_t chunks_nb = kernel.size(); - +CKKSVector& CKKSVector::conv2d_im2col_inplace( + const vector>& 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& i) { return i.empty(); }))) { + throw invalid_argument("Kernel matrix can't be empty"); } - // 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 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(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 plain_vec; plain_vec.reserve(this->_size); for (size_t i = 0; i < chunks_nb; i++) { - vector tmp(windows_nb, kernel[i]); + vector tmp(windows_nb, flatten_kernel[i]); plain_vec.insert(plain_vec.end(), tmp.begin(), tmp.end()); } diff --git a/tenseal/cpp/tensors/ckksvector.h b/tenseal/cpp/tensors/ckksvector.h index c13c0d6a..2e635108 100644 --- a/tenseal/cpp/tensors/ckksvector.h +++ b/tenseal/cpp/tensors/ckksvector.h @@ -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& kernel, size_t windows_nb); - CKKSVector& conv2d_im2col_inplace(const vector& kernel, - size_t windows_nb); + CKKSVector conv2d_im2col(const vector>& kernel, + const size_t windows_nb); + CKKSVector& conv2d_im2col_inplace(const vector>& kernel, + const size_t windows_nb); /** * Load/Save the vector from/to a serialized protobuffer. diff --git a/tenseal/cpp/tensors/utils/utils.h b/tenseal/cpp/tensors/utils/utils.h index 89724e28..b94a0da1 100644 --- a/tenseal/cpp/tensors/utils/utils.h +++ b/tenseal/cpp/tensors/utils/utils.h @@ -1,6 +1,8 @@ #ifndef TENSEAL_UTILS_UTILS_H #define TENSEAL_UTILS_UTILS_H +#include +#include #include #include "seal/seal.h" @@ -10,6 +12,107 @@ namespace tenseal { using namespace seal; using namespace std; +/** + * horizontally scan matrix (vector of vectors) + **/ +template +void horizontal_scan(const vector>& src, vector& 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& 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 +void vertical_scan(const vector>& src, vector& 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& 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 +size_t im2col(const vector>& src, vector>& 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& 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(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 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. */ diff --git a/tests/python/tenseal/tensors/test_ckks_vector.py b/tests/python/tenseal/tensors/test_ckks_vector.py index 9146c623..d40160f7 100644 --- a/tests/python/tenseal/tensors/test_ckks_vector.py +++ b/tests/python/tenseal/tensors/test_ckks_vector.py @@ -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) @@ -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) diff --git a/tutorials/experimental/Convolution.ipynb b/tutorials/experimental/Convolution.ipynb index 41e7b9cd..40144400 100644 --- a/tutorials/experimental/Convolution.ipynb +++ b/tutorials/experimental/Convolution.ipynb @@ -49,31 +49,6 @@ " )" ] }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Memory strided Image Block to Columns implementation" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "metadata": {}, - "outputs": [], - "source": [ - "def memory_strided_im2col(x, kernel_shape, stride=1):\n", - " # Infer shapes\n", - " x_h, x_w = x.shape\n", - " k_h, k_w = kernel_shape\n", - " # Assuming Padding=0, Stride=1\n", - " out_h, out_w = ((x.shape[0] - kernel_shape[0])//stride + 1,\n", - " (x.shape[1] - kernel_shape[1])//stride + 1)\n", - "\n", - " windows = view_as_windows(x, kernel_shape, step=stride)\n", - " return windows.reshape(out_h * out_w, k_h * k_w)" - ] - }, { "cell_type": "markdown", "metadata": {}, @@ -83,7 +58,7 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": 4, "metadata": {}, "outputs": [ { @@ -95,10 +70,9 @@ " [ 5 6 7 8]\n", " [ 9 10 11 12]\n", " [13 14 15 16]]\n", - "kernel (3, 3)\n", - "[[1 2 3]\n", - " [4 5 6]\n", - " [7 8 9]]\n" + "kernel (2, 2)\n", + "[[1 2]\n", + " [3 4]]\n" ] } ], @@ -106,7 +80,7 @@ "# input image dimension n * n\n", "x_size = 4\n", "# kernel dimension n * n\n", - "k_size = 3\n", + "k_size = 2\n", "# stride\n", "stride = 1\n", "\n", @@ -133,7 +107,7 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": 5, "metadata": {}, "outputs": [], "source": [ @@ -151,44 +125,14 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "For each convolution layer, a communication between the client and server is required. The server send the ciphertext (encrypted vector) to the client which is the input of the next convolution layer, in order to decrypt it and apply im2col (Image Block to Column) on the that input." - ] - }, - { - "cell_type": "code", - "execution_count": 7, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "new_x shape: (4, 9)\n", - "[[ 1 2 3 5 6 7 9 10 11]\n", - " [ 2 3 4 6 7 8 10 11 12]\n", - " [ 5 6 7 9 10 11 13 14 15]\n", - " [ 6 7 8 10 11 12 14 15 16]]\n" - ] - } - ], - "source": [ - "new_x = memory_strided_im2col(x, kernel.shape, stride)\n", - "print(\"new_x shape: \", new_x.shape)\n", - "print(new_x)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "After that the client encode and encrypt the input matrix in a vertical scan (column-major) and send it back to the server.\n", + "For each convolution layer, a communication between the client and server is required. The server send the ciphertext (encrypted vector) to the client which is the input of the next convolution layer, in order to decrypt it and apply im2col (Image Block to Column) on the that input.\n", "\n", - "new_x.flatten(order='F') is equivalent to new_x.T.flatten()\n" + "After that the client encode and encrypt the input matrix in a vertical scan (column-major) and send it back to the server." ] }, { "cell_type": "code", - "execution_count": 8, + "execution_count": 6, "metadata": { "scrolled": true }, @@ -197,43 +141,25 @@ "name": "stdout", "output_type": "stream", "text": [ + "windows number: 9\n", + "ckksvector size: 36\n", "9\n", - "(4, 16)\n", - "[[ 1 2 3 5 6 7 9 10 11 0 0 0 0 0 0 0]\n", - " [ 2 3 4 6 7 8 10 11 12 0 0 0 0 0 0 0]\n", - " [ 5 6 7 9 10 11 13 14 15 0 0 0 0 0 0 0]\n", - " [ 6 7 8 10 11 12 14 15 16 0 0 0 0 0 0 0]]\n", - "flatten_kernel_size: 9\n", - "windows number: 4\n", - "ckksvector size: 64\n", - "4\n", "y_enc\n", - "[348.000048024785, 393.00005275438355, 528.0000706569963, 573.0000769331538]\n", - "CPU times: user 41.8 ms, sys: 0 ns, total: 41.8 ms\n", - "Wall time: 40.7 ms\n" + "[44.00000566594027, 54.00000698409561, 64.00000865907336, 84.0000112607995, 94.0000125997811, 104.00001396028475, 124.00001662478137, 134.00001797146558, 144.00001931406726]\n", + "CPU times: user 25.7 ms, sys: 0 ns, total: 25.7 ms\n", + "Wall time: 24.6 ms\n" ] } ], "source": [ "%%time \n", - "# pad the input\n", - "print(kernel.size)\n", - "next_power2 = pow(2, math.ceil(math.log2(kernel.size)))\n", - "pad_width = next_power2 - kernel.size\n", - "padded_x = np.pad(new_x, ((0, 0), (0, pad_width)))\n", - "print(padded_x.shape)\n", - "print(padded_x)\n", "\n", - "x_enc = ts.ckks_vector(context, padded_x.flatten(order='F').tolist())\n", - "windows_nb = padded_x.shape[0]\n", - "print(\"flatten_kernel_size: \", kernel.size)\n", + "x_enc, windows_nb = ts.im2col_encoding(context, x, kernel.shape[0], kernel.shape[1], stride)\n", + "\n", "print(\"windows number: \", windows_nb)\n", "print(\"ckksvector size: \", x_enc.size())\n", "\n", - "# pad the kernel\n", - "padded_kernel = np.pad(kernel.flatten(), (0, pad_width))\n", - "\n", - "y_enc = x_enc.conv2d_im2col(padded_kernel.tolist(), windows_nb)\n", + "y_enc = x_enc.conv2d_im2col(kernel.tolist(), windows_nb)\n", "\n", "print(y_enc.size())\n", "y_plain = y_enc.decrypt()\n", @@ -251,7 +177,7 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": 7, "metadata": {}, "outputs": [ { @@ -259,7 +185,7 @@ "output_type": "stream", "text": [ "y_toch\n", - "[348. 393. 528. 573.]\n" + "[ 44. 54. 64. 84. 94. 104. 124. 134. 144.]\n" ] } ],