Update 2.pruning_criteria.ipynb

Author: ironartisan

docs/notebook/ch03/2.pruning_criteria.ipynb

@@ -389,11 +389,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 49,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "# 定义一个LeNet网络\n",
+    "# 为避免前面的操作影响后续结果，重新定义一个LeNet网络，和前面一致\n",
     "class LeNet(nn.Module):\n",
     "    def __init__(self, num_classes=10):\n",
     "        super(LeNet, self).__init__()\n",
@@ -451,13 +451,20 @@
     "model.load_state_dict(checkpoint)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "基于梯度幅度的修剪标准"
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 52,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "# 基于梯度幅度的修剪标准\n",
+    "# 修剪整个模型的权重，传入整个模型\n",
     "def gradient_magnitude_pruning(model, percentile):\n",
     "    for name, param in model.named_parameters():\n",
     "        if 'weight' in name:\n",
@@ -467,10 +474,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 53,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
+    "# 修剪局部模型权重，传入某一层的权重\n",
     "@torch.no_grad()\n",
     "def gradient_magnitude_pruning(weight, gradient, percentile=0.5):\n",
     "    num_elements = weight.numel()\n",
@@ -491,7 +499,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 54,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -514,7 +522,7 @@
     }
    ],
    "source": [
-    "# 使用示例\n",
+    "# 使用示例，这里以fc2层的权重为例\n",
     "percentile = 0.5\n",
     "gradient_magnitude_pruning(model.fc2.weight, gradients['fc2.weight'], percentile)"
    ]