Skip to content

Commit

Permalink
fix PR comments and modify opset names enum (change to str enum and r…
Browse files Browse the repository at this point in the history
…emoved OPSET prefix)
  • Loading branch information
ofirgo committed Jan 5, 2025
1 parent 151522d commit 37c09bb
Show file tree
Hide file tree
Showing 29 changed files with 686 additions and 656 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@

from model_compression_toolkit import DefaultDict
from model_compression_toolkit.constants import MIN_THRESHOLD
from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework.attach2fw import \
CustomOpsetLayers


class QuantizationErrorMethod(Enum):
Expand Down Expand Up @@ -88,7 +90,7 @@ class QuantizationConfig:
concat_threshold_update: bool = False
activation_bias_correction: bool = False
activation_bias_correction_threshold: float = 0.0
custom_tpc_opset_to_layer: Optional[Dict[str, Tuple[List[Any], Optional[Dict[str, DefaultDict]]]]] = None
custom_tpc_opset_to_layer: Optional[Dict[str, CustomOpsetLayers]] = None


# Default quantization configuration the library use.
Expand Down
114 changes: 57 additions & 57 deletions model_compression_toolkit/target_platform_capabilities/schema/v1.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,63 +24,63 @@


class OperatorSetNames(str, Enum):
OPSET_CONV = "Conv"
OPSET_DEPTHWISE_CONV = "DepthwiseConv2D"
OPSET_CONV_TRANSPOSE = "ConvTranspose"
OPSET_FULLY_CONNECTED = "FullyConnected"
OPSET_CONCATENATE = "Concatenate"
OPSET_STACK = "Stack"
OPSET_UNSTACK = "Unstack"
OPSET_GATHER = "Gather"
OPSET_EXPAND = "Expend"
OPSET_BATCH_NORM = "BatchNorm"
OPSET_L2NORM = "L2Norm"
OPSET_RELU = "ReLU"
OPSET_RELU6 = "ReLU6"
OPSET_LEAKY_RELU = "LEAKYReLU"
OPSET_ELU = "Elu"
OPSET_HARD_TANH = "HardTanh"
OPSET_ADD = "Add"
OPSET_SUB = "Sub"
OPSET_MUL = "Mul"
OPSET_DIV = "Div"
OPSET_MIN = "Min"
OPSET_MAX = "Max"
OPSET_PRELU = "PReLU"
OPSET_ADD_BIAS = "AddBias"
OPSET_SWISH = "Swish"
OPSET_SIGMOID = "Sigmoid"
OPSET_SOFTMAX = "Softmax"
OPSET_LOG_SOFTMAX = "LogSoftmax"
OPSET_TANH = "Tanh"
OPSET_GELU = "Gelu"
OPSET_HARDSIGMOID = "HardSigmoid"
OPSET_HARDSWISH = "HardSwish"
OPSET_FLATTEN = "Flatten"
OPSET_GET_ITEM = "GetItem"
OPSET_RESHAPE = "Reshape"
OPSET_UNSQUEEZE = "Unsqueeze"
OPSET_SQUEEZE = "Squeeze"
OPSET_PERMUTE = "Permute"
OPSET_TRANSPOSE = "Transpose"
OPSET_DROPOUT = "Dropout"
OPSET_SPLIT_CHUNK = "SplitChunk"
OPSET_MAXPOOL = "MaxPool"
OPSET_AVGPOOL = "AvgPool"
OPSET_SIZE = "Size"
OPSET_SHAPE = "Shape"
OPSET_EQUAL = "Equal"
OPSET_ARGMAX = "ArgMax"
OPSET_TOPK = "TopK"
OPSET_FAKE_QUANT = "FakeQuant"
OPSET_COMBINED_NON_MAX_SUPPRESSION = "CombinedNonMaxSuppression"
OPSET_ZERO_PADDING2d = "ZeroPadding2D"
OPSET_CAST = "Cast"
OPSET_RESIZE = "Resize"
OPSET_PAD = "Pad"
OPSET_FOLD = "Fold"
OPSET_STRIDED_SLICE = "StridedSlice"
OPSET_SSD_POST_PROCESS = "SSDPostProcess"
CONV = "Conv"
DEPTHWISE_CONV = "DepthwiseConv2D"
CONV_TRANSPOSE = "ConvTranspose"
FULLY_CONNECTED = "FullyConnected"
CONCATENATE = "Concatenate"
STACK = "Stack"
UNSTACK = "Unstack"
GATHER = "Gather"
EXPAND = "Expend"
BATCH_NORM = "BatchNorm"
L2NORM = "L2Norm"
RELU = "ReLU"
RELU6 = "ReLU6"
LEAKY_RELU = "LEAKYReLU"
ELU = "Elu"
HARD_TANH = "HardTanh"
ADD = "Add"
SUB = "Sub"
MUL = "Mul"
DIV = "Div"
MIN = "Min"
MAX = "Max"
PRELU = "PReLU"
ADD_BIAS = "AddBias"
SWISH = "Swish"
SIGMOID = "Sigmoid"
SOFTMAX = "Softmax"
LOG_SOFTMAX = "LogSoftmax"
TANH = "Tanh"
GELU = "Gelu"
HARDSIGMOID = "HardSigmoid"
HARDSWISH = "HardSwish"
FLATTEN = "Flatten"
GET_ITEM = "GetItem"
RESHAPE = "Reshape"
UNSQUEEZE = "Unsqueeze"
SQUEEZE = "Squeeze"
PERMUTE = "Permute"
TRANSPOSE = "Transpose"
DROPOUT = "Dropout"
SPLIT_CHUNK = "SplitChunk"
MAXPOOL = "MaxPool"
AVGPOOL = "AvgPool"
SIZE = "Size"
SHAPE = "Shape"
EQUAL = "Equal"
ARGMAX = "ArgMax"
TOPK = "TopK"
FAKE_QUANT = "FakeQuant"
COMBINED_NON_MAX_SUPPRESSION = "CombinedNonMaxSuppression"
ZERO_PADDING2D = "ZeroPadding2D"
CAST = "Cast"
RESIZE = "Resize"
PAD = "Pad"
FOLD = "Fold"
STRIDED_SLICE = "StridedSlice"
SSD_POST_PROCESS = "SSDPostProcess"

@classmethod
def get_values(cls):
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import Dict, Tuple, List, Any, Optional
from typing import Dict, NamedTuple, List, Any, Optional

from model_compression_toolkit import DefaultDict
from model_compression_toolkit.logger import Logger
Expand All @@ -8,6 +8,20 @@
OperationsSetToLayers


class CustomOpsetLayers(NamedTuple):
"""
This struct is used to define a set of operators to map to a custom operator set defined in the TPC.
Args:
operators: a list of framework operators to map to a certain custom opset name.
attr_mapping: a mapping between an opset name to a mapping between its attributes' general names and names in
the framework.
"""

operators: List
attr_mapping: Optional[Dict[str, DefaultDict]]


class AttachTpcToFramework:

def __init__(self):
Expand All @@ -19,7 +33,7 @@ def __init__(self):
self._opset2attr_mapping = None # Mapping of operation sets to their corresponding framework-specific layers

def attach(self, tpc_model: TargetPlatformModel,
custom_opset2layer: Optional[Dict[str, Tuple[List[Any], Optional[Dict[str, DefaultDict]]]]] = None
custom_opset2layer: Optional[Dict[str, CustomOpsetLayers]] = None
) -> TargetPlatformCapabilities:
"""
Attaching a TargetPlatformModel which includes a platform capabilities description to specific
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -44,85 +44,85 @@ def __init__(self):
super().__init__()

self._opset2layer = {
OperatorSetNames.OPSET_CONV: [Conv2D, tf.nn.conv2d],
OperatorSetNames.OPSET_DEPTHWISE_CONV: [DepthwiseConv2D, tf.nn.depthwise_conv2d],
OperatorSetNames.OPSET_CONV_TRANSPOSE: [Conv2DTranspose, tf.nn.conv2d_transpose],
OperatorSetNames.OPSET_FULLY_CONNECTED: [Dense],
OperatorSetNames.OPSET_CONCATENATE: [tf.concat, Concatenate],
OperatorSetNames.OPSET_STACK: [tf.stack],
OperatorSetNames.OPSET_UNSTACK: [tf.unstack],
OperatorSetNames.OPSET_GATHER: [tf.gather, tf.compat.v1.gather],
OperatorSetNames.OPSET_EXPAND: [],
OperatorSetNames.OPSET_BATCH_NORM: [BatchNormalization, tf.nn.batch_normalization],
OperatorSetNames.OPSET_RELU: [tf.nn.relu, ReLU, LayerFilterParams(Activation, activation="relu")],
OperatorSetNames.OPSET_RELU6: [tf.nn.relu6],
OperatorSetNames.OPSET_LEAKY_RELU: [tf.nn.leaky_relu, LeakyReLU, LayerFilterParams(Activation, activation="leaky_relu")],
OperatorSetNames.OPSET_HARD_TANH: [LayerFilterParams(Activation, activation="hard_tanh")],
OperatorSetNames.OPSET_ADD: [tf.add, Add],
OperatorSetNames.OPSET_SUB: [tf.subtract, Subtract],
OperatorSetNames.OPSET_MUL: [tf.math.multiply, Multiply],
OperatorSetNames.OPSET_DIV: [tf.math.divide, tf.math.truediv],
OperatorSetNames.OPSET_MIN: [tf.math.minimum, Minimum],
OperatorSetNames.OPSET_MAX: [tf.math.maximum, Maximum],
OperatorSetNames.OPSET_PRELU: [PReLU],
OperatorSetNames.OPSET_SWISH: [tf.nn.swish, LayerFilterParams(Activation, activation="swish")],
OperatorSetNames.OPSET_HARDSWISH: [LayerFilterParams(Activation, activation="hard_swish")],
OperatorSetNames.OPSET_SIGMOID: [tf.nn.sigmoid, LayerFilterParams(Activation, activation="sigmoid")],
OperatorSetNames.OPSET_TANH: [tf.nn.tanh, LayerFilterParams(Activation, activation="tanh")],
OperatorSetNames.OPSET_GELU: [tf.nn.gelu, LayerFilterParams(Activation, activation="gelu")],
OperatorSetNames.OPSET_HARDSIGMOID: [tf.keras.activations.hard_sigmoid,
LayerFilterParams(Activation, activation="hard_sigmoid")],
OperatorSetNames.OPSET_FLATTEN: [Flatten],
OperatorSetNames.OPSET_GET_ITEM: [tf.__operators__.getitem],
OperatorSetNames.OPSET_RESHAPE: [Reshape, tf.reshape],
OperatorSetNames.OPSET_PERMUTE: [Permute],
OperatorSetNames.OPSET_TRANSPOSE: [tf.transpose],
OperatorSetNames.OPSET_UNSQUEEZE: [tf.expand_dims],
OperatorSetNames.OPSET_SQUEEZE: [tf.squeeze],
OperatorSetNames.OPSET_DROPOUT: [Dropout],
OperatorSetNames.OPSET_SPLIT_CHUNK: [tf.split],
OperatorSetNames.OPSET_MAXPOOL: [MaxPooling2D, tf.nn.avg_pool2d],
OperatorSetNames.OPSET_AVGPOOL: [AveragePooling2D],
OperatorSetNames.OPSET_SIZE: [tf.size],
OperatorSetNames.OPSET_RESIZE: [tf.image.resize],
OperatorSetNames.OPSET_PAD: [tf.pad, Cropping2D],
OperatorSetNames.OPSET_FOLD: [tf.space_to_batch_nd],
OperatorSetNames.OPSET_SHAPE: [tf.shape, tf.compat.v1.shape],
OperatorSetNames.OPSET_EQUAL: [tf.math.equal],
OperatorSetNames.OPSET_ARGMAX: [tf.math.argmax],
OperatorSetNames.OPSET_TOPK: [tf.nn.top_k],
OperatorSetNames.OPSET_FAKE_QUANT: [tf.quantization.fake_quant_with_min_max_vars],
OperatorSetNames.OPSET_COMBINED_NON_MAX_SUPPRESSION: [tf.image.combined_non_max_suppression],
OperatorSetNames.OPSET_ZERO_PADDING2d: [ZeroPadding2D],
OperatorSetNames.OPSET_CAST: [tf.cast],
OperatorSetNames.OPSET_STRIDED_SLICE: [tf.strided_slice],
OperatorSetNames.OPSET_ELU: [tf.nn.elu, LayerFilterParams(Activation, activation="elu")],
OperatorSetNames.OPSET_SOFTMAX: [tf.nn.softmax, Softmax,
LayerFilterParams(Activation, activation="softmax")],
OperatorSetNames.OPSET_LOG_SOFTMAX: [tf.nn.log_softmax],
OperatorSetNames.OPSET_ADD_BIAS: [tf.nn.bias_add],
OperatorSetNames.OPSET_L2NORM: [tf.math.l2_normalize],
OperatorSetNames.CONV: [Conv2D, tf.nn.conv2d],
OperatorSetNames.DEPTHWISE_CONV: [DepthwiseConv2D, tf.nn.depthwise_conv2d],
OperatorSetNames.CONV_TRANSPOSE: [Conv2DTranspose, tf.nn.conv2d_transpose],
OperatorSetNames.FULLY_CONNECTED: [Dense],
OperatorSetNames.CONCATENATE: [tf.concat, Concatenate],
OperatorSetNames.STACK: [tf.stack],
OperatorSetNames.UNSTACK: [tf.unstack],
OperatorSetNames.GATHER: [tf.gather, tf.compat.v1.gather],
OperatorSetNames.EXPAND: [],
OperatorSetNames.BATCH_NORM: [BatchNormalization, tf.nn.batch_normalization],
OperatorSetNames.RELU: [tf.nn.relu, ReLU, LayerFilterParams(Activation, activation="relu")],
OperatorSetNames.RELU6: [tf.nn.relu6],
OperatorSetNames.LEAKY_RELU: [tf.nn.leaky_relu, LeakyReLU, LayerFilterParams(Activation, activation="leaky_relu")],
OperatorSetNames.HARD_TANH: [LayerFilterParams(Activation, activation="hard_tanh")],
OperatorSetNames.ADD: [tf.add, Add],
OperatorSetNames.SUB: [tf.subtract, Subtract],
OperatorSetNames.MUL: [tf.math.multiply, Multiply],
OperatorSetNames.DIV: [tf.math.divide, tf.math.truediv],
OperatorSetNames.MIN: [tf.math.minimum, Minimum],
OperatorSetNames.MAX: [tf.math.maximum, Maximum],
OperatorSetNames.PRELU: [PReLU],
OperatorSetNames.SWISH: [tf.nn.swish, LayerFilterParams(Activation, activation="swish")],
OperatorSetNames.HARDSWISH: [LayerFilterParams(Activation, activation="hard_swish")],
OperatorSetNames.SIGMOID: [tf.nn.sigmoid, LayerFilterParams(Activation, activation="sigmoid")],
OperatorSetNames.TANH: [tf.nn.tanh, LayerFilterParams(Activation, activation="tanh")],
OperatorSetNames.GELU: [tf.nn.gelu, LayerFilterParams(Activation, activation="gelu")],
OperatorSetNames.HARDSIGMOID: [tf.keras.activations.hard_sigmoid,
LayerFilterParams(Activation, activation="hard_sigmoid")],
OperatorSetNames.FLATTEN: [Flatten],
OperatorSetNames.GET_ITEM: [tf.__operators__.getitem],
OperatorSetNames.RESHAPE: [Reshape, tf.reshape],
OperatorSetNames.PERMUTE: [Permute],
OperatorSetNames.TRANSPOSE: [tf.transpose],
OperatorSetNames.UNSQUEEZE: [tf.expand_dims],
OperatorSetNames.SQUEEZE: [tf.squeeze],
OperatorSetNames.DROPOUT: [Dropout],
OperatorSetNames.SPLIT_CHUNK: [tf.split],
OperatorSetNames.MAXPOOL: [MaxPooling2D, tf.nn.avg_pool2d],
OperatorSetNames.AVGPOOL: [AveragePooling2D],
OperatorSetNames.SIZE: [tf.size],
OperatorSetNames.RESIZE: [tf.image.resize],
OperatorSetNames.PAD: [tf.pad, Cropping2D],
OperatorSetNames.FOLD: [tf.space_to_batch_nd],
OperatorSetNames.SHAPE: [tf.shape, tf.compat.v1.shape],
OperatorSetNames.EQUAL: [tf.math.equal],
OperatorSetNames.ARGMAX: [tf.math.argmax],
OperatorSetNames.TOPK: [tf.nn.top_k],
OperatorSetNames.FAKE_QUANT: [tf.quantization.fake_quant_with_min_max_vars],
OperatorSetNames.COMBINED_NON_MAX_SUPPRESSION: [tf.image.combined_non_max_suppression],
OperatorSetNames.ZERO_PADDING2D: [ZeroPadding2D],
OperatorSetNames.CAST: [tf.cast],
OperatorSetNames.STRIDED_SLICE: [tf.strided_slice],
OperatorSetNames.ELU: [tf.nn.elu, LayerFilterParams(Activation, activation="elu")],
OperatorSetNames.SOFTMAX: [tf.nn.softmax, Softmax,
LayerFilterParams(Activation, activation="softmax")],
OperatorSetNames.LOG_SOFTMAX: [tf.nn.log_softmax],
OperatorSetNames.ADD_BIAS: [tf.nn.bias_add],
OperatorSetNames.L2NORM: [tf.math.l2_normalize],
}

if FOUND_SONY_CUSTOM_LAYERS:
self._opset2layer[OperatorSetNames.OPSET_SSD_POST_PROCESS] = [SSDPostProcess]
self._opset2layer[OperatorSetNames.SSD_POST_PROCESS] = [SSDPostProcess]
else:
# If Custom layers is not installed then we don't want the user to fail, but just ignore custom layers
# in the initialized framework TPC
self._opset2layer[OperatorSetNames.OPSET_SSD_POST_PROCESS] = []
self._opset2layer[OperatorSetNames.SSD_POST_PROCESS] = []

self._opset2attr_mapping = {
OperatorSetNames.OPSET_CONV: {
OperatorSetNames.CONV: {
KERNEL_ATTR: DefaultDict(default_value=KERAS_KERNEL),
BIAS_ATTR: DefaultDict(default_value=BIAS)},
OperatorSetNames.OPSET_CONV_TRANSPOSE: {
OperatorSetNames.CONV_TRANSPOSE: {
KERNEL_ATTR: DefaultDict(default_value=KERAS_KERNEL),
BIAS_ATTR: DefaultDict(default_value=BIAS)},
OperatorSetNames.OPSET_DEPTHWISE_CONV: {
OperatorSetNames.DEPTHWISE_CONV: {
KERNEL_ATTR: DefaultDict({
DepthwiseConv2D: KERAS_DEPTHWISE_KERNEL,
tf.nn.depthwise_conv2d: KERAS_DEPTHWISE_KERNEL}, default_value=KERAS_KERNEL),
BIAS_ATTR: DefaultDict(default_value=BIAS)},
OperatorSetNames.OPSET_FULLY_CONNECTED: {
OperatorSetNames.FULLY_CONNECTED: {
KERNEL_ATTR: DefaultDict(default_value=KERAS_KERNEL),
BIAS_ATTR: DefaultDict(default_value=BIAS)}}
Loading

0 comments on commit 37c09bb

Please sign in to comment.