0712-1513

diffusion.ipynb

@@ -33,7 +33,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 31,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -77,7 +77,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 32,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -95,11 +95,26 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 34,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "9bbf7e9db9ce426d9c59d6f6d8e8df29",
+       "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"
+    }
+   ],
    "source": [
-    "# notebook_login()"
+    "notebook_login()"
    ]
   },
   {
@@ -817,75 +832,99 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 29,
+   "execution_count": 9,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "outputs/model_state-N500-device0\n",
-      "outputs/model_state-N500-device1\n",
-      "len(model_states) = 2\n",
-      "epoch\n",
-      "unet_state_dict\n",
-      "epoch 9 9\n",
-      "odict_keys(['token_embedding.weight', 'token_embedding.bias', 'time_embed.0.weight', 'time_embed.0.bias', 'time_embed.2.weight', 'time_embed.2.bias', 'input_blocks.0.0.weight', 'input_blocks.0.0.bias', 'input_blocks.1.0.in_layers.0.weight', 'input_blocks.1.0.in_layers.0.bias', 'input_blocks.1.0.in_layers.2.weight', 'input_blocks.1.0.in_layers.2.bias', 'input_blocks.1.0.emb_layers.1.weight', 'input_blocks.1.0.emb_layers.1.bias', 'input_blocks.1.0.out_layers.0.weight', 'input_blocks.1.0.out_layers.0.bias', 'input_blocks.1.0.out_layers.3.weight', 'input_blocks.1.0.out_layers.3.bias', 'input_blocks.1.1.norm.weight', 'input_blocks.1.1.norm.bias', 'input_blocks.1.1.qkv.weight', 'input_blocks.1.1.qkv.bias', 'input_blocks.1.1.proj_out.weight', 'input_blocks.1.1.proj_out.bias', 'input_blocks.2.0.in_layers.0.weight', 'input_blocks.2.0.in_layers.0.bias', 'input_blocks.2.0.in_layers.2.weight', 'input_blocks.2.0.in_layers.2.bias', 'input_blocks.2.0.emb_layers.1.weight', 'input_blocks.2.0.emb_layers.1.bias', 'input_blocks.2.0.out_layers.0.weight', 'input_blocks.2.0.out_layers.0.bias', 'input_blocks.2.0.out_layers.3.weight', 'input_blocks.2.0.out_layers.3.bias', 'input_blocks.2.1.norm.weight', 'input_blocks.2.1.norm.bias', 'input_blocks.2.1.qkv.weight', 'input_blocks.2.1.qkv.bias', 'input_blocks.2.1.proj_out.weight', 'input_blocks.2.1.proj_out.bias', 'input_blocks.3.0.op.weight', 'input_blocks.3.0.op.bias', 'input_blocks.4.0.in_layers.0.weight', 'input_blocks.4.0.in_layers.0.bias', 'input_blocks.4.0.in_layers.2.weight', 'input_blocks.4.0.in_layers.2.bias', 'input_blocks.4.0.emb_layers.1.weight', 'input_blocks.4.0.emb_layers.1.bias', 'input_blocks.4.0.out_layers.0.weight', 'input_blocks.4.0.out_layers.0.bias', 'input_blocks.4.0.out_layers.3.weight', 'input_blocks.4.0.out_layers.3.bias', 'input_blocks.4.0.skip_connection.weight', 'input_blocks.4.0.skip_connection.bias', 'input_blocks.5.0.in_layers.0.weight', 'input_blocks.5.0.in_layers.0.bias', 'input_blocks.5.0.in_layers.2.weight', 'input_blocks.5.0.in_layers.2.bias', 'input_blocks.5.0.emb_layers.1.weight', 'input_blocks.5.0.emb_layers.1.bias', 'input_blocks.5.0.out_layers.0.weight', 'input_blocks.5.0.out_layers.0.bias', 'input_blocks.5.0.out_layers.3.weight', 'input_blocks.5.0.out_layers.3.bias', 'input_blocks.6.0.op.weight', 'input_blocks.6.0.op.bias', 'input_blocks.7.0.in_layers.0.weight', 'input_blocks.7.0.in_layers.0.bias', 'input_blocks.7.0.in_layers.2.weight', 'input_blocks.7.0.in_layers.2.bias', 'input_blocks.7.0.emb_layers.1.weight', 'input_blocks.7.0.emb_layers.1.bias', 'input_blocks.7.0.out_layers.0.weight', 'input_blocks.7.0.out_layers.0.bias', 'input_blocks.7.0.out_layers.3.weight', 'input_blocks.7.0.out_layers.3.bias', 'input_blocks.7.0.skip_connection.weight', 'input_blocks.7.0.skip_connection.bias', 'input_blocks.8.0.in_layers.0.weight', 'input_blocks.8.0.in_layers.0.bias', 'input_blocks.8.0.in_layers.2.weight', 'input_blocks.8.0.in_layers.2.bias', 'input_blocks.8.0.emb_layers.1.weight', 'input_blocks.8.0.emb_layers.1.bias', 'input_blocks.8.0.out_layers.0.weight', 'input_blocks.8.0.out_layers.0.bias', 'input_blocks.8.0.out_layers.3.weight', 'input_blocks.8.0.out_layers.3.bias', 'middle_block.0.in_layers.0.weight', 'middle_block.0.in_layers.0.bias', 'middle_block.0.in_layers.2.weight', 'middle_block.0.in_layers.2.bias', 'middle_block.0.emb_layers.1.weight', 'middle_block.0.emb_layers.1.bias', 'middle_block.0.out_layers.0.weight', 'middle_block.0.out_layers.0.bias', 'middle_block.0.out_layers.3.weight', 'middle_block.0.out_layers.3.bias', 'middle_block.1.norm.weight', 'middle_block.1.norm.bias', 'middle_block.1.qkv.weight', 'middle_block.1.qkv.bias', 'middle_block.1.proj_out.weight', 'middle_block.1.proj_out.bias', 'middle_block.2.in_layers.0.weight', 'middle_block.2.in_layers.0.bias', 'middle_block.2.in_layers.2.weight', 'middle_block.2.in_layers.2.bias', 'middle_block.2.emb_layers.1.weight', 'middle_block.2.emb_layers.1.bias', 'middle_block.2.out_layers.0.weight', 'middle_block.2.out_layers.0.bias', 'middle_block.2.out_layers.3.weight', 'middle_block.2.out_layers.3.bias', 'output_blocks.0.0.in_layers.0.weight', 'output_blocks.0.0.in_layers.0.bias', 'output_blocks.0.0.in_layers.2.weight', 'output_blocks.0.0.in_layers.2.bias', 'output_blocks.0.0.emb_layers.1.weight', 'output_blocks.0.0.emb_layers.1.bias', 'output_blocks.0.0.out_layers.0.weight', 'output_blocks.0.0.out_layers.0.bias', 'output_blocks.0.0.out_layers.3.weight', 'output_blocks.0.0.out_layers.3.bias', 'output_blocks.0.0.skip_connection.weight', 'output_blocks.0.0.skip_connection.bias', 'output_blocks.1.0.in_layers.0.weight', 'output_blocks.1.0.in_layers.0.bias', 'output_blocks.1.0.in_layers.2.weight', 'output_blocks.1.0.in_layers.2.bias', 'output_blocks.1.0.emb_layers.1.weight', 'output_blocks.1.0.emb_layers.1.bias', 'output_blocks.1.0.out_layers.0.weight', 'output_blocks.1.0.out_layers.0.bias', 'output_blocks.1.0.out_layers.3.weight', 'output_blocks.1.0.out_layers.3.bias', 'output_blocks.1.0.skip_connection.weight', 'output_blocks.1.0.skip_connection.bias', 'output_blocks.2.0.in_layers.0.weight', 'output_blocks.2.0.in_layers.0.bias', 'output_blocks.2.0.in_layers.2.weight', 'output_blocks.2.0.in_layers.2.bias', 'output_blocks.2.0.emb_layers.1.weight', 'output_blocks.2.0.emb_layers.1.bias', 'output_blocks.2.0.out_layers.0.weight', 'output_blocks.2.0.out_layers.0.bias', 'output_blocks.2.0.out_layers.3.weight', 'output_blocks.2.0.out_layers.3.bias', 'output_blocks.2.0.skip_connection.weight', 'output_blocks.2.0.skip_connection.bias', 'output_blocks.2.1.conv.weight', 'output_blocks.2.1.conv.bias', 'output_blocks.3.0.in_layers.0.weight', 'output_blocks.3.0.in_layers.0.bias', 'output_blocks.3.0.in_layers.2.weight', 'output_blocks.3.0.in_layers.2.bias', 'output_blocks.3.0.emb_layers.1.weight', 'output_blocks.3.0.emb_layers.1.bias', 'output_blocks.3.0.out_layers.0.weight', 'output_blocks.3.0.out_layers.0.bias', 'output_blocks.3.0.out_layers.3.weight', 'output_blocks.3.0.out_layers.3.bias', 'output_blocks.3.0.skip_connection.weight', 'output_blocks.3.0.skip_connection.bias', 'output_blocks.4.0.in_layers.0.weight', 'output_blocks.4.0.in_layers.0.bias', 'output_blocks.4.0.in_layers.2.weight', 'output_blocks.4.0.in_layers.2.bias', 'output_blocks.4.0.emb_layers.1.weight', 'output_blocks.4.0.emb_layers.1.bias', 'output_blocks.4.0.out_layers.0.weight', 'output_blocks.4.0.out_layers.0.bias', 'output_blocks.4.0.out_layers.3.weight', 'output_blocks.4.0.out_layers.3.bias', 'output_blocks.4.0.skip_connection.weight', 'output_blocks.4.0.skip_connection.bias', 'output_blocks.5.0.in_layers.0.weight', 'output_blocks.5.0.in_layers.0.bias', 'output_blocks.5.0.in_layers.2.weight', 'output_blocks.5.0.in_layers.2.bias', 'output_blocks.5.0.emb_layers.1.weight', 'output_blocks.5.0.emb_layers.1.bias', 'output_blocks.5.0.out_layers.0.weight', 'output_blocks.5.0.out_layers.0.bias', 'output_blocks.5.0.out_layers.3.weight', 'output_blocks.5.0.out_layers.3.bias', 'output_blocks.5.0.skip_connection.weight', 'output_blocks.5.0.skip_connection.bias', 'output_blocks.5.1.conv.weight', 'output_blocks.5.1.conv.bias', 'output_blocks.6.0.in_layers.0.weight', 'output_blocks.6.0.in_layers.0.bias', 'output_blocks.6.0.in_layers.2.weight', 'output_blocks.6.0.in_layers.2.bias', 'output_blocks.6.0.emb_layers.1.weight', 'output_blocks.6.0.emb_layers.1.bias', 'output_blocks.6.0.out_layers.0.weight', 'output_blocks.6.0.out_layers.0.bias', 'output_blocks.6.0.out_layers.3.weight', 'output_blocks.6.0.out_layers.3.bias', 'output_blocks.6.0.skip_connection.weight', 'output_blocks.6.0.skip_connection.bias', 'output_blocks.6.1.norm.weight', 'output_blocks.6.1.norm.bias', 'output_blocks.6.1.qkv.weight', 'output_blocks.6.1.qkv.bias', 'output_blocks.6.1.proj_out.weight', 'output_blocks.6.1.proj_out.bias', 'output_blocks.7.0.in_layers.0.weight', 'output_blocks.7.0.in_layers.0.bias', 'output_blocks.7.0.in_layers.2.weight', 'output_blocks.7.0.in_layers.2.bias', 'output_blocks.7.0.emb_layers.1.weight', 'output_blocks.7.0.emb_layers.1.bias', 'output_blocks.7.0.out_layers.0.weight', 'output_blocks.7.0.out_layers.0.bias', 'output_blocks.7.0.out_layers.3.weight', 'output_blocks.7.0.out_layers.3.bias', 'output_blocks.7.0.skip_connection.weight', 'output_blocks.7.0.skip_connection.bias', 'output_blocks.7.1.norm.weight', 'output_blocks.7.1.norm.bias', 'output_blocks.7.1.qkv.weight', 'output_blocks.7.1.qkv.bias', 'output_blocks.7.1.proj_out.weight', 'output_blocks.7.1.proj_out.bias', 'output_blocks.8.0.in_layers.0.weight', 'output_blocks.8.0.in_layers.0.bias', 'output_blocks.8.0.in_layers.2.weight', 'output_blocks.8.0.in_layers.2.bias', 'output_blocks.8.0.emb_layers.1.weight', 'output_blocks.8.0.emb_layers.1.bias', 'output_blocks.8.0.out_layers.0.weight', 'output_blocks.8.0.out_layers.0.bias', 'output_blocks.8.0.out_layers.3.weight', 'output_blocks.8.0.out_layers.3.bias', 'output_blocks.8.0.skip_connection.weight', 'output_blocks.8.0.skip_connection.bias', 'output_blocks.8.1.norm.weight', 'output_blocks.8.1.norm.bias', 'output_blocks.8.1.qkv.weight', 'output_blocks.8.1.qkv.bias', 'output_blocks.8.1.proj_out.weight', 'output_blocks.8.1.proj_out.bias', 'out.0.weight', 'out.0.bias', 'out.2.weight', 'out.2.bias'])\n",
-      "exactly same\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
-    "import torch\n",
-    "import os\n",
+    "# import torch\n",
+    "# import os\n",
     "\n",
-    "def compare_models(num_gpus):\n",
-    "    model_states = []\n",
+    "# def compare_models(num_gpus):\n",
+    "#     model_states = []\n",
     "    \n",
-    "    for gpu_id in range(num_gpus):\n",
-    "        filename = f\"outputs/model_state-N500-device{gpu_id}\"\n",
-    "        if os.path.exists(filename):\n",
-    "            state_dict = torch.load(filename, map_location='cpu')\n",
-    "            model_states.append(state_dict)\n",
-    "            print(filename)\n",
-    "        else:\n",
-    "            print(f\"File {filename} not found!\")\n",
-    "            return False\n",
+    "#     for gpu_id in range(num_gpus):\n",
+    "#         filename = f\"outputs/model_state-N40-device{gpu_id}\"\n",
+    "#         if os.path.exists(filename):\n",
+    "#             state_dict = torch.load(filename, map_location='cpu')\n",
+    "#             model_states.append(state_dict)\n",
+    "#             print(filename)\n",
+    "#         else:\n",
+    "#             print(f\"File {filename} not found!\")\n",
+    "#             return False\n",
     "    \n",
-    "    # Compare all model state_dicts\n",
-    "    print(\"len(model_states) =\", len(model_states))\n",
-    "    base_state = model_states[0]\n",
-    "    for state in model_states[1:]:\n",
-    "        for key in base_state.keys():\n",
-    "            # print(key, base_state[key], state[key])\n",
-    "            print(key)\n",
-    "        print(\"epoch\", base_state['epoch'], state['epoch'])\n",
-    "\n",
-    "        print(base_state['unet_state_dict'].keys())\n",
-    "        for key in base_state['unet_state_dict']:\n",
-    "            # print(key)\n",
-    "            if not torch.equal(base_state['unet_state_dict'][key], state['unet_state_dict'][key]):\n",
-    "                print(\"different\")\n",
-    "                return \n",
-    "            # else:\n",
-    "        print(\"exactly same\")\n",
-    "\n",
-    "            # if key == 'epoch':\n",
-    "            #     print(base_state[key], state[key])\n",
-    "            # else:\n",
-    "            #     print(base_state[key], state[key])\n",
-    "            #     if not torch.equal(base_state[key], state[key]):\n",
-    "            #     # if not (base_state[key] == state[key]):\n",
-    "            #         print(f\"Mismatch found in parameter {key}\")\n",
-    "            #         return False\n",
+    "#     # Compare all model state_dicts\n",
+    "#     print(\"len(model_states) =\", len(model_states))\n",
+    "#     base_state = model_states[0]\n",
+    "#     for state in model_states[1:]:\n",
+    "#         for key in base_state.keys():\n",
+    "#             # print(key, base_state[key], state[key])\n",
+    "#             print(key)\n",
+    "#         print(\"epoch\", base_state['epoch'], state['epoch'])\n",
+    "\n",
+    "#         print(base_state['unet_state_dict'].keys())\n",
+    "#         for key in base_state['unet_state_dict']:\n",
+    "#             # print(key)\n",
+    "#             if not torch.equal(base_state['unet_state_dict'][key], state['unet_state_dict'][key]):\n",
+    "#                 print(\"different\")\n",
+    "#                 return \n",
+    "#             # else:\n",
+    "#         print(\"exactly same\")\n",
+    "\n",
+    "#             # if key == 'epoch':\n",
+    "#             #     print(base_state[key], state[key])\n",
+    "#             # else:\n",
+    "#             #     print(base_state[key], state[key])\n",
+    "#             #     if not torch.equal(base_state[key], state[key]):\n",
+    "#             #     # if not (base_state[key] == state[key]):\n",
+    "#             #         print(f\"Mismatch found in parameter {key}\")\n",
+    "#             #         return False\n",
     "    \n",
-    "    # print(\"All models are identical!\")\n",
-    "    # return True\n",
+    "#     # print(\"All models are identical!\")\n",
+    "#     # return True\n",
     "\n",
-    "if __name__ == \"__main__\":\n",
-    "    # epoch_to_check = 0  # specify the epoch you want to check\n",
-    "    num_gpus = 2  # specify the number of GPUs used in training\n",
-    "    compare_models(num_gpus)"
+    "# if __name__ == \"__main__\":\n",
+    "#     # epoch_to_check = 0  # specify the epoch you want to check\n",
+    "#     num_gpus = torch.cuda.device_count()  # specify the number of GPUs used in training\n",
+    "#     compare_models(num_gpus)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "test = np.random.normal(0,1,(800,1,64,64,512))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "12.5"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "(test.itemsize*test.size) / 1024/1024/1024"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "del test"
    ]
   }
  ],

diffusion.py

@@ -61,6 +61,7 @@ import torch.multiprocessing as mp
 from torch.utils.data.distributed import DistributedSampler
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.distributed import init_process_group, destroy_process_group
+import torch.distributed as dist
 
 # %%
 def ddp_setup(rank: int, world_size: int):
@@ -180,14 +181,12 @@ class DDPMScheduler(nn.Module):
             x_i = 1/torch.sqrt(self.alpha_t[i])*(x_i-eps*self.beta_t[i]/torch.sqrt(1-self.bar_alpha_t[i])) + torch.sqrt(self.beta_t[i])*z
             
             pbar_sample.update(1)
-            # pbar_sample.set_postfix(step=i)
             
-            # print("x_i.shape =", x_i.shape)
             # store only part of the intermediate steps
-            if i%20==0:# or i==0:# or i<8:
-                x_i_entire.append(x_i.detach().cpu().numpy())
-        x_i = x_i.detach().cpu().numpy()
+            # if i%20==0:# or i==0:# or i<8:
+            #     x_i_entire.append(x_i.detach().cpu().numpy())
         x_i_entire = np.array(x_i_entire)
+        x_i = x_i.detach().cpu().numpy()
         return x_i, x_i_entire
 
 
@@ -225,7 +224,7 @@ class TrainConfig:
     ###########################
     ## hardcoding these here ##
     ###########################
-    push_to_hub = True
+    push_to_hub = True 
     hub_model_id = "Xsmos/ml21cm"
     hub_private_repo = False
     dataset_name = "/storage/home/hcoda1/3/bxia34/scratch/LEN128-DIM64-CUB8.h5"
@@ -265,14 +264,14 @@ class TrainConfig:
     # seed = 0
     # save_dir = './outputs/'
 
-    save_freq = 0#.1 # the period of sampling
+    save_period = np.infty#.1 # the period of sampling
     # general parameters for the name and logger    
     # device = "cuda" if torch.cuda.is_available() else "cpu"
     lrate = 1e-4
     lr_warmup_steps = 0#5#00
     output_dir = "./outputs/"
     save_name = os.path.join(output_dir, 'model_state')
-    # save_freq = 1 #10 # the period of saving model
+    # save_period = 1 #10 # the period of saving model
     # cond = True # if training using the conditional information
     # lr_decay = False #True# if using the learning rate decay
     resume = save_name # if resume from the trained checkpoints
@@ -394,8 +393,8 @@ class DDPM21CM:
             # distributed_type="MULTI_GPU",
         )
         # print("!!!!!!!!!!!!!!!!!!!self.accelerator.device:", self.accelerator.device)
-        if self.accelerator.is_main_process:
-        # if torch.cuda.current_device() == 0:
+        # if self.accelerator.is_main_process:
+        if torch.cuda.current_device() == 0:
             if self.config.output_dir is not None:
                 os.makedirs(self.config.output_dir, exist_ok=True)
             if self.config.push_to_hub:
@@ -427,7 +426,7 @@ class DDPM21CM:
             pbar_train = tqdm(total=len(self.dataloader), disable=not self.accelerator.is_local_main_process)
             pbar_train.set_description(f"device {torch.cuda.current_device()}, Epoch {ep}")
             for i, (x, c) in enumerate(self.dataloader):
-                print(f"device {torch.cuda.current_device()}, x[:10,0,:2,:2,:2] =", x[:10,0,:2,:2,:2])
+                # print(f"device {torch.cuda.current_device()}, x[:,0,:2,0,0] =", x[:,0,:2,0,0])
                 with self.accelerator.accumulate(self.nn_model):
                     x = x.to(self.config.device)
                     xt, noise, ts = self.ddpm.add_noise(x)
@@ -460,7 +459,7 @@ class DDPM21CM:
                 self.accelerator.log(logs, step=global_step)
                 global_step += 1
 
-            # if ep == config.n_epoch-1 or (ep+1)*config.save_freq==1:
+            # if ep == config.n_epoch-1 or (ep+1)*config.save_period==1:
             self.save(ep)
 
         del self.nn_model
@@ -470,9 +469,9 @@ class DDPM21CM:
 
     def save(self, ep):
         # save model
-        if self.accelerator.is_main_process:
-        # if torch.cuda.current_device() == 0:
-            if ep == self.config.n_epoch-1 or (ep+1)*self.config.save_freq==1:
+        # if self.accelerator.is_main_process:
+        if torch.cuda.current_device() == 0:
+            if ep == self.config.n_epoch-1 or (ep+1) % self.config.save_period == 0:
                 self.nn_model.eval()
                 with torch.no_grad():
                     if self.config.push_to_hub:
@@ -488,8 +487,9 @@ class DDPM21CM:
                             'unet_state_dict': self.nn_model.module.state_dict(),
                             # 'ema_unet_state_dict': self.ema_model.state_dict(),
                             }
-                        torch.save(model_state, self.config.save_name+f"-N{self.config.num_image}-device{torch.cuda.current_device()}")
-                        print(f'device {torch.cuda.current_device()} saved model at ' + self.config.save_name+f"-N{self.config.num_image}-device{torch.cuda.current_device()}")
+                        save_name = self.config.save_name+f"-N{self.config.num_image}-epoch{ep}-device{torch.cuda.current_device()}"
+                        torch.save(model_state, save_name)
+                        print(f'device {torch.cuda.current_device()} saved model at ' + save_name)
                         # print('saved model at ' + config.save_dir + f"model_epoch_{ep}_test_{config.run_name}.pth")
 
     # def rescale(self, value, type='params', to_ranges=[0,1]):
@@ -506,7 +506,7 @@ class DDPM21CM:
         value = value * (to[1]-to[0]) + to[0]
         return value 
 
-    def sample(self, file, params:torch.tensor=None, repeat=192, ema=False, entire=False):
+    def sample(self, file, params:torch.tensor=None, num_new_img=192, ema=False, entire=False, save=False):
         # n_sample = params.shape[0]
         
         if params is None:
@@ -516,8 +516,8 @@ class DDPM21CM:
             params_backup = params.numpy().copy()
             params = self.rescale(params, self.ranges_dict['params'], to=[0,1])
 
-        print(f"sampling {repeat} images with normalized params = {params}")
-        params = params.repeat(repeat,1)
+        print(f"sampling {num_new_img} images with normalized params = {params}")
+        params = params.repeat(num_new_img,1)
         assert params.dim() == 2, "params must be a 2D torch.tensor"
         # print("params =", params)
         # print("params =", params)
@@ -526,16 +526,16 @@ class DDPM21CM:
         # del self.ema_model, self.nn
         # params = torch.tile(params, (n_sample,1)).to(device)
 
-        nn_model = ContextUnet(n_param=self.config.n_param, image_size=self.config.HII_DIM, dim=self.config.dim, stride=self.config.stride).to(self.config.device)
+        # nn_model = ContextUnet(n_param=self.config.n_param, image_size=self.config.HII_DIM, dim=self.config.dim, stride=self.config.stride).to(self.config.device)
         if ema:
-            nn_model.load_state_dict(torch.load(file)['ema_unet_state_dict'])
+            self.nn_model.load_state_dict(torch.load(file)['ema_unet_state_dict'])
         else:
-            nn_model.load_state_dict(torch.load(file)['unet_state_dict'])
+            self.nn_model.load_state_dict(torch.load(file)['unet_state_dict'])
         print(f"nn_model resumed from {file}")
         # nn_model = ContextUnet(n_param=1, image_size=28)
         # nn_model.train()
-        nn_model.to(self.ddpm.device)
-        nn_model.eval()
+        # self.nn_model.to(self.ddpm.device)
+        self.nn_model.eval()
 
         # self.ema_model = ContextUnet(n_param=config.n_param, image_size=config.HII_DIM, dim=config.dim, stride=config.stride).to(config.device)
         # self.ema_model.load_state_dict(torch.load(os.path.join(config.output_dir, f"{config.resume}"))['ema_unet_state_dict'])
@@ -543,27 +543,27 @@ class DDPM21CM:
 
         with torch.no_grad():
             x_last, x_entire = self.ddpm.sample(
-                nn_model=nn_model, 
+                nn_model=self.nn_model, 
                 params=params.to(self.config.device), 
                 device=self.config.device, 
                 guide_w=self.config.guide_w
                 )
 
-        # np.save(os.path.join(self.config.output_dir, f"{self.config.run_name}{'ema' if ema else ''}.npy"), x_last)
-        np.save(os.path.join(self.config.output_dir, f"Tvir{params_backup[0]}-zeta{params_backup[1]}-N{self.config.num_image}{'ema' if ema else ''}.npy"), x_last)
-
-        if entire:
-            np.save(os.path.join(self.config.output_dir, f"Tvir{params_backup[0]}-zeta{params_backup[1]}-N{self.config.num_image}{'ema' if ema else ''}_entire.npy"), x_last)
-# print("device =", config.device)
-
+        if save:    
+            # np.save(os.path.join(self.config.output_dir, f"{self.config.run_name}{'ema' if ema else ''}.npy"), x_last)
+            np.save(os.path.join(self.config.output_dir, f"Tvir{params_backup[0]}-zeta{params_backup[1]}-N{self.config.num_image}{'ema' if ema else ''}.npy"), x_last)
+            if entire:
+                np.save(os.path.join(self.config.output_dir, f"Tvir{params_backup[0]}-zeta{params_backup[1]}-N{self.config.num_image}{'ema' if ema else ''}_entire.npy"), x_last)
+        else:
+            return x_last
 # %%
-def main(rank, world_size):
+def train(rank, world_size):
     config = TrainConfig()
     config.world_size = world_size
 
     ddp_setup(rank, world_size)
     
-    num_image_list = [500]#[200]#[1600,3200,6400,12800,25600]
+    num_image_list = [100]#[200]#[1600,3200,6400,12800,25600]
     for i, num_image in enumerate(num_image_list):
         config.num_image = num_image
         # config.world_size = world_size
@@ -578,17 +578,11 @@ def main(rank, world_size):
 if __name__ == "__main__":
     # torch.multiprocessing.set_start_method("spawn")
     # args = (config, nn_model, ddpm, optimizer, dataloader, lr_scheduler)
-    world_size = 2#torch.cuda.device_count()
+    world_size = torch.cuda.device_count()
 
-    mp.spawn(main, args=(world_size,), nprocs=world_size, join=True)
+    mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)
     # notebook_launcher(ddpm21cm.train, num_processes=1, mixed_precision='fp16')
 
-# %%
-# torch.cuda.set_device(0)
-
-# %%
-# print(torch.cuda.__dir__())
-
 # %%
 # print("torch.cuda.is_initialized() =", torch.cuda.is_initialized())
 # print("torch.cuda.get_device_name() =", torch.cuda.get_device_name())
@@ -601,31 +595,67 @@ if __name__ == "__main__":
 # print(torch.cuda.memory())
 # print('here')
 # print(torch.cuda.memory_summary())
-
 # %% [markdown]
 # # Sampling
 
 # %%
-# if __name__ == "__main__":
-#     # num_image_list = [1600,3200,6400,12800,25600]
-#     num_image_list = [1000]
-#     # num_image_list = [3200,6400,12800,25600]
-#     # args = (config, nn_model, ddpm, optimizer, dataloader, lr_scheduler)
-#     repeat = 2
-#     config = TrainConfig()
-#     for i, num_image in enumerate(num_image_list):
-#         config.num_image = num_image
-#         ddpm21cm = DDPM21CM(config)
 
-#         ddpm21cm.sample(f"./outputs/model_state-N{num_image}", params=torch.tensor([4.4, 131.341]), repeat=repeat)
+def generate_samples(model, num_new_img, max_num_img_per_gpu, rank, world_size):
+    samples = []
+    for _ in ranges(num_new_img // max_num_img_per_gpu):    
+        sample = model.module.sample(filename, params=torch.tensor([4.4, 131.341]), num_new_img=max_num_img_per_gpu)
+        samples.append(sample)
+        # model.sample(filename, params=torch.tensor((5.6, 19.037)), num_new_img=max_num_img_per_gpu)
+        # model.sample(filename, params=torch.tensor((4.699, 30)), num_new_img=max_num_img_per_gpu)
+        # model.sample(filename, params=torch.tensor((5.477, 200)), num_new_img=max_num_img_per_gpu)
+        # model.sample(filename, params=torch.tensor((4.8, 131.341)), num_new_img=max_num_img_per_gpu)
+    samples = np.concatenate(samples, axis=0)
+
+    samples_list = [np.empty_like(samples) for _ in range(world_size)]
+    dist.all_gather_object(samples_list, samples)
+
+    if rank == 0:
+        all_samples = np.concatenate(samples_list, axis=0)
+        return all_samples
+    else:
+        return None
+
+
+def sample(rank, world_size, model, num_new_img, max_num_img_per_gpu, return_dict):
+    ddp_setup(rank, world_size)
+
+    samples = generate_samples(model, num_new_img, max_num_img_per_gpu, rank, world_size)
+
+    if rank == 0:
+        return_dict['samples'] = samples
+
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    world_size = torch.cuda.device_count()
+    # num_image_list = [1600,3200,6400,12800,25600]
+    num_image_list = [1000]
+    num_new_img = 12
+    max_num_img_per_gpu = 2
+
+    config = TrainConfig()
+    config.world_size = world_size
+
+    for num_image in num_image_list:
+        filename = f"./outputs/model_state-N{num_image}-epoch9-device0"
+        config.num_image = num_image
+        ddpm21cm = DDPM21CM(config)
 
-#         # ddpm21cm.sample(f"./outputs/model_state-N{num_image}", params=torch.tensor((5.6, 19.037)), repeat=repeat)
+        manager = np.Manager()
+        return_dict = manager.dict()
 
-#         # ddpm21cm.sample(f"./outputs/model_state-N{num_image}", params=torch.tensor((4.699, 30)), repeat=repeat)
+        mp.spawn(sample, args=(world_size, ddpm21cm, num_new_img, max_num_img_per_gpu, return_dict), nprocs=world_size, join=True)
 
-#         # ddpm21cm.sample(f"./outputs/model_state-N{num_image}", params=torch.tensor((5.477, 200)), repeat=repeat)
+        if "samples" in return_dict:
+            samples = return_dict["samples"]
+            print(f"Generated samples shape: {samples.shape}")
 
-#         # ddpm21cm.sample(f"./outputs/model_state-N{num_image}", params=torch.tensor((4.8, 131.341)), repeat=repeat)
 
 # %%
 # ls -lth outputs | head

load_h5.py

@@ -60,6 +60,13 @@ class Dataset4h5(Dataset):
         self.images = torch.from_numpy(self.images)
         print(f"images rescaled to [{self.images.min()}, {self.images.max()}]")
 
+        # print("before self.images.shape =", self.images.shape)
+        # self.images = torch.ones_like(self.images) * torch.arange(len(self.images))[:,None,None,None,None]
+        # # self.images = self.images.numpy()
+        # print("after self.images.shape =", self.images.shape)
+        # print(self.images[:6,0,:2,0,0])
+        # # self.images = self.images.numpy()
+
         cond_filter = torch.bernoulli(torch.ones(len(self.params),1)-self.drop_prob).repeat(1,self.params.shape[1]).numpy()
         self.params = torch.from_numpy(self.params*cond_filter)
         print(f"params rescaled to [{self.params.min()}, {self.params.max()}]")
@@ -98,12 +105,6 @@ class Dataset4h5(Dataset):
             self.params = f['params']['values'][self.idx]
             print("params loaded:", self.params.shape)
 
-            # print("before self.images.shape =", self.images.shape)
-            self.images = torch.ones_like(torch.from_numpy(self.images)) * torch.arange(len(self.images))[:,None,None,None,None]
-            self.images = self.images.numpy()
-            # print("after self.images.shape =", self.images.shape)
-            print(self.images[:6,0,:2,0,0])
-            # self.images = self.images.numpy()
             
             # plt.imshow(self.images[0,0,0])
             # plt.show()