pytorch_quantization.optim.helper

Helper functions for quant optimizer/trainer

pytorch_quantization.optim.helper.freeze_parameters(model, patterns)[source]

Set requires_grad to False if patterns match name

Parameters
  • model – A Module

  • patterns – A list of strings that will be used to match parameter names. If parameter name contains any pattern, it will be frozen.

pytorch_quantization.optim.helper.group_parameters(model, patterns_list, lrs=None, momentums=None, weight_decays=None)[source]

Group parameters for using per-parameters option in optimizer

Returns a list of dict that matches Pytorch optimizer fashion, see https://pytorch.org/docs/stable/optim.html#per-parameter-options for more details.

Example

>>> [
>>>    {'params': model.base.parameters()},
>>>    {'params': model.classifier.parameters(), 'lr': 1e-3}
>>> ]

Parameters will be grouped w.r.t first level of the keys_list. e.g. keys_list=[[‘conv1’, ‘conv2’], [‘conv3’]] will return 2 groups, one with conv1 and conv2 in name, and the other with conv3 in name.

If lr, momentum or weight_decay are supplied, they will be added to the group as well.

Parameters
  • model – A module

  • patterns_list – A list of list of strings. WARNING: patters must be EXCLUSIVE, the function doesn’t perform exclusive check.

  • lrs – A list of float with same length as keys_list or None.

  • momentums – A list of float with same length as keys_list or None.

  • weight_decays – A list of float with same length as keys_list or None.

Returns

param_group – A list of dict

pytorch_quantization.optim.helper.match_parameters(model, patterns)[source]

Returns an generator over module parameters if name matches key

It is useful to group parameters, and apply different functions to different group. This function provides an easy way to group them.

Parameters
  • model – A Module

  • patterns – A list of strings that will be used to match parameter names. If parameter name contains any pattern, it will be yield

Yields

param – Module parameters

pytorch_quantization.optim.helper.quant_weight_inplace(model)[source]

Make quantization inplace

Search for quantized modules including QuantConvNd and QuantLinear, make weight quantization in place using weight_quantizer.

Most publications of quantization aware training uses STE by default, which is really an approximation of derivative of the nondifferentiable quantization function, which works to some extended but by no means the F=ma of the problem. Inplace quantization can be used to implement relax-and-round, which is a common method in Discrete Optimization’s or Integer Programming.