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