diff --git a/distiller/__init__.py b/distiller/__init__.py
index 13ef908e3cc8cfa7086004c3d97d47be4f4c5c6e..b287be91fc1016c504e5a659870c11811ffe5ce7 100755
--- a/distiller/__init__.py
+++ b/distiller/__init__.py
@@ -36,10 +36,10 @@ del dict_config
del thinning
# Distiller version
-try:
- __version__ = pkg_resources.require("distiller")[0].version
-except pkg_resources.DistributionNotFound:
- __version__ = "Unknown"
+# try:
+# __version__ = pkg_resources.require("distiller")[0].version
+# except pkg_resources.DistributionNotFound:
+# __version__ = "Unknown"
def __check_pytorch_version():
diff --git a/distiller/data_loggers/collector.py b/distiller/data_loggers/collector.py
index 8cc41aab6921d8026720dd509759a96d12aff832..e13346672cb1b9b334a307e2405df30ac6e805df 100755
--- a/distiller/data_loggers/collector.py
+++ b/distiller/data_loggers/collector.py
@@ -556,6 +556,10 @@ class QuantCalibrationStatsCollector(ActivationStatsCollector):
tensor = tensor.to(torch.float32)
if not tensor.is_contiguous():
tensor = tensor.contiguous()
+
+ if (len(tensor.size()) == 0):
+ return
+
act = tensor.view(tensor.size(0), -1)
numel = act.numel()
if self.collecting_second_pass:
@@ -842,7 +846,7 @@ class RawActivationsCollector(ActivationStatsCollector):
return dir_name
-def collect_quant_stats(model, test_fn, save_dir=None, classes=None, inplace_runtime_check=False,
+def collect_quant_stats(model, trainer, dataloader, save_dir=None, classes=None, inplace_runtime_check=False,
disable_inplace_attrs=False, inplace_attr_names=('inplace',),
modules_to_collect=None):
"""
@@ -873,11 +877,13 @@ def collect_quant_stats(model, test_fn, save_dir=None, classes=None, inplace_run
inplace_attr_names=inplace_attr_names)
with collector_context(quant_stats_collector, modules_to_collect):
msglogger.info('Pass 1: Collecting min, max, avg_min, avg_max, mean')
- test_fn(model=model)
+ trainer(model, dataloader)
+ # trainer.test(model, test_dataloaders=dataloader)
# Collect Laplace distribution stats:
msglogger.info('Pass 2: Collecting b, std parameters')
quant_stats_collector.start_second_pass()
- test_fn(model=model)
+ trainer(model, dataloader)
+ # trainer.test(model, test_dataloaders=dataloader)
quant_stats_collector.stop_second_pass()
msglogger.info('Stats collection complete')
diff --git a/distiller/quantization/quantizer.py b/distiller/quantization/quantizer.py
index 19a8e83e36fd4e9eaa47544d1578b8073475f863..1972c6b2fcdc7b778dafdb61b36fdb5ba9bc9d11 100644
--- a/distiller/quantization/quantizer.py
+++ b/distiller/quantization/quantizer.py
@@ -244,6 +244,7 @@ class Quantizer(object):
curr_parameters = dict(module.named_parameters())
for param_name, param in curr_parameters.items():
n_bits = qbits.bias if param_name.endswith('bias') else qbits.wts
+
if n_bits is None:
continue
fp_attr_name = param_name
@@ -270,7 +271,7 @@ class Quantizer(object):
self.prepared = True
- msglogger.debug('Quantized model:\n\n{0}\n'.format(self.model))
+ # msglogger.debug('Quantized model:\n\n{0}\n'.format(self.model))
def _pre_prepare_model(self, dummy_input):
pass
@@ -301,6 +302,7 @@ class Quantizer(object):
else:
replace_msg(full_name)
continue
+
current_qbits = self.module_qbits_map[full_name]
# TODO - Review necessity of the block below
if current_qbits.acts is None and current_qbits.wts is None and not self.module_overrides_map[full_name]:
diff --git a/distiller/quantization/range_linear.py b/distiller/quantization/range_linear.py
index 092d4f3a3ef724c3ca8bfe7e913f13ff56bf7e62..f602fb607ed91d8c649f346d9d7a48eb09091877 100644
--- a/distiller/quantization/range_linear.py
+++ b/distiller/quantization/range_linear.py
@@ -793,6 +793,7 @@ class RangeLinearQuantParamLayerWrapper(RangeLinearQuantWrapper):
# Quantize bias
self.has_bias = hasattr(wrapped_module, 'bias') and wrapped_module.bias is not None
+
if self.has_bias and (self.fake_quant_params or not self.preset_act_stats):
b_scale, b_zero_point = _get_quant_params_from_tensor(wrapped_module.bias,
self.accum_quant_settings.num_bits,
@@ -1609,6 +1610,7 @@ class PostTrainLinearQuantizer(Quantizer):
scale_approx_mult_bits=None, inputs_quant_auto_fallback=True,
fpq_module=None, save_fp_weights=False, also_clip_weights=False):
overrides_bkp = deepcopy(overrides)
+
super(PostTrainLinearQuantizer, self).__init__(model, bits_activations=bits_activations,
bits_weights=bits_parameters, bits_bias=bits_accum,
overrides=overrides, train_with_fp_copy=False)
@@ -1633,7 +1635,7 @@ class PostTrainLinearQuantizer(Quantizer):
self.model.quantizer_metadata = {'type': type(self),
'params': {'bits_activations': bits_activations,
'bits_parameters': bits_parameters,
- 'bits_accum': bits_accum,
+ 'bits_accum': bits_parameters,
'mode': mode_dict,
'clip_acts': _enum_to_str(clip_acts),
'clip_n_stds': clip_n_stds,
diff --git a/distiller/summary_graph.py b/distiller/summary_graph.py
index bfd1ed0ed69fa3e13c05b2718627c8b4fb1100d2..5f2c819f92f324654d0d27e8cdc65a0bb590b653 100755
--- a/distiller/summary_graph.py
+++ b/distiller/summary_graph.py
@@ -87,7 +87,7 @@ class SummaryGraph(object):
device = distiller.model_device(model_clone)
dummy_input = distiller.convert_tensors_recursively_to(dummy_input, device=device)
self.dummy_input = dummy_input
- trace, _ = jit.get_trace_graph(model_clone, dummy_input, _force_outplace=True)
+ trace, _ = jit._get_trace_graph(model_clone, dummy_input, _force_outplace=True)
# As of PyTorch 1.3.0, ONNX trace optimization has an issue that results in incorrect scope names
# of nodes in the trace graph.
@@ -100,7 +100,8 @@ class SummaryGraph(object):
pre_dropout_nodes_scope_names = OrderedDict()
prev_non_dropout_op = None
- for node in trace.graph().nodes():
+
+ for node in trace.nodes():
kind = node.kind()
if 'aten' not in kind:
continue
@@ -114,18 +115,17 @@ class SummaryGraph(object):
# composing a GEMM operation; etc.
torch.onnx._optimize_trace(trace, torch.onnx.OperatorExportTypes.ONNX)
- graph = trace.graph()
self.ops = OrderedDict()
self.module_ops_map = defaultdict(list)
self.params = OrderedDict()
self.edges = []
self.temp = OrderedDict()
- in_out = list(graph.inputs()) + list(graph.outputs())
+ in_out = list(trace.inputs()) + list(trace.outputs())
for param in in_out:
self.__add_param(param)
- for node in graph.nodes():
+ for node in trace.nodes():
new_op = self.__create_op(node)
if apply_scope_name_workarounds:
@@ -196,7 +196,6 @@ class SummaryGraph(object):
self.add_footprint_attr()
self.add_arithmetic_intensity_attr()
del trace
- del graph
del model_clone
def __merge_pad_avgpool(self):
diff --git a/requirements.txt b/requirements.txt
index 197df57a56667defac3be021e2237e5281d867b9..3e4d7d6321e658304b1e70faf9024cdbf9ce69d0 100755
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,16 +1,14 @@
-pillow<7 # remove this when moving to torchvision>0.5 (see https://github.com/pytorch/vision/issues/1712#issuecomment-575036523)
-torch==1.3.1
-numpy>=1.16
-torchvision==0.4.2
-scipy>=1.3.0
+Cython
+torch==1.5.0
+torchvision==0.6
+scipy
gitpython==3.1.0
torchnet==0.0.4
-tensorflow~=1.14
pydot==1.4.1
tabulate==0.8.3
-pandas>=0.22.0
+pandas==1.1
jupyter>=1.0.0
-matplotlib~=3.0 # 3.0 is the last release to support Py3.5
+matplotlib==3.3 # 3.0 is the last release to support Py3.5
qgrid==1.1.1
graphviz==0.10.1
ipywidgets==7.4.2
@@ -19,6 +17,6 @@ pyyaml
pytest~=4.6.1
xlsxwriter>=1.2.2
pretrainedmodels==0.7.4
-scikit-learn==0.21.2
+scikit-learn
gym==0.12.5
-tqdm==4.33.0
+tqdm==4.59