mirror of
https://github.com/mit-han-lab/tinyengine.git
synced 2025-05-10 01:18:47 +08:00
168 lines
8.7 KiB
Python
168 lines
8.7 KiB
Python
# ----------------------------------------------------------------------
|
|
# Project: TinyEngine
|
|
# Title: InputResizer.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, arXiv:2206.15472
|
|
# 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
|
|
# ----------------------------------------------------------------------
|
|
|
|
import math
|
|
|
|
|
|
def _find_previous_info(layers, idx):
|
|
for layer in layers:
|
|
info = layer.get_layer_info()
|
|
if info["output_idx"] == idx:
|
|
return info
|
|
|
|
|
|
class InputResizer:
|
|
def __init__(self, layer):
|
|
self.layer = layer
|
|
|
|
def inputResize(self, input_h, input_w):
|
|
for i, layer in enumerate(self.layer):
|
|
layer_info = layer.get_layer_info()
|
|
|
|
previous_layer_info = _find_previous_info(self.layer, layer_info["input_idx"])
|
|
# we need to handle different op
|
|
op_code_str = layer_info["op"]
|
|
if i == 0:
|
|
layer_info["input_h"] = input_h
|
|
layer_info["input_w"] = input_w
|
|
_changeOPTensorSize(self.layer[i], "input", 0, layer_info["input_h"], layer_info["input_w"])
|
|
else:
|
|
if op_code_str == "SE_AVG_POOL_2D":
|
|
SEinput_h = previous_layer_info["output_h"]
|
|
SEinput_w = previous_layer_info["output_w"]
|
|
layer_info["input_h"] = SEinput_h
|
|
layer_info["input_w"] = SEinput_w
|
|
_changeOPTensorSize(self.layer[i], "input", 0, layer_info["input_h"], layer_info["input_w"])
|
|
layer_info["sample_h"] = SEinput_h
|
|
layer_info["sample_w"] = SEinput_w
|
|
else:
|
|
layer_info["input_h"] = previous_layer_info["output_h"]
|
|
layer_info["input_w"] = previous_layer_info["output_w"]
|
|
layer_info["input_c"] = previous_layer_info["output_c"]
|
|
_changeOPTensorSize(self.layer[i], "input", 0, layer_info["input_h"], layer_info["input_w"])
|
|
if op_code_str == "AVERAGE_POOL_2D":
|
|
layer_info["filter_h"] = layer_info["input_h"]
|
|
layer_info["filter_w"] = layer_info["input_w"]
|
|
layer_info["filter_c"] = layer_info["input_c"]
|
|
|
|
# handle nodes for dag op
|
|
# find the previous node
|
|
if "dagop_input0_key" in layer_info:
|
|
for op in self.layer:
|
|
l_into = op.get_layer_info()
|
|
if (
|
|
"dagop_output_key" in l_into
|
|
and l_into["dagop_output_key"] == layer_info["dagop_input0_key"]
|
|
):
|
|
layer_info["input_h"] = l_into["output_h"]
|
|
layer_info["input_w"] = l_into["output_w"]
|
|
layer_info["input_c"] = l_into["output_c"]
|
|
if "dagop_input1_key" in layer_info:
|
|
for op in self.layer:
|
|
l_into = op.get_layer_info()
|
|
if (
|
|
"dagop_output_key" in l_into
|
|
and l_into["dagop_output_key"] == layer_info["dagop_input1_key"]
|
|
):
|
|
layer_info["input_h"] = l_into["output_h"]
|
|
layer_info["input_w"] = l_into["output_w"]
|
|
layer_info["input_c"] = l_into["output_c"]
|
|
|
|
if op_code_str == "CONV_2D" or op_code_str == "DEPTHWISE_CONV_2D":
|
|
layer_info["output_h"] = math.ceil(layer_info["input_h"] / layer_info["stride_h"])
|
|
layer_info["output_w"] = math.ceil(layer_info["input_w"] / layer_info["stride_w"])
|
|
_changeOPTensorSize(self.layer[i], "output", 0, layer_info["output_h"], layer_info["output_w"])
|
|
elif op_code_str == "ADD":
|
|
layer_info["output_h"] = layer_info["input_h"]
|
|
layer_info["output_w"] = layer_info["input_w"]
|
|
layer_info["output_c"] = layer_info["input_c"]
|
|
_changeOPTensorSize(self.layer[i], "output", 0, layer_info["output_h"], layer_info["output_w"])
|
|
layer_info["input2_h"] = layer_info["input_h"]
|
|
layer_info["input2_w"] = layer_info["input_w"]
|
|
_changeOPTensorSize(self.layer[i], "input", 1, layer_info["input2_h"], layer_info["input_w"])
|
|
elif op_code_str == "SE_ELEMENT_MULT_2D":
|
|
layer_info["input2_h"] = SEinput_h
|
|
layer_info["input2_w"] = SEinput_w
|
|
_changeOPTensorSize(self.layer[i], "input", 1, layer_info["input2_h"], layer_info["input_w"])
|
|
layer_info["output_h"] = SEinput_h
|
|
layer_info["output_w"] = SEinput_w
|
|
_changeOPTensorSize(self.layer[i], "output", 0, layer_info["output_h"], layer_info["output_w"])
|
|
elif op_code_str == "UPSAMPLE":
|
|
layer_info["output_h"] = layer_info["input_h"] * layer_info["factor"]
|
|
layer_info["output_w"] = layer_info["input_w"] * layer_info["factor"]
|
|
layer_info["output_c"] = layer_info["input_c"]
|
|
_changeOPTensorSize(self.layer[i], "output", 0, layer_info["output_h"], layer_info["output_w"])
|
|
elif op_code_str == "MAX_POOL_2D":
|
|
layer_info["output_h"] = int(layer_info["input_h"] / layer_info["filter_h"])
|
|
layer_info["output_w"] = int(layer_info["input_w"] / layer_info["filter_h"])
|
|
layer_info["output_c"] = layer_info["input_c"]
|
|
_changeOPTensorSize(self.layer[i], "output", 0, layer_info["output_h"], layer_info["output_w"])
|
|
|
|
|
|
def _changeOPTensorSize(layer, tensor_type: str, tensor_idx: int, input_h: int, input_w: int):
|
|
if tensor_type == "input":
|
|
if hasattr(layer, "input_tensors") and len(layer.input_tensors) > tensor_idx:
|
|
layer.input_tensors[tensor_idx].set_input_w(input_w)
|
|
layer.input_tensors[tensor_idx].set_input_h(input_h)
|
|
elif tensor_type == "output":
|
|
if hasattr(layer, "output_tensors"):
|
|
layer.output_tensors[tensor_idx].set_input_w(input_w)
|
|
layer.output_tensors[tensor_idx].set_input_h(input_h)
|
|
|
|
|
|
class PatchResizer:
|
|
def __init__(self, layer):
|
|
self.layer = layer
|
|
|
|
# manually setting these variables for now
|
|
def patchResize(self, PatchLayers, PatchSize, PatchSize_height):
|
|
for i, layer in enumerate(self.layer):
|
|
layer_info = layer.get_layer_info()
|
|
if i < PatchLayers:
|
|
layer_info["is_patch"] = True
|
|
op_code_str = layer_info["op"]
|
|
if i == 0:
|
|
layer_info["input_h"] = PatchSize_height
|
|
layer_info["input_w"] = PatchSize
|
|
_changeOPTensorSize(self.layer[i], "input", 0, PatchSize_height, PatchSize)
|
|
else:
|
|
prev_layer_info = self.layer[i - 1].get_layer_info()
|
|
layer_info["input_h"] = prev_layer_info["output_h"]
|
|
layer_info["input_w"] = prev_layer_info["output_w"]
|
|
_changeOPTensorSize(
|
|
self.layer[i], "input", 0, prev_layer_info["output_h"], prev_layer_info["output_w"]
|
|
)
|
|
|
|
if op_code_str == "CONV_2D" or op_code_str == "DEPTHWISE_CONV_2D":
|
|
layer_info["output_h"] = math.ceil(
|
|
(layer_info["input_h"] - layer_info["kernel_h"] + 1) / layer_info["stride_h"]
|
|
)
|
|
layer_info["output_w"] = math.ceil(
|
|
(layer_info["input_w"] - layer_info["kernel_w"] + 1) / layer_info["stride_w"]
|
|
)
|
|
_changeOPTensorSize(self.layer[i], "output", 0, layer_info["output_h"], layer_info["output_w"])
|
|
elif op_code_str == "ADD":
|
|
layer_info["output_h"] = layer_info["input_h"]
|
|
layer_info["output_w"] = layer_info["input_w"]
|
|
layer_info["input2_h"] = layer_info["input_h"]
|
|
layer_info["input2_w"] = layer_info["input_w"]
|
|
_changeOPTensorSize(self.layer[i], "input", 0, layer_info["input_h"], layer_info["input_w"])
|
|
_changeOPTensorSize(self.layer[i], "input", 1, layer_info["input_h"], layer_info["input_w"])
|
|
else:
|
|
layer_info["is_patch"] = False
|