0702-1312

diffusion.ipynb

@@ -281,7 +281,7 @@
     "    lrate = 1e-4\n",
     "    lr_warmup_steps = 0#5#00\n",
     "    output_dir = \"./outputs/\"\n",
-    "    save_name = os.path.join(output_dir, 'model_state.pth')\n",
+    "    save_name = os.path.join(output_dir, 'model_state')\n",
     "    # save_freq = 1 #10 # the period of saving model\n",
     "    # cond = True # if training using the conditional information\n",
     "    # lr_decay = False #True# if using the learning rate decay\n",
@@ -460,8 +460,8 @@
     "                            'unet_state_dict': self.nn_model.state_dict(),\n",
     "                            'ema_unet_state_dict': self.ema_model.state_dict(),\n",
     "                            }\n",
-    "                        torch.save(model_state, self.config.save_name)\n",
-    "                        print('saved model at ' + self.config.save_name)\n",
+    "                        torch.save(model_state, self.config.save_name+f\"-N{self.config.num_image}\")\n",
+    "                        print('saved model at ' + self.config.save_name+f\"-N{self.config.num_image}\")\n",
     "                        # print('saved model at ' + config.save_dir + f\"model_epoch_{ep}_test_{config.run_name}.pth\")\n",
     "\n",
     "    # def rescale(self, value, type='params', to_ranges=[0,1]):\n",
@@ -537,7 +537,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b23ad327ed7d48b0856c0eb9a66be943",
+       "model_id": "6025148aa9024daaaa9f3ba7ea0c784b",
        "version_major": 2,
        "version_minor": 0
       },
@@ -563,17 +563,15 @@
      "output_type": "stream",
      "text": [
       "-------------------- round 0 ---------------------\n",
-      "resumed nn_model from ./outputs/model_state.pth\n",
       "Number of parameters for nn_model: 111048705\n",
-      "resumed ema_model from ./outputs/model_state.pth\n",
-      "run_name = 0701-1047\n",
+      "run_name = 0702-1312\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 2000 images randomly\n",
-      "images loaded: (2000, 1, 64, 64)\n"
+      "loading 1600 images randomly\n",
+      "images loaded: (1600, 1, 64, 64)\n"
      ]
     },
     {
@@ -587,20 +585,20 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "params loaded: (2000, 2)\n",
+      "params loaded: (1600, 2)\n",
       "images rescaled to [-1.0, 1.1335339546203613]\n",
-      "params rescaled to [0.00010391510804133771, 0.9998883049763877]\n"
+      "params rescaled to [0.0001702067256199591, 0.9998215201715621]\n"
      ]
     },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6b39caf3d0c74a389e15275ad0422549",
+       "model_id": "82818076b9634cf2821ddc51d00752b7",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -609,12 +607,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "08c1cd82b84d46988cad13d1679f15dc",
+       "model_id": "e439288486394473af821f74b9df8128",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -623,12 +621,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1dd978d29d0c41efa15af7d8b85c2554",
+       "model_id": "4e7f1080ae9c44c1bab77db59431976e",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -637,12 +635,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a0c894fe81a54e2899a4efdb6cb838ab",
+       "model_id": "9332319e53044145914249fa79781b55",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -651,12 +649,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "cf595c2ffd3a4a74b0a22a546848f961",
+       "model_id": "9427a193e4c24989a8c24202493991ee",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -665,12 +663,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "010dc738cc4040d9b5589ef1621b2db8",
+       "model_id": "f2100a02d6874dd39e74265651d68cfb",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -679,12 +677,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a880d0d06737431788b6e4dc937bf388",
+       "model_id": "8ce81cfbb0044152be24b1c75bf8381a",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -693,12 +691,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d65997de48844e35b530ce02adbe2e5b",
+       "model_id": "015753cf1d6e4b0c92efb17cc4440e37",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -707,12 +705,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f8b7d43ee1234eb8b9e238a52b6fe9ad",
+       "model_id": "bdff45db0dff4c29bb8b7f4ee2f1e92b",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -721,12 +719,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "35ca65ab005e4c3db3163abe39e9fcf5",
+       "model_id": "7499ebc69b984c508a1da52e298c5485",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -735,12 +733,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b27fac35150745f098f228025b95fc10",
+       "model_id": "1e5055f2bddd4d13aa386258f7ae9bd9",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -749,12 +747,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "4f470fe28a074f0f9d7ee3c6163fcecf",
+       "model_id": "adb82dcd1b4a465f87d64c0fe99444e0",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -763,12 +761,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ba8a1edd0738475183dbe5ec874d4eb7",
+       "model_id": "b4206deea8e249138628a01478e8994f",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -777,12 +775,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a912fc9cfe86402cb4a51617d15815a6",
+       "model_id": "17c826779a944efd955cc516e9b04c92",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -791,12 +789,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "151ce7bd0ab44372852a6ce202ace98c",
+       "model_id": "1fed82eb826f433e9d5a85c51aea757a",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -805,12 +803,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f9748b330d324009ab5b7b8404d6849d",
+       "model_id": "01c2f68a24d84757904303fc35681cee",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -819,12 +817,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "525e0eabf5354eafa6a85cf09e8c1d30",
+       "model_id": "fbba18c4ea074bf0a7861e7a4a8a0c64",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -833,12 +831,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "23d44749ed294e8fa5611ba5653af5c3",
+       "model_id": "bd6bf27417c542d19a6e5ced172036db",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -847,12 +845,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "782ecf23b66b4a248d1ab3e27166617b",
+       "model_id": "79a49db719464cb18283c5225484e050",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -861,12 +859,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a03f4cf684fb4d67879e2e9bfd4b1cd8",
+       "model_id": "97919d69fc3a40cfbfc03286418c4016",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -875,12 +873,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "02b85ce288ce4ab0a83ea0da7ccc21a6",
+       "model_id": "9102d3133c1c49948273e2ae92ae30a5",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -889,12 +887,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1f603325944a4674bbbb846ada2a2984",
+       "model_id": "5fb87cbf52374eae8f20528ea02ffbcd",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -903,12 +901,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8f0dc0491a3749c9b8284cab27686961",
+       "model_id": "c4317d2231484a68a00b4d458694f289",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -917,12 +915,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f545ac9f48ef4434b4c6537f1c77f3ab",
+       "model_id": "423e00e991d241f08d53ed1ab13c5ee1",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -931,12 +929,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ce605dce528c4ca69b8dea4858216c92",
+       "model_id": "d67d402de8fe4f0493527fc53106b1d8",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -945,12 +943,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d907f2fefb974be38a050677785bd1a7",
+       "model_id": "4df6775e00264b53bd3945860c18dc22",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -959,12 +957,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f041d4784e5740c881929426c477408a",
+       "model_id": "0445848a37ac46a8b88b24ee34a4fb88",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -973,12 +971,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0f4130e3aa104e578812e22b98778193",
+       "model_id": "9d25fe13f99b411293ca0b62b16c9935",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -987,12 +985,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0bc61fe4d64b47abbe97bb358cce85eb",
+       "model_id": "ed5fb6ff03694a199b6d62a66709a745",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1001,12 +999,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1ae54adaa4434bfa93a748e52d259139",
+       "model_id": "e5fd2eb6029d47e89208a08f9df196b6",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1015,12 +1013,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a032ece6fc2740cd8506608fbef80325",
+       "model_id": "f472161eceff48bea9f8ed68c4fb67dc",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1029,12 +1027,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "cc2a8c07ec6047f7b99466c1e5e4d310",
+       "model_id": "67e983dac33841edbcdb91fe331933e8",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1043,12 +1041,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "5cf08a55fefd4859b2cf35cf859c4869",
+       "model_id": "61f079cd7523453abf4887be98b0a6a8",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1057,12 +1055,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f336833fd44e40eb98eb9e42ed011cf6",
+       "model_id": "a7bdb93a4bc148dcbd367e05a2ac46d8",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1071,12 +1069,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e2843d610d50429a8cc65e9c4051a9c5",
+       "model_id": "73f1305e213140859346ebb239f9d7bf",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1085,12 +1083,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9302becd557a4d5faa8ccde0aacbb4bc",
+       "model_id": "d142bdd2d1c64c4091207338267cea71",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1099,12 +1097,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "66c8271cfece47918ecab1522cb5e3e7",
+       "model_id": "434b794f81984cb8a8ba540d05eb938a",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1113,12 +1111,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "491cbe8db73d41bc8e1a980e6343086b",
+       "model_id": "0a33673984684e8ca65e8cc65d4cc1f5",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1127,12 +1125,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "26623ba0a2424c448c91f7359f81c6cd",
+       "model_id": "b9d5166185a04053a06f369619a8654f",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1141,12 +1139,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "aea1c9aabd7848e19ee0dea60ba43966",
+       "model_id": "7fcbd0f7a09e474b8cd04904ae5fd4ab",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1155,12 +1153,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7dd06fec865945aa8d8181312c2f08d8",
+       "model_id": "a189ea99daa2489c81733a3d25ccadb8",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1169,12 +1167,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "645917cca25a40e8b311cf0fca4b5fac",
+       "model_id": "87e779fce4fc47e580debfd737790b86",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1183,12 +1181,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "dd7b77cb6200422f87c7ca0a3218b008",
+       "model_id": "1e44cc8124a94f0dabb5489a06fe3ad3",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1197,12 +1195,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a2441972859a4739bf804ac7960c33cd",
+       "model_id": "97f9d2a71669403ab9672205739af7c8",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1211,12 +1209,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "5852dd77bf6e409baa1f3d049e41fdb3",
+       "model_id": "be24e4fc75794090980b3501f99d3949",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1225,12 +1223,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "64d524b098b94eff8b8b715fad31e614",
+       "model_id": "bcbecbae045746a883fb6cd1aaf0f8ff",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1239,12 +1237,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8592a5488c024a6d952ce2aae60dc5a5",
+       "model_id": "9cb1a9c35dde4aebb1e4ade969e11c57",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1253,12 +1251,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "55e79db64846496f937b8429ba61f01f",
+       "model_id": "7b6fe2498b504053866cd20060e641ac",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1267,12 +1265,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c17e75bcb5ea4cf99d8860ff01721cc1",
+       "model_id": "12eb61759b264bac902285108cad196a",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1281,12 +1279,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "adb7ec9dfd2c469680b49d2d5c337b71",
+       "model_id": "378297e998ee40dd88cbff6f6ecdd849",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/40 [00:00<?, ?it/s]"
+       "  0%|          | 0/32 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1294,12 +1292,14 @@
     }
    ],
    "source": [
+    "num_image_list = [1600,3200]#,6400,12800,25600]\n",
+    "\n",
     "if __name__ == \"__main__\":\n",
     "    # args = (config, nn_model, ddpm, optimizer, dataloader, lr_scheduler)\n",
-    "    num_round = 1\n",
-    "    for i in range(num_round):\n",
+    "    for i, num_image in enumerate(num_image_list):\n",
     "        print(f\" round {i} \".center(50, '-'))\n",
     "        ddpm21cm = DDPM21CM()\n",
+    "        ddpm21cm.config.num_image = num_image\n",
     "        print(f\"run_name = {ddpm21cm.config.run_name}\")\n",
     "        notebook_launcher(ddpm21cm.train, num_processes=1)"
    ]
@@ -1313,13 +1313,14 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "total 980M\n",
+      "total 968M\n",
+      "-rw-r--r--  1 bxia34 848M Jul  1 10:59 model_state.pth\n",
+      "drwxr-xr-x 12 bxia34 4.0K Jul  1 10:50 \u001b[0m\u001b[01;34mlogs\u001b[0m/\n",
+      "-rw-r--r--  1 bxia34 3.1M Jul  1 10:37 Tvir4.800000190734863-zeta131.34100341796875-N32000.npy\n",
       "-rw-r--r--  1 bxia34 3.1M Jul  1 04:47 Tvir5.4770002365112305-zeta200.0-N32000.npy\n",
       "-rw-r--r--  1 bxia34 3.1M Jul  1 04:29 Tvir4.698999881744385-zeta30.0-N32000.npy\n",
       "-rw-r--r--  1 bxia34 3.1M Jul  1 04:11 Tvir5.599999904632568-zeta19.03700065612793-N32000.npy\n",
       "-rw-r--r--  1 bxia34 3.1M Jul  1 03:53 Tvir4.400000095367432-zeta131.34100341796875-N32000.npy\n",
-      "-rw-r--r--  1 bxia34 848M Jul  1 03:35 model_state.pth\n",
-      "drwxr-xr-x 11 bxia34 4.0K Jul  1 01:04 \u001b[0m\u001b[01;34mlogs\u001b[0m/\n",
       "-rw-r--r--  1 bxia34 3.1M Jul  1 00:23 Tvir4.800000190734863-zeta131.34100341796875-N20000.npy\n",
       "-rw-r--r--  1 bxia34 3.1M Jul  1 00:05 Tvir5.4770002365112305-zeta200.0-N20000.npy\n",
       "-rw-r--r--  1 bxia34 3.1M Jun 30 23:47 Tvir4.698999881744385-zeta30.0-N20000.npy\n",
@@ -1386,115 +1387,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "dce3b8a47c404689b08fd2c9e811cd8f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/1000 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
-   "source": [
-    "ddpm21cm = DDPM21CM()\n",
-    "ddpm21cm.sample(\"./outputs/model_state.pth\", params=torch.tensor([4.4, 131.341]))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "resumed nn_model from ./outputs/model_state.pth\n",
-      "Number of parameters for nn_model: 111048705\n",
-      "resumed ema_model from ./outputs/model_state.pth\n",
-      "sampling 192 images with normalized params = tensor([[0.8000, 0.0377]])\n",
-      "nn_model resumed from ./outputs/model_state.pth\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "390067ed3a1c463a9573d82a1ba66bcc",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/1000 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
-   "source": [
-    "ddpm21cm = DDPM21CM()\n",
-    "ddpm21cm.sample(\"./outputs/model_state.pth\", params=torch.tensor((5.6, 19.037)))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "resumed nn_model from ./outputs/model_state.pth\n",
-      "Number of parameters for nn_model: 111048705\n",
-      "resumed ema_model from ./outputs/model_state.pth\n",
-      "sampling 192 images with normalized params = tensor([[0.3495, 0.0833]])\n",
-      "nn_model resumed from ./outputs/model_state.pth\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "5ae0f37108724e6f90136b646bfd691b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/1000 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
-   "source": [
-    "ddpm21cm = DDPM21CM()\n",
-    "ddpm21cm.sample(\"./outputs/model_state.pth\", params=torch.tensor((4.699, 30)))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "resumed nn_model from ./outputs/model_state.pth\n",
-      "Number of parameters for nn_model: 111048705\n",
-      "resumed ema_model from ./outputs/model_state.pth\n",
-      "sampling 192 images with normalized params = tensor([[0.7385, 0.7917]])\n",
-      "nn_model resumed from ./outputs/model_state.pth\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f62cdd2196de4ad6aa237e22c7a45446",
+       "model_id": "ebb1ef2f95274dbf8ccddd40b8ee931a",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1507,44 +1400,24 @@
     }
    ],
    "source": [
-    "ddpm21cm = DDPM21CM()\n",
-    "ddpm21cm.sample(\"./outputs/model_state.pth\", params=torch.tensor((5.477, 200)))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "resumed nn_model from ./outputs/model_state.pth\n",
-      "Number of parameters for nn_model: 111048705\n",
-      "resumed ema_model from ./outputs/model_state.pth\n",
-      "sampling 192 images with normalized params = tensor([[0.4000, 0.5056]])\n",
-      "nn_model resumed from ./outputs/model_state.pth\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "2668c98390d6464fa9d913643fc3d998",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/1000 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
-   "source": [
-    "ddpm21cm = DDPM21CM()\n",
-    "ddpm21cm.sample(\"./outputs/model_state.pth\", params=torch.tensor((4.8, 131.341)))"
+    "if __name__ == \"__main__\":\n",
+    "    # args = (config, nn_model, ddpm, optimizer, dataloader, lr_scheduler)\n",
+    "    repeat = 800\n",
+    "    for i, num_image in enumerate(num_image_list):\n",
+    "        ddpm21cm = DDPM21CM()\n",
+    "        ddpm21cm.sample(f\"./outputs/model_state-N{num_image}\", params=torch.tensor([4.4, 131.341]), repeat=repeat)\n",
+    "\n",
+    "        ddpm21cm = DDPM21CM()\n",
+    "        ddpm21cm.sample(f\"./outputs/model_state-N{num_image}\", params=torch.tensor((5.6, 19.037)), repeat=repeat)\n",
+    "\n",
+    "        ddpm21cm = DDPM21CM()\n",
+    "        ddpm21cm.sample(f\"./outputs/model_state-N{num_image}\", params=torch.tensor((4.699, 30)), repeat=repeat)\n",
+    "\n",
+    "        ddpm21cm = DDPM21CM()\n",
+    "        ddpm21cm.sample(f\"./outputs/model_state-N{num_image}\", params=torch.tensor((5.477, 200)), repeat=repeat)\n",
+    "\n",
+    "        ddpm21cm = DDPM21CM()\n",
+    "        ddpm21cm.sample(f\"./outputs/model_state-N{num_image}\", params=torch.tensor((4.8, 131.341)), repeat=repeat)"
    ]
   },
   {
@@ -1556,14 +1429,19 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "total 980M\n",
+      "total 995M\n",
+      "-rw-r--r--  1 bxia34 pace-jw254 3.1M Jul  1 12:31 Tvir4.800000190734863-zeta131.34100341796875-N2000.npy\n",
+      "-rw-r--r--  1 bxia34 pace-jw254 3.1M Jul  1 12:12 Tvir5.4770002365112305-zeta200.0-N2000.npy\n",
+      "-rw-r--r--  1 bxia34 pace-jw254 3.1M Jul  1 11:54 Tvir4.698999881744385-zeta30.0-N2000.npy\n",
+      "-rw-r--r--  1 bxia34 pace-jw254 3.1M Jul  1 11:35 Tvir5.599999904632568-zeta19.03700065612793-N2000.npy\n",
+      "-rw-r--r--  1 bxia34 pace-jw254 3.1M Jul  1 11:17 Tvir4.400000095367432-zeta131.34100341796875-N2000.npy\n",
+      "-rw-r--r--  1 bxia34 pace-jw254 848M Jul  1 10:59 model_state.pth\n",
+      "drwxr-xr-x 12 bxia34 pace-jw254 4.0K Jul  1 10:50 \u001b[0m\u001b[01;34mlogs\u001b[0m/\n",
       "-rw-r--r--  1 bxia34 pace-jw254 3.1M Jul  1 10:37 Tvir4.800000190734863-zeta131.34100341796875-N32000.npy\n",
       "-rw-r--r--  1 bxia34 pace-jw254 3.1M Jul  1 04:47 Tvir5.4770002365112305-zeta200.0-N32000.npy\n",
       "-rw-r--r--  1 bxia34 pace-jw254 3.1M Jul  1 04:29 Tvir4.698999881744385-zeta30.0-N32000.npy\n",
       "-rw-r--r--  1 bxia34 pace-jw254 3.1M Jul  1 04:11 Tvir5.599999904632568-zeta19.03700065612793-N32000.npy\n",
       "-rw-r--r--  1 bxia34 pace-jw254 3.1M Jul  1 03:53 Tvir4.400000095367432-zeta131.34100341796875-N32000.npy\n",
-      "-rw-r--r--  1 bxia34 pace-jw254 848M Jul  1 03:35 model_state.pth\n",
-      "drwxr-xr-x 11 bxia34 pace-jw254 4.0K Jul  1 01:04 \u001b[0m\u001b[01;34mlogs\u001b[0m/\n",
       "-rw-r--r--  1 bxia34 pace-jw254 3.1M Jul  1 00:23 Tvir4.800000190734863-zeta131.34100341796875-N20000.npy\n",
       "-rw-r--r--  1 bxia34 pace-jw254 3.1M Jul  1 00:05 Tvir5.4770002365112305-zeta200.0-N20000.npy\n",
       "-rw-r--r--  1 bxia34 pace-jw254 3.1M Jun 30 23:47 Tvir4.698999881744385-zeta30.0-N20000.npy\n",