wehub-resource-sync
129 lines
4.2 KiB
Python
129 lines
4.2 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
from .em import EM, EmptyClusterResolveError
|
|
|
|
|
|
class PQ(EM):
|
|
"""
|
|
Quantizes the layer weights W with the standard Product Quantization
|
|
technique. This learns a codebook of codewords or centroids of size
|
|
block_size from W. For further reference on using PQ to quantize
|
|
neural networks, see "And the Bit Goes Down: Revisiting the Quantization
|
|
of Neural Networks", Stock et al., ICLR 2020.
|
|
|
|
PQ is performed in two steps:
|
|
(1) The matrix W (weights or fully-connected or convolutional layer)
|
|
is reshaped to (block_size, -1).
|
|
- If W is fully-connected (2D), its columns are split into
|
|
blocks of size block_size.
|
|
- If W is convolutional (4D), its filters are split along the
|
|
spatial dimension.
|
|
(2) We apply the standard EM/k-means algorithm to the resulting reshaped matrix.
|
|
|
|
Args:
|
|
- W: weight matrix to quantize of size (in_features x out_features)
|
|
- block_size: size of the blocks (subvectors)
|
|
- n_centroids: number of centroids
|
|
- n_iter: number of k-means iterations
|
|
- eps: for cluster reassignment when an empty cluster is found
|
|
- max_tentatives for cluster reassignment when an empty cluster is found
|
|
- verbose: print information after each iteration
|
|
|
|
Remarks:
|
|
- block_size be compatible with the shape of W
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
W,
|
|
block_size,
|
|
n_centroids=256,
|
|
n_iter=20,
|
|
eps=1e-6,
|
|
max_tentatives=30,
|
|
verbose=True,
|
|
):
|
|
self.block_size = block_size
|
|
W_reshaped = self._reshape(W)
|
|
super(PQ, self).__init__(
|
|
W_reshaped,
|
|
n_centroids=n_centroids,
|
|
n_iter=n_iter,
|
|
eps=eps,
|
|
max_tentatives=max_tentatives,
|
|
verbose=verbose,
|
|
)
|
|
|
|
def _reshape(self, W):
|
|
"""
|
|
Reshapes the matrix W as expained in step (1).
|
|
"""
|
|
|
|
# fully connected: by convention the weight has size out_features x in_features
|
|
if len(W.size()) == 2:
|
|
self.out_features, self.in_features = W.size()
|
|
assert (
|
|
self.in_features % self.block_size == 0
|
|
), "Linear: n_blocks must be a multiple of in_features"
|
|
return (
|
|
W.reshape(self.out_features, -1, self.block_size)
|
|
.permute(2, 1, 0)
|
|
.flatten(1, 2)
|
|
)
|
|
|
|
# convolutional: we reshape along the spatial dimension
|
|
elif len(W.size()) == 4:
|
|
self.out_channels, self.in_channels, self.k_h, self.k_w = W.size()
|
|
assert (
|
|
self.in_channels * self.k_h * self.k_w
|
|
) % self.block_size == 0, (
|
|
"Conv2d: n_blocks must be a multiple of in_channels * k_h * k_w"
|
|
)
|
|
return (
|
|
W.reshape(self.out_channels, -1, self.block_size)
|
|
.permute(2, 1, 0)
|
|
.flatten(1, 2)
|
|
)
|
|
# not implemented
|
|
else:
|
|
raise NotImplementedError(W.size())
|
|
|
|
def encode(self):
|
|
"""
|
|
Performs self.n_iter EM steps.
|
|
"""
|
|
|
|
self.initialize_centroids()
|
|
for i in range(self.n_iter):
|
|
try:
|
|
self.step(i)
|
|
except EmptyClusterResolveError:
|
|
break
|
|
|
|
def decode(self):
|
|
"""
|
|
Returns the encoded full weight matrix. Must be called after
|
|
the encode function.
|
|
"""
|
|
|
|
# fully connected case
|
|
if "k_h" not in self.__dict__:
|
|
return (
|
|
self.centroids[self.assignments]
|
|
.reshape(-1, self.out_features, self.block_size)
|
|
.permute(1, 0, 2)
|
|
.flatten(1, 2)
|
|
)
|
|
|
|
# convolutional case
|
|
else:
|
|
return (
|
|
self.centroids[self.assignments]
|
|
.reshape(-1, self.out_channels, self.block_size)
|
|
.permute(1, 0, 2)
|
|
.reshape(self.out_channels, self.in_channels, self.k_h, self.k_w)
|
|
)
|