def set_mask(self, mask):
if mask is None:
return
channel = mask.sum().item()
mid_channel = make_divisible(channel // self.reduction, 8)
exp_mask = _get_channel_mask(self.se.expand.weight.data, mid_channel)
self.se.reduction.set_mask(mask, exp_mask)
self.se.expand.set_mask(exp_mask, mask)
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, self.C * expansion)
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)