def get_input_oper(inp_block_def: str):
parts = inp_block_def.split('#')
if len(parts) != 2:
raise Exception('Invalid definition of stack with inputs.')
input_name, oper_name = parts
input_name = input_name + "#"
if len(oper_name) > 0:
input_oper_class = NetBuilderConfig.get_input_operation(oper_name)
else:
input_oper_class = NoOpInput
oper = input_oper_class()
return input_name, oper
def decode_block_args(args: list):
"""
Parses block definition. For example the representation 'k3_s2' is decoded to {'kernel_size': 3, 'stride': 2}
"""
args_defs = {}
repeats = 1
for arg in args:
parser = NetBuilderConfig.find_parser(arg)
arg_inner = arg[len(parser.code):]
if parser.arg_name == KEY_REPEATS_NUM:
repeats = parser.func(arg_inner)
else:
args_defs[parser.arg_name] = parser.func(arg_inner)
return args_defs, repeats
def make_block(arch_args, in_chs):
"""
Creates a block instance
"""
block_type_code = arch_args['block_type']
out_chs = arch_args.get('out_chs')
if out_chs is None:
out_chs = in_chs
arch_args['name'] = clean_name(arch_args['name'])
block_type_def = NetBuilderConfig.get_block_type(block_type_code)
args = block_type_def.transform_args_fn(in_chs, arch_args)
block = block_type_def.block_type_fn(**args)
return block, out_chs
def create_execution_item(self, block_args: dict, in_chs: list):
name = block_args.get('name')
block_type = block_args.get('block_type')
# operation
if block_type == BLOCK_TYPE_OPER:
oper_name = block_args['oper_name']
oper_args = block_args['oper_args']
oper_class = NetBuilderConfig.get_operation(oper_name)
oper = oper_class(oper_args)
input_tensor_chs = []
for tn in oper.get_input_tenors_names():
ei = self.find_exec_item(tn)
input_tensor_chs.append(ei.out_chs)
if len(input_tensor_chs) > 0:
out_chs = oper.get_out_chs_number(input_tensor_chs)
else:
out_chs = oper.get_out_chs_number(in_chs)
ei = ExecutionItem(name, oper, out_chs, is_oper=True)
return ei
# layer
else:
assert len(in_chs) == 1
block, out_chs = make_block(block_args, in_chs[0])
ei = ExecutionItem(name, block, out_chs, is_oper=False)
return ei
import tensorflow as tf
from tf_netbuilder import NetBuilderConfig
from tf_netbuilder.parser import KEY_REPEATS_NUM
from tf_netbuilder.activations import ACTIVATION_FUNCS
from tf_netbuilder.layers import InvertedResidual, ConvBnAct, SmallerConv7x7
from tf_netbuilder.layers.prep_funcs import prepare_ir_args, prepare_cn_args, prepare_cn2_args, prepare_hd_args, \
prepare_avgpool_args, prepare_maxpool_args
from tf_netbuilder.operations import ConcatOper, MulOper, UpscaleX2, \
NormalizationMinus05Plus05Input, NormalizationMinus1Plus1Input
# registering parsers for arguments of all blocks
NetBuilderConfig.add_parser("n", lambda arg: ACTIVATION_FUNCS[arg], "activation")
NetBuilderConfig.add_parser("s", lambda arg: int(arg), "strides")
NetBuilderConfig.add_parser("c", lambda arg: int(arg), "out_chs")
NetBuilderConfig.add_parser("e", lambda arg: float(eval(arg)), "exp_ratio")
NetBuilderConfig.add_parser("se", lambda arg: int(arg), "se_factor")
NetBuilderConfig.add_parser("k", lambda arg: int(arg), "kernel_size")
NetBuilderConfig.add_parser("bn", lambda arg: True, "use_batch_norm") # no arg, just the flag
# below is the internal parser and variable determining how many times a given layer should be repeated
NetBuilderConfig.add_parser("r", lambda arg: int(arg), KEY_REPEATS_NUM)
# registering blocks types - layers
NetBuilderConfig.add_block_type("ir", InvertedResidual, prepare_ir_args)
NetBuilderConfig.add_block_type("cn", ConvBnAct, prepare_cn_args)
NetBuilderConfig.add_block_type("cn2", SmallerConv7x7, prepare_cn2_args)
NetBuilderConfig.add_block_type("hd", tf.keras.layers.Conv2D, prepare_hd_args)
NetBuilderConfig.add_block_type("avgpool", tf.keras.layers.AveragePooling2D, prepare_avgpool_args)
def register_tf_netbuilder_extensions():
NetBuilderConfig.add_block_type("cnreg", tf.keras.layers.Conv2D,
prepare_cnreg),
NetBuilderConfig.add_block_type("cn3", tf.keras.layers.Conv2D,
prepare_cn3_args),