Increase batch size

pomonam · Mar 20, 2024 · 11eb42a · 11eb42a
1 parent a378089
commit 11eb42a
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 13 deletions.
diff --git a/examples/cifar/analyze.py b/examples/cifar/analyze.py
@@ -125,6 +125,7 @@ def main():
         per_device_batch_size=None,
         factor_args=factor_args,
         overwrite_output_dir=True,
+        initial_per_device_batch_size_attempt=8192,
     )
     analyzer.compute_pairwise_scores(
         scores_name="pairwise",
@@ -134,6 +135,7 @@ def main():
         train_dataset=train_dataset,
         per_device_query_batch_size=500,
         overwrite_output_dir=True,
+        initial_per_device_train_batch_size_attempt=8192,
     )
     scores = analyzer.load_pairwise_scores("pairwise")
     print(scores)

diff --git a/examples/uci/tutorial.ipynb b/examples/uci/tutorial.ipynb
@@ -1178,29 +1178,28 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 36,
+   "execution_count": 38,
    "id": "1a22d21e-3d87-4a88-9665-517487f71ef6",
    "metadata": {},
    "outputs": [
     {
-     "ename": "NameError",
-     "evalue": "name 'summed_scores' is not defined",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[36], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28mlen\u001b[39m(get_topk_keep_indices(\u001b[43msummed_scores\u001b[49m, topk\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m10\u001b[39m)), \u001b[38;5;28mlen\u001b[39m(get_topk_keep_indices(summed_scores, topk\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m20\u001b[39m))\n",
-      "\u001b[0;31mNameError\u001b[0m: name 'summed_scores' is not defined"
-     ]
+     "data": {
+      "text/plain": [
+       "(917, 907)"
+      ]
+     },
+     "execution_count": 38,
+     "metadata": {},
+     "output_type": "execute_result"
     }
    ],
    "source": [
-    "len(get_topk_keep_indices(summed_scores, topk=10)), len(get_topk_keep_indices(summed_scores, topk=20))"
+    "len(get_topk_keep_indices(scores[\"all_modules\"].sum(dim=0), topk=10)), len(get_topk_keep_indices(scores[\"all_modules\"].sum(dim=0), topk=20))"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 39,
    "id": "b90e2a6a-e6bc-4865-a0f8-4e8d6d8012c6",
    "metadata": {},
    "outputs": [],
@@ -1209,7 +1208,7 @@
     "if_removed_loss_lst = []\n",
     "\n",
     "for topk in topk_lst:\n",
-    "    keep_indices = get_topk_keep_indices(summed_scores, topk=topk)\n",
+    "    keep_indices = get_topk_keep_indices(scores[\"all_modules\"].sum(dim=0), topk=topk)\n",
     "    \n",
     "    new_loss = 0.\n",
     "    for _ in range(num_iter):\n",