def test_unify_all_predefinied_archs(self):
""" Initial check for arch definitions """
arch_factory = fbnet_modeldef_cls.MODEL_ARCH
self.assertGreater(len(arch_factory), 0)
for name, arch in arch_factory.items():
with self.subTest(arch=name):
print("Unifiying {}".format(name))
fbnet_builder.unify_arch_def(arch, ["blocks"])
def _create_builder(arch_name_or_def: typing.Union[str, dict]):
if isinstance(arch_name_or_def, str):
assert arch_name_or_def in modeldef.MODEL_ARCH, (
f"Invalid arch name {arch_name_or_def}, "
f"available names: {modeldef.MODEL_ARCH.keys()}")
arch_def = modeldef.MODEL_ARCH[arch_name_or_def]
else:
assert isinstance(arch_name_or_def, dict)
arch_def = arch_name_or_def
arch_def = mbuilder.unify_arch_def(arch_def, ["blocks"])
scale_factor = 1.0
width_divisor = 1
bn_info = {"name": "bn", "momentum": 0.003}
drop_out = 0.2
arch_def["dropout_ratio"] = drop_out
builder = mbuilder.FBNetBuilder(width_ratio=scale_factor,
bn_args=bn_info,
width_divisor=width_divisor)
builder.add_basic_args(**arch_def.get("basic_args", {}))
return builder, arch_def
def _build_model(arch_def, dim_in):
arch_def = fbnet_builder.unify_arch_def(arch_def, ["blocks"])
torch.manual_seed(0)
builder = fbnet_builder.FBNetBuilder(1.0)
model = builder.build_blocks(arch_def["blocks"], dim_in=dim_in)
model.eval()
return model
def test_count_strides(self):
e6 = {"expansion": 6}
arch_def = {
"blocks1": [
# [op, c, s, n, ...]
# stage 0
[("conv_k3", 32, 2, 1)],
# stage 1
[("ir_k3", 64, 2, 2, e6), ("ir_k5", 96, 1, 1, e6)],
],
"blocks2": [
# [op, c, s, n, ...]
# stage 0
[("conv_k3", 32, 2, 1)],
# stage 1
[("ir_k3", 64, -2, 2, e6), ("ir_k5", 96, -2, 1, e6)],
],
}
unified_arch = mbuilder.unify_arch_def(arch_def,
["blocks1", "blocks2"])
gt_strides_blocks1 = [2, 2, 1, 1]
gt_strides_blocks2 = [2, 0.5, 1, 0.5]
count_strides1 = mbuilder.count_stride_each_block(
unified_arch["blocks1"])
count_strides2 = mbuilder.count_stride_each_block(
unified_arch["blocks2"])
self.assertEqual(gt_strides_blocks1, count_strides1)
self.assertEqual(gt_strides_blocks2, count_strides2)
all_strides1 = mbuilder.count_strides(unified_arch["blocks1"])
all_strides2 = mbuilder.count_strides(unified_arch["blocks2"])
self.assertEqual(all_strides1, 4)
self.assertEqual(all_strides2, 0.5)
def test_count_stages(self):
e6 = {"expansion": 6}
arch_def = {
"blocks1": [
# [op, c, s, n, ...]
# stage 0
[("conv_k3", 32, 2, 1)],
# stage 1
[("ir_k3", 64, 2, 2, e6), ("ir_k5", 96, 1, 1, e6)],
# stage 2
[("ir_k3", 64, -2, 2, e6), ("ir_k5", 96, -2, 1, e6)],
]
}
arch_def = mbuilder.unify_arch_def(arch_def, ["blocks1"])
num_stages = mbuilder.get_num_stages(arch_def["blocks1"])
self.assertEqual(num_stages, 3)
gt_stage_out_channels = [32, 96, 96]
stage_out_channels = mbuilder.get_stages_dim_out(arch_def["blocks1"])
self.assertEqual(stage_out_channels, gt_stage_out_channels)
gt_blocks_in_stage = [1, 3, 3]
blocks_in_stage = mbuilder.get_num_blocks_in_stage(arch_def["blocks1"])
self.assertEqual(blocks_in_stage, gt_blocks_in_stage)
def _build_model(arch_def, dim_in):
arch_def = fbnet_builder.unify_arch_def(arch_def, ["blocks"])
torch.manual_seed(0)
builder = fbnet_builder.FBNetBuilder(1.0)
builder.add_basic_args(**arch_def.get("basic_args", {}))
model = builder.build_blocks(arch_def["blocks"], dim_in=dim_in)
model.eval()
return model, builder
def _create_and_run(self, arch_name, model_arch):
arch = fbnet_builder.unify_arch_def(model_arch, ["blocks"])
builder = fbnet_builder.FBNetBuilder()
model = builder.build_blocks(arch["blocks"], dim_in=3)
model.eval()
res = model_arch.get("input_size", 224)
inputs = torch.zeros([1, 3, res, res])
output = flops_utils.print_model_flops(model, inputs)
self.assertEqual(output.shape[0], 1)
def build_fbnet(cfg, name, in_channels):
"""
Create a FBNet module using FBNet V2 builder.
Args:
cfg (CfgNode): the config that contains MODEL.FBNET_V2.
name (str): the key in arch_def that represents a subpart of network
in_channels (int): input channel size
Returns:
nn.Sequential: the first return is a nn.Sequential, each element
corresponds a stage in arch_def.
List[ShapeSpec]: the second return is a list of ShapeSpec containing the
output channels and accumulated strides for that stage.
"""
builder, raw_arch_def = _get_fbnet_builder_and_arch_def(cfg)
# Reset the last_depth for this builder (might have been cached), this is
# the only mutable member variable.
builder.last_depth = in_channels
# NOTE: Each sub part of the model consists of several stages and each stage
# has several blocks. "Raw" arch_def (Dict[str, List[List[Tuple]]]) uses a
# list of stages to describe the architecture, which is more compact and
# thus written as builtin metadata (inside FBNetV2ModelArch) or config
# (MODEL.FBNET_V2.ARCH_DEF). "Unified" arch_def (Dict[str, List[Dict]])
# uses a list blocks from all stages instead, which is recognized by builder.
arch_def = mbuilder.unify_arch_def(raw_arch_def, [name])
arch_def = {name: arch_def[name]}
logger.info(
"Build FBNet using unified arch_def:\n{}".format(
format_dict_expanding_list_values(arch_def)
)
)
arch_def_blocks = arch_def[name]
stages = []
trunk_stride_per_stage = _get_stride_per_stage(arch_def_blocks)
shape_spec_per_stage = []
for i, stride_i in enumerate(trunk_stride_per_stage):
stages.append(
builder.build_blocks(
arch_def_blocks,
stage_indices=[i],
prefix_name=FBNET_BUILDER_IDENTIFIER + "_",
)
)
shape_spec_per_stage.append(
ShapeSpec(
channels=builder.last_depth,
stride=stride_i,
)
)
return FBNetModule(*stages), shape_spec_per_stage
def _create_builder(arch_name_or_def: typing.Union[str, dict]):
if isinstance(arch_name_or_def,
str) and arch_name_or_def in modeldef.MODEL_ARCH:
arch_def = modeldef.MODEL_ARCH[arch_name_or_def]
elif isinstance(arch_name_or_def, str):
try:
arch_def = json.loads(arch_name_or_def)
assert isinstance(arch_def,
dict), f"Invalid arch type {arch_name_or_def}"
except ValueError:
assert arch_name_or_def in modeldef.MODEL_ARCH, (
f"Invalid arch name {arch_name_or_def}, "
f"available names: {modeldef.MODEL_ARCH.keys()}")
else:
assert isinstance(arch_name_or_def, dict)
arch_def = arch_name_or_def
arch_def = mbuilder.unify_arch_def(arch_def, ["blocks"])
builder = mbuilder.FBNetBuilder()
builder.add_basic_args(**arch_def.get("basic_args", {}))
return builder, arch_def
def test_unify_arch(self):
e6 = {"expansion": 6}
dw_skip_bnrelu = {"dw_skip_bnrelu": True}
arch_def = {
"blocks": [
# [op, c, s, n, ...]
# stage 0
[("conv_k3", 32, 2, 1)],
# stage 1
[
("ir_k3", 64, 2, 2, e6, dw_skip_bnrelu),
("ir_k5", 96, 1, 1, e6),
],
],
"backbone": [0],
"heads": [1],
}
gt_unified_arch = {
"blocks": [
{
"stage_idx": 0,
"block_idx": 0,
"block_op": "conv_k3",
"block_cfg": {
"out_channels": 32,
"stride": 2
},
},
{
"stage_idx": 1,
"block_idx": 0,
"block_op": "ir_k3",
"block_cfg": {
"out_channels": 64,
"stride": 2,
"expansion": 6,
"dw_skip_bnrelu": True,
},
},
{
"stage_idx": 1,
"block_idx": 1,
"block_op": "ir_k3",
"block_cfg": {
"out_channels": 64,
"stride": 1,
"expansion": 6,
"dw_skip_bnrelu": True,
},
},
{
"stage_idx": 1,
"block_idx": 2,
"block_op": "ir_k5",
"block_cfg": {
"out_channels": 96,
"stride": 1,
"expansion": 6,
},
},
],
"backbone": [0],
"heads": [1],
}
unified_arch = mbuilder.unify_arch_def(arch_def, ["blocks"])
self.assertEqual(unified_arch, gt_unified_arch)
