0708-1342

context_unet.py

@@ -327,8 +327,10 @@ class ContextUnet(nn.Module):
                 channel_mult = (1, 1, 2, 3, 4)
             elif image_size == 64:
                 channel_mult = (1, 1, 2, 2, 4, 4)#(1, 2, 3, 4)
+            elif image_size == 32:
+                channel_mult = (1, 2, 2, 4)
             elif image_size == 28:
-                channel_mult = (1, 2)#(1, 2, 3, 4)
+                channel_mult = (1, 2, 4)#(1, 2, 3, 4)
             else:
                 raise ValueError(f"unsupported image size: {image_size}")
         # else:

diffusion.ipynb

@@ -74,24 +74,9 @@
    "cell_type": "code",
    "execution_count": 2,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9d207dc0a7a24bccb2f419954b64ddc0",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "VBox(children=(HTML(value='<center> <img\\nsrc=https://huggingface.co/front/assets/huggingface_logo-noborder.sv…"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
+   "outputs": [],
    "source": [
-    "notebook_login()"
+    "# notebook_login()"
    ]
   },
   {
@@ -272,12 +257,12 @@
     "\n",
     "    # dim = 2\n",
     "    dim = 3\n",
-    "    stride = (2,2) if dim == 2 else (2,2,2)\n",
+    "    stride = (2,2) if dim == 2 else (2,2,1)\n",
     "    num_image = 2000#32000#20000#15000#7000#25600#3000#10000#1000#10000#5000#2560#800#2560\n",
-    "    batch_size = 2#2#50#20#2#100 # 10\n",
+    "    batch_size = 2#50#20#2#100 # 10\n",
     "    n_epoch = 10#50#20#20#2#5#25 # 120\n",
-    "    HII_DIM = 64\n",
-    "    num_redshift = 128#64#512#256#256#64#512#128\n",
+    "    HII_DIM = 32#64\n",
+    "    num_redshift = 4#128#64#512#256#256#64#512#128\n",
     "    channel = 1\n",
     "    img_shape = (channel, HII_DIM, num_redshift) if dim == 2 else (channel, HII_DIM, HII_DIM, num_redshift)\n",
     "\n",
@@ -579,17 +564,16 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Number of parameters for nn_model: 306285057\n",
-      "---------------- num_image = 1000 ----------------\n",
-      "run_name = 0706-2157\n",
+      "Number of parameters for nn_model: 190142209\n",
+      "---------------- num_image = 100 -----------------\n",
+      "run_name = 0708-1342\n",
       "Launching training on one GPU.\n",
       "dataset content: <KeysViewHDF5 ['brightness_temp', 'density', 'kwargs', 'params', 'redshifts_distances', 'seeds', 'xH_box']>\n",
       "51200 images can be loaded\n",
       "field.shape = (64, 64, 514)\n",
       "params keys = [b'ION_Tvir_MIN', b'HII_EFF_FACTOR']\n",
-      "loading 1000 images randomly\n",
-      "images loaded: (1000, 1, 64, 64, 128)\n",
-      "params loaded: (1000, 2)\n"
+      "loading 100 images randomly\n",
+      "images loaded: (100, 1, 32, 32, 4)\n"
      ]
     },
     {
@@ -603,19 +587,20 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "images rescaled to [-1.0, 1.2339041233062744]\n",
-      "params rescaled to [0.0006788479393025145, 0.9997530171043563]\n"
+      "params loaded: (100, 2)\n",
+      "images rescaled to [-1.0, 1.2789411544799805]\n",
+      "params rescaled to [0.004197723271926046, 0.9944779188934443]\n"
      ]
     },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1d069efc366d480ba2515135ccb18f6c",
+       "model_id": "a9bfadac7d3841c9a5a8c3440649c4f0",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/500 [00:00<?, ?it/s]"
+       "  0%|          | 0/50 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -624,12 +609,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "93abe7efc5044d3796f67a0d243058b2",
+       "model_id": "9df4310f213742d9a7aae110fca32403",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/500 [00:00<?, ?it/s]"
+       "  0%|          | 0/50 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -638,12 +623,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "75e22a8998244becb916714cf4b8053e",
+       "model_id": "399101df4a5f4de8a4a3f155b3ade75b",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/500 [00:00<?, ?it/s]"
+       "  0%|          | 0/50 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -652,12 +637,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7364f0924d2246e49fa02a596e0afde8",
+       "model_id": "ea4834c350594a9c9cbd87727a88a6b8",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/500 [00:00<?, ?it/s]"
+       "  0%|          | 0/50 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -666,12 +651,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a5c57b0059154c2695e7fc0d426bb0d9",
+       "model_id": "b7b0d9a8c2ad456387dc1b053550c702",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/500 [00:00<?, ?it/s]"
+       "  0%|          | 0/50 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -680,12 +665,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b0c9ee9c876144eab582e196e71ed83c",
+       "model_id": "17e73c3722d64ae895f337a7379b5225",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/500 [00:00<?, ?it/s]"
+       "  0%|          | 0/50 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -694,12 +679,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d2956adb8f4c45938f2d8f93baca2998",
+       "model_id": "0a743e1a2db2445d93533c9ec5ed921f",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/500 [00:00<?, ?it/s]"
+       "  0%|          | 0/50 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -708,12 +693,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1f3031b4cd1040829ce3b457ee30e464",
+       "model_id": "cfda06e79a0e4f8b8172e9314263fb5b",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/500 [00:00<?, ?it/s]"
+       "  0%|          | 0/50 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -722,12 +707,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b4dcf375da4c4c4492d308e4bdc19358",
+       "model_id": "8cb50f206a844a9da91197c2a9ed715b",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/500 [00:00<?, ?it/s]"
+       "  0%|          | 0/50 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -736,12 +721,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "50c6936f333a451394d6761fa9a085be",
+       "model_id": "dbaed562dee44cddb3b0d17f439464b6",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/500 [00:00<?, ?it/s]"
+       "  0%|          | 0/50 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -749,7 +734,7 @@
     }
    ],
    "source": [
-    "num_image_list = [1000]#[200]#[1600,3200,6400,12800,25600]\n",
+    "num_image_list = [100]#[1000]#[200]#[1600,3200,6400,12800,25600]\n",
     "if __name__ == \"__main__\":\n",
     "    # torch.multiprocessing.set_start_method(\"spawn\")\n",
     "    # args = (config, nn_model, ddpm, optimizer, dataloader, lr_scheduler)\n",
@@ -759,7 +744,9 @@
     "        ddpm21cm = DDPM21CM(config)\n",
     "        print(f\" num_image = {ddpm21cm.config.num_image} \".center(50, '-'))\n",
     "        print(f\"run_name = {ddpm21cm.config.run_name}\")\n",
-    "        notebook_launcher(ddpm21cm.train, num_processes=1, mixed_precision='fp16')"
+    "        notebook_launcher(\n",
+    "            ddpm21cm.train, num_processes=1#, mixed_precision='fp16'\n",
+    "            )"
    ]
   },
   {
@@ -767,6 +754,37 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Number of parameters for nn_model: 306285057\n",
+      "sampling 2 images with normalized params = tensor([[0.2000, 0.5056]])\n",
+      "nn_model resumed from ./outputs/model_state-N1000\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "69eb2e5d3375414cab966a9c8db91901",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/1000 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
    "source": [
     "if __name__ == \"__main__\":\n",
     "    # num_image_list = [1600,3200,6400,12800,25600]\n",
@@ -792,22 +810,47 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 19,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "total 13G\n",
+      "-rw-r--r--  1 bxia34 pace-jw254 4.1M Jul  6 23:59 Tvir4.400000095367432-zeta131.34100341796875-N1000.npy\n",
+      "-rw-r--r--  1 bxia34 pace-jw254 2.3G Jul  6 23:06 model_state-N1000\n",
+      "drwxr-xr-x 56 bxia34 pace-jw254 4.0K Jul  6 22:09 \u001b[0m\u001b[01;34mlogs\u001b[0m/\n",
+      "-rw-r--r--  1 bxia34 pace-jw254 4.1M Jul  6 21:39 Tvir4.400000095367432-zeta131.34100341796875-N50.npy\n",
+      "-rw-r--r--  1 bxia34 pace-jw254 2.3G Jul  6 21:25 model_state-N50\n",
+      "-rw-r--r--  1 bxia34 pace-jw254 193K Jul  6 20:46 Tvir4.400000095367432-zeta131.34100341796875-N20.npy\n",
+      "-rw-r--r--  1 bxia34 pace-jw254 848M Jul  6 20:45 model_state-N20\n",
+      "-rw-r--r--  1 bxia34 pace-jw254 6.1M Jul  5 14:44 Tvir4.400000095367432-zeta131.34100341796875-N200.npy\n",
+      "-rw-r--r--  1 bxia34 pace-jw254 2.3G Jul  5 12:20 model_state-N200\n"
+     ]
+    }
+   ],
    "source": [
     "ls -lth outputs | head"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 21,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "samples.shape = (2, 1, 64, 64, 128)\n"
+     ]
+    }
+   ],
    "source": [
     "def plot_grid(samples, c=None, row=1, col=2):\n",
     "    print(\"samples.shape =\", samples.shape)\n",
-    "    for j in range(samples.shape[2]):\n",
+    "    for j in range(samples.shape[4]):\n",
     "        plt.figure(figsize = (12,6), dpi=400)\n",
     "        for i in range(len(samples)):\n",
     "            plt.subplot(row,col,i+1)\n",
@@ -819,7 +862,7 @@
     "        # plt.suptitle('simulations')\n",
     "        plt.tight_layout()\n",
     "        plt.subplots_adjust(wspace=0, hspace=0)\n",
-    "        plt.savefig(f\"test3D-{j:02d}.png\")\n",
+    "        plt.savefig(f\"test3D-{j:03d}.png\")\n",
     "        plt.close()\n",
     "        # plt.show()\n",
     "    \n",