05153506

context_unet.py

@@ -63,44 +63,52 @@ AvgPool = {
 }
 
 class Downsample(nn.Module):
-    def __init__(self, channels, use_conv, out_channels=None, dim=2, stride=(2,2)):
+    def __init__(self, channels, use_conv, out_channels=None, dim=2, stride=(2,2), use_checkpoint=False):
         super().__init__()
         self.channels = channels
         self.out_channels = out_channels or channels
-        # stride = config.stride
+        self.use_checkpoint = use_checkpoint
+        self.dim = dim
         if use_conv:
-            # print("conv")
             self.op = Conv[dim](channels, self.out_channels, 3, stride=stride, padding=1)
         else:
-            # print("pool")
             assert channels == self.out_channels
             self.op = AvgPool[dim](kernel_size=stride, stride=stride)
 
     def forward(self, x):
         assert x.shape[1] == self.channels
-        return self.op(x)
+        if self.use_checkpoint and isinstance(self.op, Conv[self.dim]):
+            print(f"checkpoint working in Downsample")
+            return checkpoint.checkpoint(self.op, x)
+        else:
+            return self.op(x)
 
 class Upsample(nn.Module):
-    def __init__(self, channels, use_conv, out_channels=None, dim=2, stride=(2,2)):
+    def __init__(self, channels, use_conv, out_channels=None, dim=2, stride=(2,2), use_checkpoint=False):
         super().__init__()
         self.channels = channels
         self.out_channels = out_channels
         self.use_conv = use_conv
         self.stride = stride
+        self.use_checkpoint = use_checkpoint
+
         if self.use_conv:
             self.conv = Conv[dim](self.channels, self.out_channels, 3, padding=1)
 
     def forward(self, x):
         assert x.shape[1] == self.channels
-        # stride = config.stride
-        # print(torch.tensor(x.shape[2:]))
-        # print(torch.tensor(stride))
         shape = torch.tensor(x.shape[2:]) * torch.tensor(self.stride)
         shape = tuple(shape.detach().numpy())
         # print(shape)
         x = F.interpolate(x, shape, mode='nearest')
+
         if self.use_conv:
-            x = self.conv(x)
+            if self.use_checkpoint:
+                print(f"checkpoint working in upsample")
+                return checkpoint.checkpoint(self.conv, x)
+            else:
+                x = self.conv(x)
+
         return x
 
 def zero_module(module):
@@ -335,6 +343,7 @@ class ContextUnet(nn.Module):
         #dtype = torch.float32,
         ):
         super().__init__()
+        #self.use_checkpoint = use_checkpoint
 
         if channel_mult == None:
             if image_size == 512:
@@ -433,7 +442,13 @@ class ContextUnet(nn.Module):
                             stride = stride,
                         )
                         if resblock_updown
-                        else Downsample(ch, conv_resample, out_channels=out_ch, dim=dim, stride=stride)
+                        else Downsample(ch, 
+                                        conv_resample, 
+                                        out_channels=out_ch, 
+                                        dim=dim, 
+                                        stride=stride, 
+                                        #use_checkpoint=use_checkpoint,
+                                        )
                     )
                 )
                 ch = out_ch
@@ -519,7 +534,13 @@ class ContextUnet(nn.Module):
                             stride = stride,
                         )
                         if resblock_updown
-                        else Upsample(ch, conv_resample, out_channels=out_ch, dim=dim, stride=stride)
+                        else Upsample(ch, 
+                                      conv_resample, 
+                                      out_channels=out_ch, 
+                                      dim=dim, 
+                                      stride=stride,
+                                      #use_checkpoint=use_checkpoint,
+                                      )
                     )
                     ds //= 2
                 self.output_blocks.append(TimestepEmbedSequential(*layers))

diffusion.py

@@ -23,7 +23,7 @@ import copy
 from tqdm.auto import tqdm
 # from diffusers import UNet2DModel#, UNet3DConditionModel
 # from diffusers import DDPMScheduler
-import datetime
+from datetime import datetime
 from pathlib import Path
 #from diffusers.optimization import get_cosine_schedule_with_warmup
 #from accelerate import notebook_launcher, Accelerator
@@ -241,8 +241,8 @@ class TrainConfig:
     world_size = 1#torch.cuda.device_count()
     # repeat = 2
 
-    #dim = 2
-    dim = 3#2
+    dim = 2
+    #dim = 3#2
     stride = (2,4) if dim == 2 else (2,2,4)
     num_image = 32#0#0#640#320#6400#3000#480#1200#120#3000#300#3000#6000#30#60#6000#1000#2000#20000#15000#7000#25600#3000#10000#1000#10000#5000#2560#800#2560
     batch_size = 1#1#10#50#10#50#20#50#1#2#50#20#2#100 # 10
@@ -360,7 +360,7 @@ def get_gpu_info(device):
 
 class DDPM21CM:
     def __init__(self, config):
-        config.run_name = datetime.datetime.now().strftime("%d%H%M%S") # the unique name of each experiment
+        config.run_name = datetime.now().strftime("%d%H%M%S") # the unique name of each experiment
         self.config = config
         self.ddpm = DDPMScheduler(betas=(1e-4, 0.02), num_timesteps=config.num_timesteps, img_shape=config.img_shape, device=config.device, config=config,)#, dtype=config.dtype
 
@@ -380,7 +380,7 @@ class DDPM21CM:
             #self.nn_model.module.to(config.dtype)
             print(f"{config.run_name} cuda:{torch.cuda.current_device()}/{self.config.global_rank} resumed nn_model from {config.resume} with {sum(x.numel() for x in self.nn_model.parameters())} parameters".center(self.config.str_len,'+'))
         else:
-            print(f"{config.run_name} cuda:{torch.cuda.current_device()}/{self.config.global_rank} initialized nn_model randomly with {sum(x.numel() for x in self.nn_model.parameters())} parameters".center(self.config.str_len,'+'))
+            print(f"{config.run_name} cuda:{torch.cuda.current_device()}/{self.config.global_rank} initialized nn_model randomly with {sum(x.numel() for x in self.nn_model.parameters())} parameters, {datetime.now().strftime('%d-%H:%M:%S.%f')}".center(self.config.str_len,'+'))
 
         # whether to use ema
         if config.ema:
@@ -412,9 +412,10 @@ class DDPM21CM:
             drop_prob=self.config.drop_prob, 
             dim=self.config.dim,
             ranges_dict=self.ranges_dict,
-            num_workers=min(2,len(os.sched_getaffinity(0))//self.config.world_size),
+            num_workers=min(1,len(os.sched_getaffinity(0))//self.config.world_size),
             str_len = self.config.str_len,
             )
+        #print(f"cuda:{torch.cuda.current_device()}/{self.config.global_rank}: Dataset4h5 done")
 
         dataloader_start = time()
         self.dataloader = DataLoader(
@@ -488,7 +489,7 @@ class DDPM21CM:
         global_step = 0
         for ep in range(self.config.n_epoch):
             self.ddpm.train()
-            pbar_train = tqdm(total=len(self.dataloader), file=sys.stderr, disable=True)#, mininterval=self.config.pbar_update_step)#, disable=True)#not self.accelerator.is_local_main_process)
+            pbar_train = tqdm(total=len(self.dataloader), file=sys.stderr)#, disable=True)#, mininterval=self.config.pbar_update_step)#, disable=True)#not self.accelerator.is_local_main_process)
             pbar_train.set_description(f"{socket.gethostbyname(socket.gethostname())} cuda:{torch.cuda.current_device()}/{self.config.global_rank} Epoch {ep}")
             epoch_start = time()
             for i, (x, c) in enumerate(self.dataloader):
@@ -507,6 +508,10 @@ class DDPM21CM:
                     loss = F.mse_loss(noise, noise_pred)
                     loss = loss / self.config.gradient_accumulation_steps
 
+                    #print(f"loss = {loss}")
+                    if torch.isnan(loss).any():
+                        raise ValueError(f"{socket.gethostbyname(socket.gethostname())} cuda:{torch.cuda.current_device()}/{self.config.global_rank} Epoch {ep}, loss: {loss}")
+
                 # scaler backward propogation
                 self.scaler.scale(loss).backward()
                 #loss.backward()
@@ -624,7 +629,7 @@ class DDPM21CM:
         #print(f"x_last.dtype = {x_last.dtype}")
         if save:    
             # np.save(os.path.join(self.config.output_dir, f"{self.config.run_name}{'ema' if ema else ''}.npy"), x_last)
-            savetime = datetime.datetime.now().strftime("%d%H%M%S")
+            savetime = datetime.now().strftime("%d%H%M%S")
             savename = os.path.join(self.config.output_dir, f"Tvir{params_backup[0]:.3f}-zeta{params_backup[1]:.3f}-N{self.config.num_image}-device{self.config.global_rank}-{os.path.basename(self.config.resume)}-{savetime}{'ema' if ema else ''}.npy")
             if not os.path.exists(self.config.output_dir):
                 os.makedirs(self.config.output_dir)
@@ -721,7 +726,7 @@ if __name__ == "__main__":
     ############################ training ################################
     if args.train:
         config.dataset_name = args.train
-        print(f" training, ip_addr = {socket.gethostbyname(socket.gethostname())}, master_addr = {master_addr}, local_world_size = {local_world_size}, world_size = {world_size} ".center(config.str_len,'#'))
+        print(f" training, ip = {socket.gethostbyname(socket.gethostname())}, local_world_size = {local_world_size}, world_size = {world_size}, {datetime.now().strftime('%d-%H:%M:%S.%f')} ".center(config.str_len,'#'))
         mp.spawn(
                 train, 
                 args=(world_size, local_world_size, master_addr, master_port, config), 

load_h5.py

@@ -20,7 +20,7 @@ import os
 # from diffusers import DDPMScheduler
 # from diffusers.utils import make_image_grid
 from time import time
-import datetime
+from datetime import datetime
 import concurrent.futures
 import psutil
 # from pathlib import Path
@@ -115,25 +115,33 @@ class Dataset4h5(Dataset):
         elif self.dim == 3:
             self.images = np.empty((self.num_image, 1, self.HII_DIM, self.HII_DIM, self.num_redshift), dtype=np.float32)
         # self.num_workers = len(os.sched_getaffinity(0))//torch.cuda.device_count()
-        concurrent_init_start = time()
-        with concurrent.futures.ProcessPoolExecutor(max_workers=self.num_workers) as executor:
-            concurrent_init_end = time()
-            print(f" {socket.gethostbyname(socket.gethostname())} cuda:{torch.cuda.current_device()}/{self.global_rank}, concurrently loading by {self.num_workers}/{len(os.sched_getaffinity(0))} workers, initialized after {concurrent_init_end-concurrent_init_start:.3f}s ".center(self.str_len, '-'))
-            futures = [None] * self.num_workers
-            for i, idx in enumerate(np.array_split(self.idx, self.num_workers)):
-                executor_start = time()
-                futures[i] = executor.submit(self.read_data_chunk, self.dir_name, idx, torch.cuda.current_device(), concurrent_init_end, executor_start)
 
+        concurrent_init_start = time()
+        if self.num_workers == 1:
+            print(f"{socket.gethostbyname(socket.gethostname())} cuda:{torch.cuda.current_device()}/{self.global_rank}, loading by {self.num_workers} workers, {datetime.now().strftime('%d-%H:%M:%S.%f')}".center(self.str_len, '-'))
+            self.images, self.params = self.read_data_chunk(self.dir_name, self.idx, torch.cuda.current_device(), concurrent_init_start, concurrent_init_start)
+            self.params = self.params.astype(self.images.dtype)
             concurrent_start = time()
-            start_idx = 0
-            for future in concurrent.futures.as_completed(futures):
-                images, params = future.result()
-                batch_size = params.shape[0]        
-                self.images[start_idx:start_idx+batch_size] = images
-                self.params[start_idx:start_idx+batch_size] = params
-                start_idx += batch_size
-            concurrent_end = time()
-            print(f" {socket.gethostbyname(socket.gethostname())} cuda:{torch.cuda.current_device()}/{self.global_rank}, {start_idx} images {self.images.shape} & params {self.params.shape} loaded after {concurrent_start-concurrent_init_start:.3f}/{concurrent_end-concurrent_start:.3f}s ".center(self.str_len, '-'))
+            print(f"{socket.gethostbyname(socket.gethostname())} cuda:{torch.cuda.current_device()}/{self.global_rank}, images {self.images.shape} & params {self.params.shape} loaded after {concurrent_start-concurrent_init_start:.3f}s, {datetime.now().strftime('%d-%H:%M:%S.%f')}".center(self.str_len, '-'))
+        else:
+            with concurrent.futures.ProcessPoolExecutor(max_workers=self.num_workers) as executor:
+                concurrent_init_end = time()
+                print(f" {socket.gethostbyname(socket.gethostname())} cuda:{torch.cuda.current_device()}/{self.global_rank}, concurrently loading by {self.num_workers}/{len(os.sched_getaffinity(0))} workers, initialized after {concurrent_init_end-concurrent_init_start:.3f}s ".center(self.str_len, '-'))
+                futures = [None] * self.num_workers
+                for i, idx in enumerate(np.array_split(self.idx, self.num_workers)):
+                    executor_start = time()
+                    futures[i] = executor.submit(self.read_data_chunk, self.dir_name, idx, torch.cuda.current_device(), concurrent_init_end, executor_start)
+    
+                concurrent_start = time()
+                start_idx = 0
+                for future in concurrent.futures.as_completed(futures):
+                    images, params = future.result()
+                    batch_size = params.shape[0]        
+                    self.images[start_idx:start_idx+batch_size] = images
+                    self.params[start_idx:start_idx+batch_size] = params
+                    start_idx += batch_size
+                concurrent_end = time()
+                print(f" {socket.gethostbyname(socket.gethostname())} cuda:{torch.cuda.current_device()}/{self.global_rank}, {start_idx} images {self.images.shape} & params {self.params.shape} loaded after {concurrent_start-concurrent_init_start:.3f}/{concurrent_end-concurrent_start:.3f}s ".center(self.str_len, '-'))
 
         transform_start = time()
         if self.transform:
@@ -164,7 +172,7 @@ class Dataset4h5(Dataset):
             param_start = time()
             params = f['params']['values'][idx]
             param_end = time()
-            print(f"{socket.gethostbyname(socket.gethostname())}, cuda:{torch.cuda.current_device()}/{self.global_rank}, CPU-pid {cpu_num}-{pid}: images {images.shape} & params {params.shape} loaded after {executor_start-concurrent_init_end:.3f}/{set_device-executor_start:.3f}/{open_h5py-set_device:.3f}/{images_start-open_h5py:.3f}s + {images_end-images_start:.3f}s & {param_end-param_start:.3f}s")
+            print(f"cuda:{torch.cuda.current_device()}/{self.global_rank}, CPU:{cpu_num}, images {images.shape} & params {params.shape} loaded after {executor_start-concurrent_init_end:.3f}/{set_device-executor_start:.3f}/{open_h5py-set_device:.3f}/{images_start-open_h5py:.3f}s + {images_end-images_start:.3f}s & {param_end-param_start:.3f}s")
 
         return images, params
 

perlmutter_diffusion.sbatch

@@ -3,9 +3,9 @@
 #SBATCH -J diffusion
 #SBATCH -C gpu&hbm80g
 #SBATCH -q regular #shared 
-#SBATCH -N10
+#SBATCH -N1
 #SBATCH --gpus-per-node=4
-#SBATCH -t 30:20:00
+#SBATCH -t 02:30:00
 #SBATCH --ntasks-per-node=1
 #SBATCH -oReport-%j 
 #SBATCH --mail-type=BEGIN,END,FAIL
@@ -36,12 +36,12 @@ cat $0
 #nvidia-smi
 
 srun python diffusion.py \
-    --num_image 640 \
-    --batch_size 1 \
+    --num_image 6400 \
+    --batch_size 128 \
     --n_epoch 50 \
     --num_new_img_per_gpu 20 \
     --max_num_img_per_gpu 4 \
-    --channel_mult 1 1 2 2 4 4 \
+    --channel_mult 0.5 1 2 2 4 8 \
     --gradient_accumulation_steps 1 \
     --autocast 1 \
     --use_checkpoint 1 \

tensorboard.ipynb

@@ -2,21 +2,12 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 1,
    "id": "ae45e44e-a11c-43ef-b830-c7a58a72f51e",
    "metadata": {
     "tags": []
    },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "The tensorboard extension is already loaded. To reload it, use:\n",
-      "  %reload_ext tensorboard\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "import nersc_tensorboard_helper\n",
     "%load_ext tensorboard"
@@ -24,30 +15,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 2,
    "id": "a5c088b8-5051-402f-b4ec-2b684ad5a952",
    "metadata": {},
    "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "Reusing TensorBoard on port 45739 (pid 1821871), started 2 days, 2:32:50 ago. (Use '!kill 1821871' to kill it.)"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
     {
      "data": {
       "text/html": [
        "\n",
-       "      <iframe id=\"tensorboard-frame-905898ab07792b79\" width=\"100%\" height=\"800\" frameborder=\"0\">\n",
+       "      <iframe id=\"tensorboard-frame-497d865784f48ad7\" width=\"100%\" height=\"800\" frameborder=\"0\">\n",
        "      </iframe>\n",
        "      <script>\n",
        "        (function() {\n",
-       "          const frame = document.getElementById(\"tensorboard-frame-905898ab07792b79\");\n",
+       "          const frame = document.getElementById(\"tensorboard-frame-497d865784f48ad7\");\n",
        "          const url = new URL(\"/\", window.location);\n",
-       "          const port = 45739;\n",
+       "          const port = 38971;\n",
        "          if (port) {\n",
        "            url.port = port;\n",
        "          }\n",
@@ -70,14 +52,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 3,
    "id": "2f76c0a9-2218-4073-86aa-f4f655d7642f",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/html": [
-       "<a href=\"https://jupyter.nersc.gov/user/binxia/perlmutter-login-node-base/proxy/45739/\">https://jupyter.nersc.gov/user/binxia/perlmutter-login-node-base/proxy/45739/</a>"
+       "<a href=\"https://jupyter.nersc.gov/user/binxia/perlmutter-login-node-base/proxy/38971/\">https://jupyter.nersc.gov/user/binxia/perlmutter-login-node-base/proxy/38971/</a>"
       ],
       "text/plain": [
        "<IPython.core.display.HTML object>"