Initializes locality-sensitive hashing with FALCONN to find nearest neighbors in training data.

def init_lsh(self):
    """
    Initializes locality-sensitive hashing with FALCONN to find nearest neighbors in training data.
    """
    self.query_objects = {
    }  # contains the object that can be queried to find nearest neighbors at each layer.
    # mean of training data representation per layer (that needs to be substracted before LSH).
    self.centers = {}
    for layer in self.layers:
      assert self.nb_tables >= self.neighbors

      # Normalize all the lenghts, since we care about the cosine similarity.
      self.train_activations_lsh[layer] /= np.linalg.norm(
          self.train_activations_lsh[layer], axis=1).reshape(-1, 1)

      # Center the dataset and the queries: this improves the performance of LSH quite a bit.
      center = np.mean(self.train_activations_lsh[layer], axis=0)
      self.train_activations_lsh[layer] -= center
      self.centers[layer] = center

      # LSH parameters
      params_cp = falconn.LSHConstructionParameters()
      params_cp.dimension = len(self.train_activations_lsh[layer][1])
      params_cp.lsh_family = falconn.LSHFamily.CrossPolytope
      params_cp.distance_function = falconn.DistanceFunction.EuclideanSquared
      params_cp.l = self.nb_tables
      params_cp.num_rotations = 2  # for dense set it to 1; for sparse data set it to 2
      params_cp.seed = 5721840
      # we want to use all the available threads to set up
      params_cp.num_setup_threads = 0
      params_cp.storage_hash_table = falconn.StorageHashTable.BitPackedFlatHashTable

      # we build 18-bit hashes so that each table has
      # 2^18 bins; this is a good choice since 2^18 is of the same
      # order of magnitude as the number of data points
      falconn.compute_number_of_hash_functions(self.number_bits, params_cp)

      print('Constructing the LSH table')
      table = falconn.LSHIndex(params_cp)
      table.setup(self.train_activations_lsh[layer])

      # Parse test feature vectors and find k nearest neighbors
      query_object = table.construct_query_object()
      query_object.set_num_probes(self.nb_tables)
      self.query_objects[layer] = query_object