def generate_profiling_primitives(
self,
base_channels,
mult_ratio,
strides,
acts=None,
use_ses=None,
primitive_type="mobilenet_v2_block",
spatial_size=224,
stem_stride=2,
stem_type="conv_3x3",
sample=None,
as_dict=True,
):
channels = [make_divisible(c * mult_ratio, 8) for c in base_channels]
primitives = []
acts = acts or [
None,
] * len(strides)
use_ses = use_ses or [
None,
] * len(strides)
sizes = [round(spatial_size / stem_stride + 1e-3)]
for s in strides:
sizes += [round(sizes[-1] / s + 1e-3)]
stem_prim = Prim(stem_type, spatial_size, 3, channels[0], stem_stride,
True)
first_fixed_prim = Prim(primitive_type,
spatial_size / stem_stride,
channels[0],
channels[1],
strides[0],
True,
kernel_size=3,
activation=acts[0],
use_se=use_ses[0],
expansion=1)
primitives += [stem_prim, first_fixed_prim]
for w, k, stage, i in self._traverse_search_space(sample):
primitives.append(
Prim(
primitive_type,
sizes[stage + i],
channels[stage + i],
channels[stage + 1],
1 if i else strides[stage],
True,
kernel_size=k,
activation=acts[stage],
use_se=use_ses[stage],
expansion=w,
))
if self.fixed_primitives is not None:
primitives += self.fixed_primitives
primitives = list(set(primitives))
if as_dict:
primitives = [dict(p._asdict()) for p in primitives]
return primitives
def __init__(
self,
expansion,
C,
C_out,
stride,
kernel_sizes=(3, 5, 7),
do_kernel_transform=True,
affine=True,
activation="relu",
schedule_cfg=None,
):
FlexibleBlock.__init__(self, schedule_cfg)
self.activation = activation
C_inner = make_divisible(C * expansion, 8)
self.kernel_sizes = sorted(kernel_sizes)
self.kernel_size = self.kernel_sizes[-1]
self.do_kernel_transform = do_kernel_transform
self.affine = affine
inv_bottleneck = None
if expansion != 1:
inv_bottleneck = nn.Sequential(
FlexiblePointLinear(C, C_inner, 1, 1, 0),
FlexibleBatchNorm2d(C_inner, affine=affine),
get_op(activation)(),
)
depth_wise = nn.Sequential(
FlexibleDepthWiseConv(
C_inner,
self.kernel_sizes,
stride,
do_kernel_transform=do_kernel_transform,
),
FlexibleBatchNorm2d(C_inner, affine=affine),
get_op(activation)(),
)
point_linear = nn.Sequential(
FlexiblePointLinear(C_inner, C_out, 1, 1, 0),
FlexibleBatchNorm2d(C_out, affine=affine),
)
super(FlexibleMobileNetV2Block, self).__init__(
expansion,
C,
C_out,
stride,
self.kernel_size,
affine,
activation,
inv_bottleneck,
depth_wise,
point_linear,
)
self.reset_mask()
def __init__(
self,
expansion,
C,
C_out,
stride,
kernel_size,
affine,
activation="relu",
inv_bottleneck=None,
depth_wise=None,
point_linear=None,
):
super(MobileNetV2Block, self).__init__()
self.expansion = expansion
self.C = C
self.C_out = C_out
self.C_inner = make_divisible(C * expansion, 8)
self.stride = stride
self.kernel_size = kernel_size
self.act_fn = get_op(activation)()
self.inv_bottleneck = None
if expansion != 1:
self.inv_bottleneck = inv_bottleneck or nn.Sequential(
nn.Conv2d(C, self.C_inner, 1, 1, 0, bias=False),
nn.BatchNorm2d(self.C_inner), self.act_fn)
self.depth_wise = depth_wise or nn.Sequential(
nn.Conv2d(self.C_inner,
self.C_inner,
self.kernel_size,
stride,
padding=self.kernel_size // 2,
bias=False), nn.BatchNorm2d(self.C_inner), self.act_fn)
self.point_linear = point_linear or nn.Sequential(
nn.Conv2d(self.C_inner, C_out, 1, 1, 0, bias=False),
nn.BatchNorm2d(C_out))
self.shortcut = nn.Sequential()
self.has_conv_shortcut = False
if stride == 1 and C != C_out:
self.has_conv_shortcut = True
self.shortcut = nn.Sequential(
nn.Conv2d(C,
C_out,
kernel_size=1,
stride=1,
padding=0,
bias=False),
nn.BatchNorm2d(C_out),
)
def rollout_to_primitives(cls,
rollout,
primitive_type,
spatial_size,
strides,
base_channels,
mult_ratio=1.,
stem_type="conv_3x3",
stem_stride=2,
**kwargs):
acts = kwargs.get("acts", [None] * len(rollout.depth))
use_ses = kwargs.get("use_ses", [None] * len(rollout.depth))
stem_stride = kwargs.get("stem_stride", 2)
sizes = [
round(spatial_size / stem_stride /
reduce(lambda p, q: p * q, strides[:i]))
for i in range(1,
len(strides) + 1)
]
channels = [make_divisible(mult_ratio * c, 8) for c in base_channels]
primitives = [
Prim(stem_type, spatial_size, 3, channels[0], stem_stride, True)
]
for i, (depth, size, s, c_in, c_out, act, se) in enumerate(
zip(
rollout.depth,
sizes,
strides,
channels[:-1],
channels[1:],
acts,
use_ses,
)):
for j, width, kernel in zip(range(depth), rollout.width[i],
rollout.kernel[i]):
if j > 0:
c_in = c_out
s = 1
primitives.append(
Prim(
primitive_type,
size,
c_in,
c_out,
s,
True,
kernel_size=kernel,
activation=act,
use_se=se,
expansion=width,
))
return primitives
def set_mask(self, expansion, kernel_size):
mask = None
if expansion is not None and expansion != self.expansion:
filters = self.point_linear[0].weight.data
mask = _get_channel_mask(filters,
make_divisible(self.C * expansion, 8))
if self.inv_bottleneck:
self.inv_bottleneck[0].set_mask(None, mask)
self.inv_bottleneck[1].set_mask(mask)
self.depth_wise[0].set_mask(mask, kernel_size)
self.depth_wise[1].set_mask(mask)
self.point_linear[0].set_mask(mask, None)
def __init__(self,
expansion,
C,
C_out,
stride,
affine,
kernel_size=3,
relu6=False):
super(MobileNetBlock, self).__init__()
C_inner = self.C_inner = make_divisible(expansion * C, 8)
self.stride = stride
self.relu6 = relu6
self.activation = F.relu6 if self.relu6 else F.relu
self.conv1 = nn.Conv2d(C,
C_inner,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.bn1 = nn.BatchNorm2d(C_inner, affine=affine)
padding = int((kernel_size - 1) / 2)
self.conv2 = nn.Conv2d(C_inner,
C_inner,
kernel_size=kernel_size,
stride=stride,
padding=padding,
groups=C_inner,
bias=False)
self.bn2 = nn.BatchNorm2d(C_inner, affine=affine)
self.conv3 = nn.Conv2d(C_inner,
C_out,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.bn3 = nn.BatchNorm2d(C_out, affine=affine)
self.shortcut = nn.Sequential()
self.has_conv_shortcut = False
if stride == 1 and C != C_out:
self.has_conv_shortcut = True
self.shortcut = nn.Sequential(
nn.Conv2d(C,
C_out,
kernel_size=1,
stride=1,
padding=0,
bias=False),
nn.BatchNorm2d(C_out, affine=affine),
)
def __init__(
self,
device,
blocks=[1, 4, 4, 4, 4, 4],
strides=[1, 2, 2, 2, 1, 2],
expansions=[1, 6, 6, 6, 6, 6],
layer_channels=[16, 16, 24, 40, 80, 112, 160, 960, 1280],
mult_ratio=1.0,
kernel_sizes=[3, 5, 7],
do_kernel_transform=True,
use_ses=[False, False, True, False, True, True],
acts=["relu", "relu", "relu", "h_swish", "h_swish", "h_swish"],
num_classes=10,
block_type="mbv3_block",
pretrained_path=None,
stem_stride=2,
schedule_cfg=None,
):
super(MobileNetV3Arch, self).__init__(
device,
blocks,
strides,
expansions,
layer_channels,
mult_ratio,
kernel_sizes,
do_kernel_transform,
num_classes,
block_type,
pretrained_path,
schedule_cfg,
)
self.block_initializer = FlexibleBlock.get_class_(block_type)
self.channels = [
make_divisible(c * mult_ratio, 8) for c in layer_channels
]
self.stem_stride = stem_stride
self.stem = nn.Sequential(
nn.Conv2d(3,
self.channels[0],
kernel_size=3,
stride=self.stem_stride,
padding=1,
bias=False),
nn.BatchNorm2d(self.channels[0]),
get_op("h_swish")(),
)
expect(
blocks[0] == expansions[0] == 1,
"The first conv layer should have single block and no expansion.",
ValueError,
)
self.mult_ratio = mult_ratio
self.use_ses = use_ses
self.acts = acts
self.cells = [
self.make_stage(
self.channels[0],
self.channels[1],
self.blocks[0],
self.strides[0],
self.expansions[0],
[3],
self.use_ses[0],
self.acts[0],
)
]
for i, depth in enumerate(self.blocks[1:], 1):
self.cells.append(
self.make_stage(
self.channels[i],
self.channels[i + 1],
depth,
self.strides[i],
self.expansions[i],
self.kernel_sizes,
self.use_ses[i],
self.acts[i],
))
self.cells = nn.ModuleList(self.cells)
self.conv_head = nn.Sequential(
nn.Conv2d(self.channels[-3],
self.channels[-2],
1,
1,
0,
bias=False),
nn.BatchNorm2d(self.channels[-2]),
get_op("h_swish")(),
)
self.conv_final = nn.Sequential(
nn.Conv2d(self.channels[-2],
self.channels[-1],
1,
1,
0,
bias=False),
get_op("h_swish")(),
)
self.classifier = nn.Linear(self.channels[-1], num_classes)
if self.pretrained_path:
state_dict = torch.load(self.pretrained_path, "cpu")
if state_dict["classifier.weight"].shape[0] != self.num_classes:
del state_dict["classifier.weight"]
del state_dict["classifier.bias"]
self.logger.info(self.load_state_dict(state_dict, strict=False))
self.to(self.device)
def rollout_to_primitives(
cls,
rollout,
primitive_type,
spatial_size,
strides,
base_channels,
mult_ratio=1.,
stem_type="conv_3x3",
# stem_stride=2,
**kwargs):
acts = kwargs.get("acts", [None] * len(rollout.depth))
use_ses = kwargs.get("use_ses", [None] * len(rollout.depth))
stem_stride = kwargs.get("stem_stride", 2)
sizes = [round(spatial_size / stem_stride + 1e-3)]
for s in strides:
sizes += [round(sizes[-1] / s + 1e-3)]
in_sizes = [spatial_size] + sizes
# import ipdb; ipdb.set_trace()
channels = [make_divisible(mult_ratio * c, 8) for c in base_channels]
primitives = [
Prim(stem_type, spatial_size, 3, channels[0], stem_stride, True)
]
for i, (depth, in_size, size, s, c_in, c_out, act, se) in enumerate(
zip(
rollout.depth,
in_sizes[1:],
sizes[1:],
strides,
channels[:-1],
channels[1:],
acts,
use_ses,
)):
# `size` indicates the output size of this block
# if the stride equals 2, the input size of this block should be `size * s`
for j, width, kernel in zip(range(depth), rollout.width[i],
rollout.kernel[i]):
# if j > 0:
# c_in = c_out
# s = 1
# print(">>>", size, s, size * s)
primitives.append(
Prim(
primitive_type,
in_size if j == 0 else size,
c_in if j == 0 else c_out,
c_out,
s if j == 0 else 1,
True,
kernel_size=kernel,
activation=act,
use_se=se,
expansion=width,
))
return primitives
