def test_sparsity():
zeros = torch.zeros(2, 3, 5, 6)
print(distiller.sparsity(zeros))
assert distiller.sparsity(zeros) == 1.0
assert distiller.sparsity_3D(zeros) == 1.0
assert distiller.density_3D(zeros) == 0.0
ones = torch.zeros(12, 43, 4, 6)
ones.fill_(1)
assert distiller.sparsity(ones) == 0.0
def get_layer_macs(self, layer):
"""Return the number of MACs required to compute <layer>'s Convolution"""
if layer is None:
return 0
conv_module = distiller.model_find_module(self.model, layer.name)
# MACs = volume(OFM) * (#IFM * K^2)
dense_macs = (conv_module.out_channels * layer.ofm_h * layer.ofm_w) * (conv_module.in_channels * layer.k**2)
if PERFORM_THINNING:
return dense_macs
# If we didn't physically remove structures, we need to use the structural sparsity to compute MACs
conv_pname = layer.name + ".weight"
conv_p = distiller.model_find_param(self.model, conv_pname)
# return dense_macs * distiller.density_ch(conv_p) # Channel pruning
return dense_macs * distiller.density_3D(conv_p) # Filter pruning
def test_sparsity():
zeros = torch.zeros(2, 3, 5, 6)
print(distiller.sparsity(zeros))
assert distiller.sparsity(zeros) == 1.0
assert distiller.sparsity_3D(zeros) == 1.0
assert distiller.density_3D(zeros) == 0.0
ones = torch.ones(12, 43, 4, 6)
assert distiller.sparsity(ones) == 0.0
x = torch.tensor([[1., 2., 0, 4., 0], [1., 2., 0, 4., 0]])
assert distiller.density(x) == 0.6
assert distiller.density_cols(x, transposed=False) == 0.6
assert distiller.sparsity_rows(x, transposed=False) == 0
x = torch.tensor([[0., 0., 0], [1., 4., 0], [1., 2., 0], [0., 0., 0]])
assert distiller.density(x) == 4 / 12
assert distiller.sparsity_rows(x, transposed=False) == 0.5
assert common.almost_equal(distiller.sparsity_cols(x, transposed=False),
1 / 3)
assert common.almost_equal(distiller.sparsity_rows(x), 1 / 3)
def collect_conv_details(model,
dataset,
perform_thinning,
layers_to_prune=None):
dummy_input = distiller.get_dummy_input(dataset)
g = SummaryGraph(model, dummy_input)
conv_layers = OrderedDict()
total_macs = 0
total_params = 0
for id, (name, m) in enumerate(model.named_modules()):
if isinstance(m, torch.nn.Conv2d):
conv = SimpleNamespace()
conv.t = len(conv_layers)
conv.k = m.kernel_size[0]
conv.stride = m.stride
# Use the SummaryGraph to obtain some other details of the models
conv_op = g.find_op(normalize_module_name(name))
assert conv_op is not None
conv.weights_vol = conv_op['attrs']['weights_vol']
total_params += conv.weights_vol
conv.macs = conv_op['attrs']['MACs']
conv_pname = name + ".weight"
conv_p = distiller.model_find_param(model, conv_pname)
if not perform_thinning:
#conv.macs *= distiller.density_ch(conv_p) # Channel pruning
conv.macs *= distiller.density_3D(conv_p) # Filter pruning
total_macs += conv.macs
conv.ofm_h = g.param_shape(conv_op['outputs'][0])[2]
conv.ofm_w = g.param_shape(conv_op['outputs'][0])[3]
conv.ifm_h = g.param_shape(conv_op['inputs'][0])[2]
conv.ifm_w = g.param_shape(conv_op['inputs'][0])[3]
conv.name = name
conv.id = id
if layers_to_prune is None or name in layers_to_prune:
conv_layers[len(conv_layers)] = conv
return conv_layers, total_macs, total_params
