wehub-resource-sync
338 lines
11 KiB
Python
338 lines
11 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.
|
|
|
|
import logging
|
|
import re
|
|
from operator import attrgetter, itemgetter
|
|
|
|
import numpy as np
|
|
import torch.distributed as dist
|
|
import torch.nn as nn
|
|
|
|
from .modules import PQConv2d, PQEmbedding, PQLinear
|
|
from .pq import PQ
|
|
|
|
|
|
def quantize_model_(
|
|
model,
|
|
size_tracker,
|
|
layers_to_quantize,
|
|
block_sizes_config,
|
|
n_centroids_config,
|
|
step=0,
|
|
n_iter=15,
|
|
eps=1e-6,
|
|
max_tentatives=100,
|
|
verbose=True,
|
|
):
|
|
"""
|
|
Quantize a model in-place by stages. All the targeted
|
|
layers are replaced by their quantized counterpart,
|
|
and the model is ready for the finetuning of the
|
|
centroids in a standard training loop (no modifications
|
|
required). Note that we do not quantize biases.
|
|
|
|
Args:
|
|
- model: a nn.Module
|
|
- size_tracker: useful for tracking quatization statistics
|
|
- layers_to_quantize: a list containing regexps for
|
|
filtering the layers to quantize at each stage according
|
|
to their name (as in model.named_parameters())
|
|
- block_sizes_config: dict like
|
|
{
|
|
'Conv2d': ('kernel_size', {'(3, 3)': 9, '(1, 1)': 4}),
|
|
'Linear': ('in_features', {'*': 8})
|
|
}
|
|
For instance, all conv2d layers with kernel size 3x3 have
|
|
a block size of 9 and all Linear layers are quantized with
|
|
a block size of 8, irrespective of their size.
|
|
- n_centroids_config: dict like
|
|
{
|
|
'Conv2d': ('kernel_size', {'*': 256}),
|
|
'Linear': ('in_features', {'*': 256})
|
|
}
|
|
For instance, all conv2d layers are quantized with 256 centroids
|
|
- step: the layers to quantize inplace corresponding
|
|
to layers_to_quantize[step]
|
|
"""
|
|
|
|
quantized_layers = get_layers(model, layers_to_quantize[step])
|
|
|
|
for layer in quantized_layers:
|
|
|
|
# book-keeping
|
|
is_master_process = (not dist.is_initialized()) or (
|
|
dist.is_initialized() and dist.get_rank() == 0
|
|
)
|
|
verbose = verbose and is_master_process
|
|
|
|
# get block size and centroids
|
|
module = attrgetter(layer)(model)
|
|
block_size = get_param(module, layer, block_sizes_config)
|
|
n_centroids = get_param(module, layer, n_centroids_config)
|
|
if verbose:
|
|
logging.info(
|
|
f"Quantizing layer {layer} with block size {block_size} and {n_centroids} centroids"
|
|
)
|
|
|
|
# quantize layer
|
|
weight = module.weight.data.clone()
|
|
is_bias = "bias" in [x[0] for x in module.named_parameters()]
|
|
bias = module.bias.data.clone() if is_bias else None
|
|
quantizer = PQ(
|
|
weight,
|
|
block_size,
|
|
n_centroids=n_centroids,
|
|
n_iter=n_iter,
|
|
eps=eps,
|
|
max_tentatives=max_tentatives,
|
|
verbose=verbose,
|
|
)
|
|
|
|
# quantization performed on all GPUs with same seed
|
|
quantizer.encode()
|
|
centroids = quantizer.centroids.contiguous()
|
|
assignments = quantizer.assignments.contiguous()
|
|
|
|
# broadcast results to make sure weights are up-to-date
|
|
if dist.is_initialized():
|
|
dist.broadcast(centroids, 0)
|
|
dist.broadcast(assignments, 0)
|
|
|
|
# instantiate the quantized counterpart
|
|
if isinstance(module, nn.Linear):
|
|
out_features, in_features = map(
|
|
lambda k: module.__dict__[k], ["out_features", "in_features"]
|
|
)
|
|
quantized_module = PQLinear(
|
|
centroids, assignments, bias, in_features, out_features
|
|
)
|
|
elif isinstance(module, nn.Embedding):
|
|
num_embeddings, embedding_dim = map(
|
|
lambda k: module.__dict__[k], ["num_embeddings", "embedding_dim"]
|
|
)
|
|
quantized_module = PQEmbedding(
|
|
centroids, assignments, num_embeddings, embedding_dim
|
|
)
|
|
elif isinstance(module, nn.Conv2d):
|
|
out_channels, in_channels, kernel_size = map(
|
|
lambda k: module.__dict__[k],
|
|
["out_channels", "in_channels", "kernel_size"],
|
|
)
|
|
stride, padding, dilation, groups, padding_mode = map(
|
|
lambda k: module.__dict__[k],
|
|
["stride", "padding", "dilation", "groups", "padding_mode"],
|
|
)
|
|
|
|
quantized_module = PQConv2d(
|
|
centroids,
|
|
assignments,
|
|
bias,
|
|
in_channels,
|
|
out_channels,
|
|
kernel_size,
|
|
stride=stride,
|
|
padding=padding,
|
|
dilation=dilation,
|
|
groups=groups,
|
|
padding_mode=padding_mode,
|
|
)
|
|
else:
|
|
raise ValueError(f"Module {module} not yet supported for quantization")
|
|
|
|
# replace layer by its quantized counterpart
|
|
attrsetter(layer)(model, quantized_module)
|
|
|
|
# update statistics
|
|
size_tracker.update(weight, block_size, n_centroids)
|
|
|
|
# return name of quantized layers
|
|
return quantized_layers
|
|
|
|
|
|
def get_layers(model, filter_regexp):
|
|
"""
|
|
Filters out the layers according to a regexp. Note that
|
|
we omit biases.
|
|
|
|
Args:
|
|
- model: a nn.Module
|
|
- filter_regexp: a regexp to filter the layers to keep
|
|
according to their name in model.named_parameters().
|
|
For instance, the regexp:
|
|
|
|
down_layers\\.[123456]\\.(conv[12]|identity\\.conv))
|
|
|
|
is keeping blocks down_layers from 1 to 6, and inside
|
|
each block is keeping conv1, conv2 and identity.conv.
|
|
|
|
Remarks:
|
|
- We add (module\\.)? at the beginning of the regexp to
|
|
account for the possible use of nn.parallel.DataParallel
|
|
"""
|
|
|
|
# get all parameter names
|
|
all_layers = map(itemgetter(0), model.named_parameters())
|
|
|
|
# remove biases
|
|
all_layers = filter(lambda x: "bias" not in x, all_layers)
|
|
|
|
# remove .weight in all other names (or .weight_orig is spectral norm)
|
|
all_layers = map(lambda x: x.replace(".weight_orig", ""), all_layers)
|
|
all_layers = map(lambda x: x.replace(".weight", ""), all_layers)
|
|
|
|
# return filtered layers
|
|
filter_regexp = "(module\\.)?" + "(" + filter_regexp + ")"
|
|
r = re.compile(filter_regexp)
|
|
|
|
return list(filter(r.match, all_layers))
|
|
|
|
|
|
def get_param(module, layer_name, param_config):
|
|
"""
|
|
Given a quantization configuration, get the right parameter
|
|
for the module to be quantized.
|
|
|
|
Args:
|
|
- module: a nn.Module
|
|
- layer_name: the name of the layer
|
|
- param_config: a dict like
|
|
{
|
|
'Conv2d': ('kernel_size', {'(3, 3)': 9, '(1, 1)': 4}),
|
|
'Linear': ('in_features', {'*': 8})
|
|
}
|
|
For instance, all conv2d layers with kernel size 3x3 have
|
|
a block size of 9 and all Linear layers are quantized with
|
|
a block size of 8, irrespective of their size.
|
|
|
|
Remarks:
|
|
- if 'fuzzy_name' is passed as a parameter, layers whose layer_name
|
|
include 'fuzzy_name' will be assigned the given parameter.
|
|
In the following example, conv.expand layers will have a block
|
|
size of 9 while conv.reduce will have a block size of 4 and all
|
|
other layers will have a block size of 2.
|
|
{
|
|
'Conv2d': ('fuzzy_name', {'expand': 9, 'reduce': 4, '*': 2}),
|
|
'Linear': ('fuzzy_name', {'classifier': 8, 'projection': 4})
|
|
}
|
|
|
|
"""
|
|
|
|
layer_type = module.__class__.__name__
|
|
|
|
if layer_type not in param_config:
|
|
raise KeyError(f"Layer type {layer_type} not in config for layer {module}")
|
|
|
|
feature, params = param_config[module.__class__.__name__]
|
|
|
|
if feature != "fuzzy_name":
|
|
feature_value = str(getattr(module, feature))
|
|
if feature_value not in params:
|
|
if "*" in params:
|
|
feature_value = "*"
|
|
else:
|
|
raise KeyError(
|
|
f"{feature}={feature_value} not in config for layer {module}"
|
|
)
|
|
else:
|
|
feature_values = [name for name in params if name in layer_name]
|
|
if len(feature_values) == 0:
|
|
if "*" in params:
|
|
feature_value = "*"
|
|
else:
|
|
raise KeyError(f"name={layer_name} not in config for {module}")
|
|
else:
|
|
feature_value = feature_values[0]
|
|
|
|
return params[feature_value]
|
|
|
|
|
|
class SizeTracker(object):
|
|
"""
|
|
Class to keep track of the compressed network size with iPQ.
|
|
|
|
Args:
|
|
- model: a nn.Module
|
|
|
|
Remarks:
|
|
- The compressed size is the sum of three components
|
|
for each layer in the network:
|
|
(1) Storing the centroids given by iPQ in fp16
|
|
(2) Storing the assignments of the blocks in int8
|
|
(3) Storing all non-compressed elements such as biases
|
|
- This cost in only valid if we use 256 centroids (then
|
|
indexing can indeed by done with int8).
|
|
"""
|
|
|
|
def __init__(self, model):
|
|
self.model = model
|
|
self.size_non_compressed_model = self.compute_size()
|
|
self.size_non_quantized = self.size_non_compressed_model
|
|
self.size_index = 0
|
|
self.size_centroids = 0
|
|
self.n_quantized_layers = 0
|
|
|
|
def compute_size(self):
|
|
"""
|
|
Computes the size of the model (in MB).
|
|
"""
|
|
|
|
res = 0
|
|
for _, p in self.model.named_parameters():
|
|
res += p.numel()
|
|
return res * 4 / 1024 / 1024
|
|
|
|
def update(self, W, block_size, n_centroids):
|
|
"""
|
|
Updates the running statistics when quantizing a new layer.
|
|
"""
|
|
|
|
# bits per weights
|
|
bits_per_weight = np.log2(n_centroids) / block_size
|
|
self.n_quantized_layers += 1
|
|
|
|
# size of indexing the subvectors of size block_size (in MB)
|
|
size_index_layer = bits_per_weight * W.numel() / 8 / 1024 / 1024
|
|
self.size_index += size_index_layer
|
|
|
|
# size of the centroids stored in float16 (in MB)
|
|
size_centroids_layer = n_centroids * block_size * 2 / 1024 / 1024
|
|
self.size_centroids += size_centroids_layer
|
|
|
|
# size of non-compressed layers, e.g. LayerNorms or biases (in MB)
|
|
size_uncompressed_layer = W.numel() * 4 / 1024 / 1024
|
|
self.size_non_quantized -= size_uncompressed_layer
|
|
|
|
def __repr__(self):
|
|
size_compressed = (
|
|
self.size_index + self.size_centroids + self.size_non_quantized
|
|
)
|
|
compression_ratio = self.size_non_compressed_model / size_compressed # NOQA
|
|
return (
|
|
f"Non-compressed model size: {self.size_non_compressed_model:.2f} MB. "
|
|
f"After quantizing {self.n_quantized_layers} layers, size "
|
|
f"(indexing + centroids + other): {self.size_index:.2f} MB + "
|
|
f"{self.size_centroids:.2f} MB + {self.size_non_quantized:.2f} MB = "
|
|
f"{size_compressed:.2f} MB, compression ratio: {compression_ratio:.2f}x"
|
|
)
|
|
|
|
|
|
def attrsetter(*items):
|
|
def resolve_attr(obj, attr):
|
|
attrs = attr.split(".")
|
|
head = attrs[:-1]
|
|
tail = attrs[-1]
|
|
|
|
for name in head:
|
|
obj = getattr(obj, name)
|
|
return obj, tail
|
|
|
|
def g(obj, val):
|
|
for attr in items:
|
|
resolved_obj, resolved_attr = resolve_attr(obj, attr)
|
|
setattr(resolved_obj, resolved_attr, val)
|
|
|
|
return g
|