Implement im2col and update conv2d (OpenMined#134)

* implement matrix vertical and horizontal scan

* Image Block to Columns implementation

* update conv2d_im2col

* im2col_encoding python binding

* update conv2d_im2col python test

* update Convolution notebook

* linting

* use const ref to prevent the copy

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__",
 ]

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<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);
+          });
+
     py::class_<CKKSVector>(m, "CKKSVector")
         // specifying scale
         .def(py::init<shared_ptr<TenSEALContext> &, vector<double>, double>())

tenseal/cpp/tensors/ckksvector.cpp

@@ -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");
     }
 
-    // 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());
     }
 

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<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.

tenseal/cpp/tensors/utils/utils.h

@@ -1,6 +1,8 @@
 #ifndef TENSEAL_UTILS_UTILS_H
 #define TENSEAL_UTILS_UTILS_H
 
+#include <algorithm>
+#include <iterator>
 #include <memory>
 
 #include "seal/seal.h"
@@ -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) {
+    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.
 */

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)