def quantize_model(
model: Model,
qconfig_spec: Dict = None,
dtype: Union[str, Optional[torch.dtype]] = "qint8",
) -> Model:
"""Function to quantize model weights.
Args:
model (Model): model to quantize
qconfig_spec (Dict, optional): quantization config in PyTorch format. Defaults to None.
dtype (Union[str, Optional[torch.dtype]], optional): Type of weights after quantization.
Defaults to "qint8".
Returns:
Model: quantized model
"""
if isinstance(dtype, str):
type_mapping = {"qint8": torch.qint8, "quint8": torch.quint8}
try:
quantized_model = quantization.quantize_dynamic(
model.cpu(),
qconfig_spec=qconfig_spec,
dtype=type_mapping[dtype],
)
except RuntimeError:
torch.backends.quantized.engine = "qnnpack"
quantized_model = quantization.quantize_dynamic(
model.cpu(),
qconfig_spec=qconfig_spec,
dtype=type_mapping[dtype],
)
return quantized_model
def set_requires_grad(model: Model, requires_grad: Union[bool, Dict[str,
bool]]):
"""Sets the ``requires_grad`` value for all model parameters.
Example::
>>> model = SimpleModel()
>>> set_requires_grad(model, requires_grad=True)
>>> # or
>>> model = SimpleModel()
>>> set_requires_grad(model, requires_grad={""})
Args:
model: model
requires_grad (Union[bool, Dict[str, bool]]): value
"""
if isinstance(requires_grad, dict):
for name, param in model.named_parameters():
assert (name in requires_grad
), f"Parameter `{name}` does not exist in requires_grad"
param.requires_grad = requires_grad[name]
else:
requires_grad = bool(requires_grad)
for param in model.parameters():
param.requires_grad = requires_grad
def grad_norm(*, model: Model, prefix: str, norm_type: int) -> Dict:
"""Computes gradient norms for a given model.
Args:
model: model which gradients to be saved.
prefix: prefix for keys in resulting dictionary.
norm_type: norm type of gradient norm.
Returns:
Dict: dictionary in which gradient norms are stored.
"""
if isinstance(model, (DataParallel, DistributedDataParallel)):
model = model.module
total_norm = 0.0
grad_norm = {}
for tag, value in model.named_parameters():
tag = tag.replace(".", "/")
metrics_tag = f"{prefix}/{tag}"
param_norm = value.grad.data.norm(norm_type).item()
total_norm += param_norm**norm_type
grad_norm[metrics_tag] = param_norm
total_norm = total_norm**(1.0 / norm_type)
metrics_tag = f"{prefix}/total"
grad_norm[metrics_tag] = total_norm
return grad_norm
def trace_model(
model: Model,
predict_fn: Callable,
batch=None,
method_name: str = "forward",
mode: str = "eval",
requires_grad: bool = False,
opt_level: str = None,
device: Device = "cpu",
predict_params: dict = None,
) -> jit.ScriptModule:
"""Traces model using runner and batch.
Args:
model: Model to trace
predict_fn: Function to run prediction with the model provided,
takes model, inputs parameters
batch: Batch to trace the model
method_name: Model's method name that will be
used as entrypoint during tracing
mode: Mode for model to trace (``train`` or ``eval``)
requires_grad: Flag to use grads
opt_level: Apex FP16 init level, optional
device: Torch device
predict_params: additional parameters for model forward
Returns:
jit.ScriptModule: Traced model
Raises:
ValueError: if both batch and predict_fn must be specified or
mode is not in 'eval' or 'train'.
"""
if batch is None or predict_fn is None:
raise ValueError("Both batch and predict_fn must be specified.")
if mode not in ["train", "eval"]:
raise ValueError(f"Unknown mode '{mode}'. Must be 'eval' or 'train'")
predict_params = predict_params or {}
tracer = _TracingModelWrapper(model, method_name)
if opt_level is not None:
assert_fp16_available()
# If traced in AMP we need to initialize the model before calling
# the jit
# https://github.com/NVIDIA/apex/issues/303#issuecomment-493142950
from apex import amp
model = model.to(device)
model = amp.initialize(model, optimizers=None, opt_level=opt_level)
getattr(model, mode)()
set_requires_grad(model, requires_grad=requires_grad)
predict_fn(tracer, batch, **predict_params)
return tracer.tracing_result
def process_model_params(
model: Model,
layerwise_params: Dict[str, dict] = None,
no_bias_weight_decay: bool = True,
lr_scaling: float = 1.0,
) -> List[Union[torch.nn.Parameter, dict]]:
"""Gains model parameters for ``torch.optim.Optimizer``.
Args:
model: Model to process
layerwise_params: Order-sensitive dict where
each key is regex pattern and values are layer-wise options
for layers matching with a pattern
no_bias_weight_decay: If true, removes weight_decay
for all ``bias`` parameters in the model
lr_scaling: layer-wise learning rate scaling,
if 1.0, learning rates will not be scaled
Returns:
iterable: parameters for an optimizer
Example::
>>> model = catalyst.contrib.models.segmentation.ResnetUnet()
>>> layerwise_params = collections.OrderedDict([
>>> ("conv1.*", dict(lr=0.001, weight_decay=0.0003)),
>>> ("conv.*", dict(lr=0.002))
>>> ])
>>> params = process_model_params(model, layerwise_params)
>>> optimizer = torch.optim.Adam(params, lr=0.0003)
"""
params = list(model.named_parameters())
layerwise_params = layerwise_params or collections.OrderedDict()
model_params = []
for name, parameters in params:
options = {}
for pattern, pattern_options in layerwise_params.items():
if re.match(pattern, name) is not None:
# all new LR rules write on top of the old ones
options = merge_dicts(options, pattern_options)
# no bias decay from https://arxiv.org/abs/1812.01187
if no_bias_weight_decay and name.endswith("bias"):
options["weight_decay"] = 0.0
# lr linear scaling from https://arxiv.org/pdf/1706.02677.pdf
if "lr" in options:
options["lr"] *= lr_scaling
model_params.append({"params": parameters, **options})
return model_params
def get_requires_grad(model: Model):
"""Gets the ``requires_grad`` value for all model parameters.
Example::
>>> model = SimpleModel()
>>> requires_grad = get_requires_grad(model)
Args:
model: model
Returns:
requires_grad: value
"""
requires_grad = {}
for name, param in model.named_parameters():
requires_grad[name] = param.requires_grad
return requires_grad
def zero_grad(self, loss: torch.Tensor, model: Model,
optimizer: Optimizer) -> None:
"""Abstraction over ``model.zero_grad()`` step."""
model.zero_grad()
def process_components(
model: Model,
criterion: Criterion = None,
optimizer: Optimizer = None,
scheduler: Scheduler = None,
distributed_params: Dict = None,
device: Device = None,
) -> Tuple[Model, Criterion, Optimizer, Scheduler, Device]:
"""
Returns the processed model, criterion, optimizer, scheduler and device.
Args:
model: torch model
criterion: criterion function
optimizer: optimizer
scheduler: scheduler
distributed_params (dict, optional): dict with the parameters
for distributed and FP16 method
device (Device, optional): device
Returns:
tuple with processed model, criterion, optimizer, scheduler and device.
Raises:
ValueError: if device is None and TPU available,
for using TPU need to manualy move model/optimizer/scheduler
to a TPU device and pass device to a function.
NotImplementedError: if model is not nn.Module or dict for multi-gpu,
nn.ModuleDict for DataParallel not implemented yet
"""
distributed_params = distributed_params or {}
distributed_params = copy.deepcopy(distributed_params)
distributed_params.update(get_distributed_params())
if device is None and IS_XLA_AVAILABLE:
raise ValueError(
"TPU device is available. "
"Please move model, optimizer and scheduler (if present) "
"to TPU device manualy and specify a device or "
"use CPU device.")
if device is None:
device = get_device()
elif isinstance(device, str):
device = torch.device(device)
is_apex_enabled = (distributed_params.pop("apex", False)
and check_apex_available())
is_amp_enabled = (distributed_params.get("amp", False)
and check_amp_available())
if is_apex_enabled and is_amp_enabled:
raise ValueError("Both NVidia Apex and Torch.Amp are enabled. "
"You must choose only one mixed precision backend")
model: Model = maybe_recursive_call(model, "to", device=device)
if check_ddp_wrapped(model):
pass
# distributed data parallel run (ddp) (with apex support)
elif get_rank() >= 0:
assert isinstance(
model,
nn.Module), "Distributed training is not available for KV model"
local_rank = distributed_params.pop("local_rank", 0) or 0
device = f"cuda:{local_rank}"
model = maybe_recursive_call(model, "to", device=device)
syncbn = distributed_params.pop("syncbn", False)
if is_apex_enabled:
import apex
if syncbn:
model = apex.parallel.convert_syncbn_model(model)
model, optimizer = initialize_apex(model, optimizer,
**distributed_params)
model = apex.parallel.DistributedDataParallel(model)
else:
if syncbn:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
model = nn.parallel.DistributedDataParallel(
model, device_ids=[local_rank], output_device=local_rank)
# data parallel run (dp) (with apex support)
else:
# apex issue https://github.com/deepset-ai/FARM/issues/210
use_apex = (is_apex_enabled and torch.cuda.device_count() == 1) or (
is_apex_enabled and torch.cuda.device_count() > 1
and distributed_params.get("opt_level", "O0") == "O1")
if use_apex:
assert isinstance(
model,
nn.Module), "Apex training is not available for KV model"
model, optimizer = initialize_apex(model, optimizer,
**distributed_params)
if (torch.cuda.device_count() > 1 and device.type != "cpu"
and device.index is None):
if isinstance(model, nn.Module):
model = nn.DataParallel(model)
elif isinstance(model, dict):
model = {k: nn.DataParallel(v) for k, v in model.items()}
else:
raise NotImplementedError()
model: Model = maybe_recursive_call(model, "to", device=device)
return model, criterion, optimizer, scheduler, device
