0709-1331

diffusion.ipynb

@@ -259,9 +259,9 @@
     "    dim = 3\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#50#20#2#100 # 10\n",
+    "    batch_size = 2#2#50#20#2#100 # 10\n",
     "    n_epoch = 10#50#20#20#2#5#25 # 120\n",
-    "    HII_DIM = 32#64\n",
+    "    HII_DIM = 28#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",
@@ -564,16 +564,15 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Number of parameters for nn_model: 190142209\n",
+      "Number of parameters for nn_model: 160234497\n",
       "---------------- num_image = 100 -----------------\n",
-      "run_name = 0708-1342\n",
-      "Launching training on one GPU.\n",
+      "run_name = 0709-1331\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 100 images randomly\n",
-      "images loaded: (100, 1, 32, 32, 4)\n"
+      "images loaded: (100, 1, 28, 28, 4)\n"
      ]
     },
     {
@@ -588,14 +587,14 @@
      "output_type": "stream",
      "text": [
       "params loaded: (100, 2)\n",
-      "images rescaled to [-1.0, 1.2789411544799805]\n",
-      "params rescaled to [0.004197723271926046, 0.9944779188934443]\n"
+      "images rescaled to [-1.0, 1.1254141330718994]\n",
+      "params rescaled to [0.0022036265313531977, 0.9978807793709957]\n"
      ]
     },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a9bfadac7d3841c9a5a8c3440649c4f0",
+       "model_id": "ae9f12def1154f6cb1eb0fc8d1e1871c",
        "version_major": 2,
        "version_minor": 0
       },
@@ -609,7 +608,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9df4310f213742d9a7aae110fca32403",
+       "model_id": "cae19ac5ef7a4c34b6b57a6478dc159d",
        "version_major": 2,
        "version_minor": 0
       },
@@ -623,7 +622,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "399101df4a5f4de8a4a3f155b3ade75b",
+       "model_id": "b448b32948894b3c8e8780f1b6e6bf58",
        "version_major": 2,
        "version_minor": 0
       },
@@ -637,7 +636,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ea4834c350594a9c9cbd87727a88a6b8",
+       "model_id": "0271ab7c081a43ebb830dcfd3db145c1",
        "version_major": 2,
        "version_minor": 0
       },
@@ -651,7 +650,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b7b0d9a8c2ad456387dc1b053550c702",
+       "model_id": "1bdf4c28272840f496288545bfdbdb96",
        "version_major": 2,
        "version_minor": 0
       },
@@ -665,7 +664,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "17e73c3722d64ae895f337a7379b5225",
+       "model_id": "01b50340758b4b05891e51a616660eb8",
        "version_major": 2,
        "version_minor": 0
       },
@@ -679,7 +678,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0a743e1a2db2445d93533c9ec5ed921f",
+       "model_id": "20ba64496e5e4467b97428cf6dcdbeb5",
        "version_major": 2,
        "version_minor": 0
       },
@@ -693,7 +692,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "cfda06e79a0e4f8b8172e9314263fb5b",
+       "model_id": "8023e0bde0c3438fb218126c58fee954",
        "version_major": 2,
        "version_minor": 0
       },
@@ -707,7 +706,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8cb50f206a844a9da91197c2a9ed715b",
+       "model_id": "22881fbe8aff4ac1a3bde8735ab4fd24",
        "version_major": 2,
        "version_minor": 0
       },
@@ -721,7 +720,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "dbaed562dee44cddb3b0d17f439464b6",
+       "model_id": "3bf15c9ba5144fa69fdddae271c7dbee",
        "version_major": 2,
        "version_minor": 0
       },
@@ -734,7 +733,7 @@
     }
    ],
    "source": [
-    "num_image_list = [100]#[1000]#[200]#[1600,3200,6400,12800,25600]\n",
+    "num_image_list = [100]#[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",
@@ -744,47 +743,23 @@
     "        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(\n",
-    "            ddpm21cm.train, num_processes=1#, mixed_precision='fp16'\n",
-    "            )"
+    "        ddpm21cm.train()\n",
+    "        # notebook_launcher(ddpm21cm.train, num_processes=1, mixed_precision='fp16')"
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
+   "attachments": {},
+   "cell_type": "markdown",
    "metadata": {},
-   "outputs": [],
-   "source": []
+   "source": [
+    "# Sampling"
+   ]
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": null,
    "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"
-    }
-   ],
+   "outputs": [],
    "source": [
     "if __name__ == \"__main__\":\n",
     "    # num_image_list = [1600,3200,6400,12800,25600]\n",
@@ -810,43 +785,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": null,
    "metadata": {},
-   "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"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "ls -lth outputs | head"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "samples.shape = (2, 1, 64, 64, 128)\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "def plot_grid(samples, c=None, row=1, col=2):\n",
     "    print(\"samples.shape =\", samples.shape)\n",
@@ -899,118 +849,6 @@
     "#     # plt.imshow(images[0,0])\n",
     "#     # plt.show()"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# plot(\"outputs/0528-1433.npy\")\n",
-    "# plot(\"outputs/0520-2323.npy\")\n",
-    "# plot(\"outputs/0604-2353.npy\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# x = np.load(\"outputs/0528-1433.npy\")\n",
-    "# print(x.shape)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import torch\n",
-    "import torch.nn as nn\n",
-    "import time\n",
-    "\n",
-    "class MyModel(nn.Module):\n",
-    "    def __init__(self):\n",
-    "        super().__init__()\n",
-    "        self.fc = nn.Linear(100,50)\n",
-    "\n",
-    "    def forward(self, x):\n",
-    "        return self.fc(x)\n",
-    "\n",
-    "model = MyModel()\n",
-    "\n",
-    "device_count = torch.cuda.device_count()\n",
-    "print(\"device_count =\", device_count)\n",
-    "\n",
-    "if device_count > 1:\n",
-    "    print(f\"using {device_count} GPUs!\")\n",
-    "    model = nn.DataParallel(model)\n",
-    "\n",
-    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
-    "model.to(device)\n",
-    "\n",
-    "start_time = time.time()\n",
-    "for i in range(10):\n",
-    "    myinput = torch.randn(10,10,32000,100).to(device)\n",
-    "    output = model(myinput)\n",
-    "    print(output.shape)\n",
-    "# plt.imshow(myinput.cpu()[0])\n",
-    "# plt.show()\n",
-    "# plt.imshow(output.detach().cpu().numpy()[0])\n",
-    "# plt.show()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# import torch.distributed as dist\n",
-    "# dist.init_process_group(backend='nccl')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import numpy as np\n",
-    "import torch\n",
-    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
-    "\n",
-    "data = torch.randn((64,64,64))\n",
-    "\n",
-    "num_elements = data.numpy().size\n",
-    "element_size = data.numpy().itemsize\n",
-    "\n",
-    "print(data.dtype)\n",
-    "print(num_elements, element_size)\n",
-    "print(f\"total size = {num_elements*element_size/1024/1024} MB\")\n",
-    "\n",
-    "print(\"---\"*30)\n",
-    "data = data.to(torch.float64)\n",
-    "\n",
-    "num_elements = data.numpy().size\n",
-    "element_size = data.numpy().itemsize\n",
-    "\n",
-    "print(data.dtype)\n",
-    "print(num_elements, element_size)\n",
-    "print(f\"total size = {num_elements*element_size/1024/1024} MB\")\n",
-    "\n",
-    "print(\"---\"*30)\n",
-    "data = data.to(torch.float16)\n",
-    "\n",
-    "num_elements = data.numpy().size\n",
-    "element_size = data.numpy().itemsize\n",
-    "\n",
-    "print(data.dtype)\n",
-    "print(num_elements, element_size)\n",
-    "print(f\"total size = {num_elements*element_size/1024/1024} MB\")"
-   ]
   }
  ],
  "metadata": {