mirror of
https://github.com/mit-han-lab/tinyengine.git
synced 2025-05-09 17:11:01 +08:00
119 lines
3.4 KiB
Python
119 lines
3.4 KiB
Python
# ----------------------------------------------------------------------
|
|
# Project: TinyEngine
|
|
# Title: constant.py
|
|
#
|
|
# Reference papers:
|
|
# - MCUNet: Tiny Deep Learning on IoT Device, NeurIPS 2020
|
|
# - MCUNetV2: Memory-Efficient Patch-based Inference for Tiny Deep Learning, NeurIPS 2021
|
|
# - MCUNetV3: On-Device Training Under 256KB Memory, NeurIPS 2022
|
|
# Contact authors:
|
|
# - Wei-Ming Chen, wmchen@mit.edu
|
|
# - Wei-Chen Wang, wweichen@mit.edu
|
|
# - Ji Lin, jilin@mit.edu
|
|
# - Ligeng Zhu, ligeng@mit.edu
|
|
# - Song Han, songhan@mit.edu
|
|
#
|
|
# Target ISA: ARMv7E-M
|
|
# ----------------------------------------------------------------------
|
|
|
|
TTYPE_INFERNECE = "inference"
|
|
|
|
op_name_translation = {
|
|
"nn.conv2d": ["CONV_2D", "DEPTHWISE_CONV_2D", "GROUP_CONV"],
|
|
"nn.mcuconv2d": ["CONV_2D", "DEPTHWISE_CONV_2D", "GROUP_CONV"],
|
|
"nn.mcuadd": "ADD",
|
|
"add": "ADD",
|
|
"ResizeNearestNeighbor": "UPSAMPLE",
|
|
}
|
|
|
|
SKIP_OPs = {"QUANTIZE", "DEQUANTIZE", "RESHAPE"} # TODO: Handle RESHAPE during codegen
|
|
|
|
FUSE_TILE_STR = "FUSE_TILE"
|
|
FUSE_INT8CAST_STR = "FUSE_CAST"
|
|
FUSE_SGD_UPDATE_STR = "INPLACE_SGD"
|
|
FUSE_INT8CASTSLICE_STR = "FUSE_CAST_SLICE"
|
|
FUSE_WHERE_ZEROSSTR = "FUSE_WHERE_ZEROS"
|
|
INPLACE_WHERE_STR = "INPLACE_WHERE"
|
|
INPLACE_MUL_STR = "INPLACE_MUL"
|
|
INPLACE_DEPTHWISE_STRIDE2_STR = "INPLACE_DEPTHWISE_STRIDE2_MEM"
|
|
REORDER_STR = "REORDER_GCONV_TCONV"
|
|
FUSHION_CONFIG = {
|
|
FUSE_TILE_STR: True,
|
|
FUSE_INT8CAST_STR: True,
|
|
FUSE_INT8CASTSLICE_STR: True,
|
|
FUSE_SGD_UPDATE_STR: True,
|
|
FUSE_WHERE_ZEROSSTR: True,
|
|
INPLACE_WHERE_STR: True,
|
|
INPLACE_MUL_STR: True,
|
|
REORDER_STR: True,
|
|
INPLACE_DEPTHWISE_STRIDE2_STR: True,
|
|
}
|
|
INFERECE_WEIGHT_SIZE = 1
|
|
INFERECE_INT8_SIZE = 1
|
|
USE_BIT_MASK = True
|
|
USE_TTE_INT8 = True
|
|
# end: for forcelly cast to fp32
|
|
TTYPE_STATIC_WEIGHT = "static_weights"
|
|
TTYPE_TRAINING_GRADIENT = "training_gradient"
|
|
TTYPE_TRAINING_ACTIVATION = "training_activation"
|
|
TTYPE_TRAINING_WEIGHTS = "training_weight"
|
|
TTYPE_INFERNECE = "inference"
|
|
|
|
# template 3
|
|
FIGURE_CONFIG = {
|
|
"TRAIN_TENSOR_COLOR": "#68ac14", # green
|
|
"TRAIN_WEIGHT_COLOR": "#61a78f", # light green:
|
|
"TRAIN_ACTIVATION_COLOR": "#b9c8d8", # light blue
|
|
"TRAIN_GRADIENT_COLOR": "#35619f", # blue
|
|
"INFERENCE_COLOR": "#eab11f", # yellow
|
|
"FIGURE_W_INCH": 19,
|
|
"FIGURE_H_INCH": 7,
|
|
"DPI": 800,
|
|
"Y_STEP": 96,
|
|
"Y_MAX": 384,
|
|
"X_STEP": 30,
|
|
"X_MAX": 270,
|
|
"SHOW_INDEX": False,
|
|
"FONT_SIZE": 32,
|
|
"INDEX_FONT_SIZE": 24,
|
|
}
|
|
|
|
|
|
op_name_translation = {
|
|
"nn.conv2d": ["CONV_2D", "DEPTHWISE_CONV_2D", "GROUP_CONV"],
|
|
"nn.mcuconv2d": ["CONV_2D", "DEPTHWISE_CONV_2D", "GROUP_CONV"],
|
|
"nn.mcuadd": "ADD",
|
|
"nn.dense": "DENSE",
|
|
"nn.matmul": "MAT_MUL",
|
|
"add": "ADD",
|
|
"cast": "CAST",
|
|
"cast_like": "CAST_LIKE",
|
|
"divide": "DIV",
|
|
"multiply": "MUL",
|
|
"less": "LESS",
|
|
"less_equal": "LESS",
|
|
"greater": "GREATER",
|
|
"greater_equal": "GREATER",
|
|
"where": "WHERE",
|
|
"exp": "EXP",
|
|
"sum": "SUM",
|
|
"subtract": "SUBTRACT",
|
|
"tile": "TILE",
|
|
"reshape": "RESHAPE",
|
|
"reshape_like": "RESHAPE_LIKE",
|
|
"nn.relu": "RELU",
|
|
"strided_slice": "STRIDED_SLICE",
|
|
"transpose": "TRANSPOSE",
|
|
"nn.conv2d_transpose": "TRANSPOSE_CONV_2D",
|
|
"nn.bias_add": "BIAS_ADD",
|
|
"negative": "NEGATIVE",
|
|
"zeros_like": "ZEROS_LIKE",
|
|
"zeros": "ZEROS",
|
|
"ones_like": "ONES_LIKE",
|
|
"ones": "ONES",
|
|
"collapse_sum_like": "COLLAPSE_SUM_LIKE",
|
|
"nn.log_softmax": "LOG_SOFTMAX",
|
|
"nn.cross_entropy_with_logits": "NLL_LOSS",
|
|
"max": "MAX",
|
|
}
|