当前位置：首页 > news >正文

lottery-ticket-hypothesis完全指南：从MNIST数据集开始的神经网络剪枝实验

news 2026/7/21 14:24:03

lottery-ticket-hypothesis完全指南：从MNIST数据集开始的神经网络剪枝实验

【免费下载链接】lottery-ticket-hypothesisA reimplementation of "The Lottery Ticket Hypothesis" (Frankle and Carbin) on MNIST.项目地址: https://gitcode.com/gh_mirrors/lo/lottery-ticket-hypothesis

什么是彩票假说（Lottery Ticket Hypothesis）？

彩票假说（Lottery Ticket Hypothesis）是深度学习领域的一项重要发现，它揭示了神经网络中存在"中奖彩票"——即一个小型子网，当使用原始网络的初始权重进行训练时，其性能可以与完整网络相媲美。这个发现为神经网络剪枝提供了全新思路，让我们能够在保持性能的同时大幅减小模型大小。

为什么选择MNIST数据集进行实验？

MNIST数据集是机器学习领域最经典的手写数字识别数据集，包含60,000个训练样本和10,000个测试样本。选择MNIST进行彩票假说实验有以下优势：

简单直观：28x28像素的灰度图像，适合入门级实验
训练快速：普通计算机即可在短时间内完成训练
可复现性高：结果稳定，便于验证剪枝效果

在本项目中，MNIST数据集的相关配置和处理集中在以下文件：

mnist_fc/constants.py：MNIST实验的超参数设置
mnist_fc/download_data.py：下载并转换MNIST数据集
datasets/dataset_mnist.py：MNIST数据集加载和预处理

神经网络剪枝的核心概念

剪枝掩码（Pruning Masks）

剪枝的核心是创建"掩码"（masks）——一个与网络权重形状相同的二进制数组，其中1表示保留该权重，0表示剪枝该权重。在项目中，掩码的创建和管理主要通过以下文件实现：

# 掩码的保存路径定义 def masks(parent_directory): """The path where the pruning masks are stored.""" return os.path.join(parent_directory, 'masks')

foundations/paths.py中的掩码路径定义

剪枝算法

项目实现了基于权重大小的剪枝方法，通过保留权重绝对值较大的连接来构建子网：

def prune_by_percent(percents, masks, final_weights): """Return new masks that involve pruning the smallest of the final weights.""" # 实现根据百分比剪枝最小权重的逻辑

foundations/pruning.py中的剪枝算法

权重重新初始化

彩票假说的关键步骤之一是将剪枝后的子网权重重新初始化为原始网络的初始值，以验证其"中奖"特性：

# 权重重新初始化逻辑 for k, mask in masks.items(): # 保留原始初始化分布的同时应用掩码 positive = np.random.choice(init[init > 0], mask.shape) negative = np.random.choice(init[init < 0], mask.shape) presets[k] = np.where(mask, positive if positive.any() else negative, 0)

mnist_fc/reinitialize.py中的权重重新初始化

实验步骤：从零开始的彩票假说验证

1. 环境准备

首先克隆项目仓库到本地：

git clone https://gitcode.com/gh_mirrors/lo/lottery-ticket-hypothesis

2. 数据集准备

修改MNIST数据存储位置配置：

# 修改mnist_fc/locations.py文件 MNIST_LOCATION = '/path/to/your/mnist/data' # 设置数据存储路径

然后运行数据下载脚本：

python mnist_fc/download_data.py

3. 基础网络训练

运行完整网络训练脚本，获取初始权重：

python mnist_fc/train.py

训练过程中，模型会自动记录损失（loss）和准确率（accuracy）：

# 性能指标记录 for loss, it, acc in zip(data['loss'], data['iteration'], data['accuracy']): writer.write(f"{it},{loss},{acc}\n")

foundations/save_restore.py中的性能记录

4. 网络剪枝实验

执行彩票假说实验主程序：

python mnist_fc/lottery_experiment.py

该实验会自动执行多轮剪枝，每轮保留一定比例的权重，核心逻辑如下：

prune_masks = functools.partial(pruning.prune_by_percent, percents=constants.PRUNE_PERCENTS) # 执行剪枝并评估性能

mnist_fc/lottery_experiment.py中的剪枝流程

5. 重新初始化验证

为验证剪枝后的子网是否为"中奖彩票"，运行重新初始化实验：

python mnist_fc/reinitialize.py

该实验使用原始初始权重重新训练剪枝后的子网，验证其是否能达到与完整网络相当的性能。

关键代码解析

模型定义与掩码应用

项目中的基础模型类实现了掩码的应用逻辑：

def dense_layer(self, name, input_layer, units, activation=tf.nn.relu): """Mimics tf.dense_layer but masks weights and uses presets as necessary.""" if name in self._masks: mask_initializer = tf.constant_initializer(self._masks[name]) mask = tf.get_variable( name + '_mask', initializer=mask_initializer, trainable=False) weights = tf.multiply(weights, mask) # 应用掩码

foundations/model_base.py中的掩码应用

掩码的合并与操作

项目提供了掩码的并集（union）和交集（intersect）操作，用于组合不同剪枝策略的结果：

def union(*masks): """Return new masks that are the per-layer union of the provided masks.""" # 实现掩码的并集操作 def intersect(*masks): """Return new masks that are the per-layer intersection of the provided masks.""" # 实现掩码的交集操作

foundations/union.py中的掩码操作