`neural_compressor.experimental.common.criterion`¶

Initialize critetion classes.

Classes includes:: TensorFlowCrossEntropyLoss, PyTorchCrossEntropyLoss, TensorflowKnowledgeDistillationLoss, PyTorchKnowledgeDistillationLoss, PyTorchIntermediateLayersKnowledgeDistillationLoss.

Module Contents¶

Classes¶

`TensorflowCriterions`	Record criterions in TensorflowCriterions class.
`PyTorchCriterions`	Record criterions in PyTorchCriterions class.
`Criterions`	Integrate criterions of different framework.
`TensorFlowCrossEntropyLoss`	TensorFlow CrossEntropyLoss criterion.
`TensorFlowSparseCategoricalCrossentropy`	TensorFlow SparseCategoricalCrossentropyLoss criterion.
`PyTorchCrossEntropyLoss`	PyTorch CrossEntropyLoss criterion.
`KnowledgeDistillationFramework`	Knowledge Distillation Framework.
`KnowledgeDistillationLoss`	Initialize the KnowledgeDistillationLoss class.
`PyTorchKnowledgeDistillationLoss`	The PyTorchKnowledgeDistillationLoss class inherits from KnowledgeDistillationLoss.
`PyTorchKnowledgeDistillationLossWrapper`	PyTorchKnowledgeDistillationLossWrapper wraps PyTorchKnowledgeDistillationLoss.
`TensorflowKnowledgeDistillationLoss`	The TensorflowKnowledgeDistillationLoss class inherits from KnowledgeDistillationLoss.
`TensorflowKnowledgeDistillationLossWrapper`	TensorflowKnowledgeDistillationLossWrapper wraps TensorflowKnowledgeDistillationLoss.
`TensorflowKnowledgeDistillationLossExternal`	TensorflowKnowledgeDistillationLossExternal inherits from KnowledgeDistillationLoss.
`IntermediateLayersKnowledgeDistillationLoss`	The IntermediateLayersKnowledgeDistillationLoss class inherits from KnowledgeDistillationLoss.
`PyTorchIntermediateLayersKnowledgeDistillationLoss`	PyTorch Intermediate Layers Knowledge Distillation Loss.
`PyTorchIntermediateLayersKnowledgeDistillationLossWrapper`	PyTorch Intermediate Layers Knowledge Distillation Loss Wrapper.
`SelfKnowledgeDistillationLoss`	SelfKnowledge Distillation Loss.
`PyTorchSelfKnowledgeDistillationLoss`	PyTorch SelfKnowledge Distillation Loss.
`PyTorchSelfKnowledgeDistillationLossWrapper`	PyTorch SelfKnowledge Distillation Loss Wrapper.

Functions¶

criterion_registry(criterion_type, framework)

Use to register criterion classes in registry_criterions.

class neural_compressor.experimental.common.criterion.TensorflowCriterions¶

Bases: object

Record criterions in TensorflowCriterions class.

class neural_compressor.experimental.common.criterion.PyTorchCriterions¶

Bases: object

Record criterions in PyTorchCriterions class.

class neural_compressor.experimental.common.criterion.Criterions(framework)¶

Bases: object

Integrate criterions of different framework.

register(name, criterion_cls)¶

Register criterion name and result in existing criterion class.

Parameters:

name (string) – criterion name/type.
criterion_cls (string) – criterion class.

neural_compressor.experimental.common.criterion.criterion_registry(criterion_type, framework)¶

Use to register criterion classes in registry_criterions.

Parameters:

criterion_type (str) – The string of supported criterion.
framework (str) – The string of supported framework.

Returns:

The class of register.

Return type:

cls

class neural_compressor.experimental.common.criterion.TensorFlowCrossEntropyLoss(param_dict)¶

Bases: object

TensorFlow CrossEntropyLoss criterion.

class neural_compressor.experimental.common.criterion.TensorFlowSparseCategoricalCrossentropy(param_dict)¶

Bases: object

TensorFlow SparseCategoricalCrossentropyLoss criterion.

class neural_compressor.experimental.common.criterion.PyTorchCrossEntropyLoss(param_dict)¶

Bases: object

PyTorch CrossEntropyLoss criterion.

class neural_compressor.experimental.common.criterion.KnowledgeDistillationFramework(student_model=None, teacher_model=None)¶

Bases: object

Knowledge Distillation Framework.

property student_model¶: Return student model.

property teacher_model¶: Return teacher model.

class neural_compressor.experimental.common.criterion.KnowledgeDistillationLoss(temperature=1.0, loss_types=['CE', 'CE'], loss_weights=[0.5, 0.5], student_model=None, teacher_model=None)¶

Bases: KnowledgeDistillationFramework

Initialize the KnowledgeDistillationLoss class.

abstract teacher_model_forward(input, teacher_model=None)¶

Define parameters for teacher_model_forward function.

Parameters:

input (tensor, tuple or dict) – input data.
teacher_model (model, optional) – teacher model. Defaults to None.

Raises:

NotImplementedError – NotImplementedError

abstract teacher_student_loss_cal(student_outputs, teacher_outputs)¶

Define parameters for teacher_student_loss_cal function.

Parameters:

student_outputs (tensor) – student outputs
teacher_outputs (tensor) – student outputs

Raises:

NotImplementedError – NotImplementedError

abstract student_targets_loss_cal(student_outputs, targets)¶

Define parameters for student_targets_loss_cal function.

Parameters:

student_outputs (tensor) – student outputs
targets (tensor) – groud truth label

Raises:

NotImplementedError – NotImplementedError

loss_cal(student_outputs, targets)¶

Calculate loss of student model.

Parameters:

student_outputs (tensor) – student outputs
targets (tensor) – groud truth label

Returns:

loss

Return type:

tensor

loss_cal_sloss(student_outputs, teacher_outputs, student_loss)¶

Calculate all losses between student model and teacher model.

Parameters:

student_outputs (tensor) – student outputs
teacher_outputs (tensor) – teacher outputs
student_loss (tensor) – student loss

Returns:

loss

Return type:

tensor

class neural_compressor.experimental.common.criterion.PyTorchKnowledgeDistillationLoss(temperature=1.0, loss_types=['CE', 'CE'], loss_weights=[0.5, 0.5], student_model=None, teacher_model=None)¶

Bases: KnowledgeDistillationLoss

The PyTorchKnowledgeDistillationLoss class inherits from KnowledgeDistillationLoss.

SoftCrossEntropy(logits, targets)¶

Return SoftCrossEntropy.

Parameters:

logits (tensor) – output logits
targets (tensor) – ground truth label

Returns:

SoftCrossEntropy

Return type:

tensor

KullbackLeiblerDivergence(logits, targets)¶

Return KullbackLeiblerDivergence.

Parameters:

logits (tensor) – output logits
targets (tensor) – ground truth label

Returns:

KullbackLeiblerDivergence

Return type:

tensor

teacher_model_forward(input, teacher_model=None, device=None)¶

Teacher model forward.

Parameters:

input (tensor) – input data
teacher_model (torch.nn.model, optional) – teacher model. Defaults to None.
device (torch.device, optional) – device. Defaults to None.

Returns:

output

Return type:

tensor

teacher_student_loss_cal(student_outputs, teacher_outputs)¶

Calculate loss between student model and teacher model.

Parameters:

student_outputs (tensor) – student outputs
teacher_outputs (tensor) – teacher outputs

Returns:

loss

Return type:

tensor

student_targets_loss_cal(student_outputs, targets)¶

Calculate loss of student model.

Parameters:

student_outputs (tensor) – student outputs
targets (tensor) – groud truth label

Returns:

loss

Return type:

tensor

class neural_compressor.experimental.common.criterion.PyTorchKnowledgeDistillationLossWrapper(param_dict)¶

Bases: object

PyTorchKnowledgeDistillationLossWrapper wraps PyTorchKnowledgeDistillationLoss.

class neural_compressor.experimental.common.criterion.TensorflowKnowledgeDistillationLoss(temperature=1.0, loss_types=['CE', 'CE'], loss_weights=[0.5, 0.5], student_model=None, teacher_model=None)¶

Bases: KnowledgeDistillationLoss

The TensorflowKnowledgeDistillationLoss class inherits from KnowledgeDistillationLoss.

SoftCrossEntropy(targets, logits)¶

Return SoftCrossEntropy.

Parameters:

logits (tensor) – output logits
targets (tensor) – ground truth label

Returns:

SoftCrossEntropy

Return type:

tensor

teacher_model_forward(input, teacher_model=None)¶

Teacher model forward.

Parameters:

input (tensor) – input data
teacher_model (optional) – teacher model. Defaults to None.
device (torch.device, optional) – device. Defaults to None.

Returns:

output

Return type:

tensor

teacher_student_loss_cal(student_outputs, teacher_outputs)¶

Calculate loss between student model and teacher model.

Parameters:

student_outputs (tensor) – student outputs
teacher_outputs (tensor) – teacher outputs

Returns:

loss

Return type:

tensor

student_targets_loss_cal(student_outputs, targets)¶

Calculate loss of student model.

Parameters:

student_outputs (tensor) – student outputs
targets (tensor) – groud truth label

Returns:

loss

Return type:

tensor

class neural_compressor.experimental.common.criterion.TensorflowKnowledgeDistillationLossWrapper(param_dict)¶

Bases: object

TensorflowKnowledgeDistillationLossWrapper wraps TensorflowKnowledgeDistillationLoss.

class neural_compressor.experimental.common.criterion.TensorflowKnowledgeDistillationLossExternal(temperature=1.0, loss_types=['CE', 'CE'], loss_weights=[0.5, 0.5], student_model=None, teacher_model=None)¶

Bases: KnowledgeDistillationLoss

TensorflowKnowledgeDistillationLossExternal inherits from KnowledgeDistillationLoss.

teacher_model_forward(input, teacher_model=None)¶

Teacher model forward.

Parameters:

input (tensor) – input data
teacher_model (optional) – teacher model. Defaults to None.
device (optional) – device. Defaults to None.

Returns:

output

Return type:

tensor

teacher_student_loss_cal(student_outputs, teacher_outputs)¶

Calculate loss between student model and teacher model.

Parameters:

student_outputs (tensor) – student outputs
teacher_outputs (tensor) – teacher outputs

Returns:

loss

Return type:

tensor

student_targets_loss_cal(student_outputs, targets)¶

Calculate loss of student model.

Parameters:

student_outputs (tensor) – student outputs
targets (tensor) – groud truth label

Returns:

loss

Return type:

tensor

class neural_compressor.experimental.common.criterion.IntermediateLayersKnowledgeDistillationLoss(layer_mappings=[], loss_types=None, loss_weights=None, add_origin_loss=False, student_model=None, teacher_model=None)¶

Bases: KnowledgeDistillationFramework

The IntermediateLayersKnowledgeDistillationLoss class inherits from KnowledgeDistillationLoss.

abstract init_loss_funcs()¶

Init loss funcs.

Raises:: NotImplementedError – NotImplementedError

abstract init_feature_matcher(student_feature, teacher_feature)¶

Init feature matcher.

Raises:: NotImplementedError – NotImplementedError

abstract teacher_model_forward(input, teacher_model=None)¶

Teacher model forward.

Raises:: NotImplementedError – NotImplementedError

abstract loss_cal()¶

Calculate loss.

Raises:: NotImplementedError – NotImplementedError

loss_cal_sloss(student_outputs, teacher_outputs, student_loss)¶

Calculate all losses between student model and teacher model.

Parameters:

student_outputs (tensor) – student outputs
teacher_outputs (tensor) – teacher outputs
student_loss (tensor) – student loss

Returns:

loss

Return type:

tensor

clear_features()¶: Clean features in list.

class neural_compressor.experimental.common.criterion.PyTorchIntermediateLayersKnowledgeDistillationLoss(layer_mappings=[], loss_types=None, loss_weights=None, add_origin_loss=False, student_model=None, teacher_model=None)¶

Bases: IntermediateLayersKnowledgeDistillationLoss

PyTorch Intermediate Layers Knowledge Distillation Loss.

register_hooks_for_models()¶

Register hooks for models to record module output.

Raises:: AttributeError – AttributeError

remove_all_hooks()¶: Remove all hooks.

init_loss_funcs()¶: Init loss funcs.

init_feature_matcher(student_feature, teacher_feature)¶

Init feature matcher.

Parameters:

student_feature (tensor) – student feature
teacher_feature (tensor) – teacher feature

Returns:

pytorch_linear_feature_matcher

teacher_model_forward(input, teacher_model=None, device=None)¶

Define parameters for teacher_model_forward function.

Parameters:

input (tensor, tuple or dict) – input data.
teacher_model (model, optional) – teacher model. Defaults to None.

Raises:

NotImplementedError – NotImplementedError

loss_cal_sloss(student_outputs, teacher_outputs, student_loss)¶

Calculate all losses between student model and teacher model.

Parameters:

student_outputs (tensor) – student outputs
teacher_outputs (tensor) – teacher outputs
student_loss (tensor) – student loss

Returns:

loss

Return type:

tensor

loss_cal()¶

Calculate loss of student model.

Returns:: loss
Return type:: tensor

class neural_compressor.experimental.common.criterion.PyTorchIntermediateLayersKnowledgeDistillationLossWrapper(param_dict)¶

Bases: object

PyTorch Intermediate Layers Knowledge Distillation Loss Wrapper.

class neural_compressor.experimental.common.criterion.SelfKnowledgeDistillationLoss(layer_mappings=[], loss_types=None, loss_weights=None, temperature=1.0, add_origin_loss=False, student_model=None, teacher_model=None)¶

Bases: KnowledgeDistillationFramework

SelfKnowledge Distillation Loss.

abstract init_loss_funcs()¶

Init loss funcs.

Raises:: NotImplementedError – NotImplementedError

abstract teacher_model_forward(input, teacher_model=None)¶

Teacher model forward.

Raises:: NotImplementedError – NotImplementedError

abstract loss_cal(student_outputs)¶

Calculate loss.

Raises:: NotImplementedError – NotImplementedError

loss_cal_sloss(student_outputs, teacher_outputs, student_loss)¶

Calculate all losses between student model and teacher model.

Parameters:

student_outputs (dict) – student outputs
teacher_outputs (dict) – teacher outputs
student_loss (tensor) – student loss

Returns:

loss

Return type:

tensor

class neural_compressor.experimental.common.criterion.PyTorchSelfKnowledgeDistillationLoss(layer_mappings=[], loss_types=None, loss_weights=None, temperature=1.0, add_origin_loss=False, student_model=None, teacher_model=None)¶

Bases: SelfKnowledgeDistillationLoss

PyTorch SelfKnowledge Distillation Loss.

SoftCrossEntropy(logits, targets)¶

Return SoftCrossEntropy.

Parameters:

logits (tensor) – output logits
targets (tensor) – ground truth label

Returns:

SoftCrossEntropy

Return type:

tensor

KullbackLeiblerDivergence(logits, targets)¶

Return KullbackLeiblerDivergence.

Parameters:

logits (tensor) – output logits
targets (tensor) – ground truth label

Returns:

KullbackLeiblerDivergence

Return type:

tensor

L2Divergence(feature1, feature2)¶

Return L2Divergence.

Parameters:

feature1 (tensor) – feature1 value
feature2 (tensor) – feature2 value

Returns:

L2Divergence between feature1 and feature2

Return type:

tensor

init_loss_funcs()¶: Init loss funcs.

loss_cal(student_outputs)¶

Calculate loss of student model.

Parameters:: student_outputs (dict) – student outputs
Returns:: loss
Return type:: tensor

teacher_model_forward(input, teacher_model=None, device=None)¶

Teacher model forward.

Parameters:

input (tensor) – input data
teacher_model (torch.nn.model, optional) – teacher model. Defaults to None.
device (torch.device, optional) – device. Defaults to None.

Returns:

output

Return type:

tensor

class neural_compressor.experimental.common.criterion.PyTorchSelfKnowledgeDistillationLossWrapper(param_dict)¶

Bases: object

PyTorch SelfKnowledge Distillation Loss Wrapper.

neural_compressor.experimental.common.criterion¶

Module Contents¶

Classes¶

Functions¶

`neural_compressor.experimental.common.criterion`¶