2de52a8976971da4836727ba9242fedcc7474878,src/sdk/pynni/nni/compression/torch/compressor.py,Pruner,export_model,#Pruner#Any#Any#Any#Any#,213

Before Change

        if self.detect_modules_to_compress() and not self.mask_dict:
            _logger.warning("You may not use self.mask_dict in base Pruner class to record masks")
        assert model_path is not None, "model_path must be specified"
        for name, m in self.bound_model.named_modules():
            if name == "":
                continue
            masks = self.mask_dict.get(name)
            if masks is not None:
                mask_sum = masks["weight"].sum().item()
                mask_num = masks["weight"].numel()
                _logger.info("Layer: %s  Sparsity: %.2f", name, 1 - mask_sum / mask_num)
                m.weight.data = m.weight.data.mul(masks["weight"])
                if masks.__contains__("bias") and hasattr(m, "bias") and m.bias is not None:
                    m.bias.data = m.bias.data.mul(masks["bias"])
            else:
                _logger.info("Layer: %s  NOT compressed", name)
        torch.save(self.bound_model.state_dict(), model_path)
        _logger.info("Model state_dict saved to %s", model_path)
        if mask_path is not None:
            torch.save(self.mask_dict, mask_path)

After Change

            if bias_mask is not None:
                wrapper.module.bias.data = wrapper.module.bias.data.mul(bias_mask)
            // save mask to dict
            mask_dict[wrapper.name] = {"weight": weight_mask, "bias": bias_mask}

        torch.save(self.bound_model.state_dict(), model_path)
        _logger.info("Model state_dict saved to %s", model_path)
        if mask_path is not None:

Instances