0715-1309

context_unet.py

@@ -318,7 +318,7 @@ class ContextUnet(nn.Module):
         encoder_channels = None,
         dim = 2,
         stride = (2,2),
-        use_fp16 = False,
+        dtype = torch.float32,
         ):
         super().__init__()
 
@@ -351,8 +351,8 @@ class ContextUnet(nn.Module):
 
         # self.n_param = n_param
         self.model_channels = model_channels
-        self.use_fp16 = use_fp16
-        self.dtype = torch.float16 if self.use_fp16 else torch.float32
+        # self.use_fp16 = use_fp16
+        self.dtype = dtype#torch.float16 if self.use_fp16 else torch.float32
 
         self.token_embedding = nn.Linear(n_param, model_channels * 4)
 

diffusion.py

@@ -96,7 +96,7 @@ def ddp_setup(rank: int, world_size: int):
 
 # %%
 class DDPMScheduler(nn.Module):
-    def __init__(self, betas: tuple, num_timesteps: int, img_shape: list, device='cpu'):
+    def __init__(self, betas: tuple, num_timesteps: int, img_shape: list, device='cpu', dtype=torch.float32):
         super().__init__()
         
         beta_1, beta_T = betas
@@ -112,6 +112,8 @@ class DDPMScheduler(nn.Module):
         self.alpha_t = 1 - self.beta_t
         # self.bar_alpha_t = torch.exp(torch.cumsum(torch.log(self.alpha_t), dim=0))
         self.bar_alpha_t = torch.cumprod(self.alpha_t, dim=0)
+        # self.use_fp16 = use_fp16
+        self.dtype = dtype#torch.float16 if self.use_fp16 else torch.float32
 
     def add_noise(self, clean_images):
         shape = clean_images.shape
@@ -280,7 +282,8 @@ class TrainConfig:
     # params =  params
     # data_dir = './data' # data directory
 
-    use_fp16 = True
+    use_fp16 = False
+    dtype = torch.float16 if use_fp16 else torch.float32
     mixed_precision = "fp16"
     gradient_accumulation_steps = 1
 
@@ -317,10 +320,10 @@ class DDPM21CM:
         # # print("shape_loaded =", self.shape_loaded)
         # self.dataloader = DataLoader(dataset, batch_size=config.batch_size, shuffle=True)
         # del dataset
-        self.ddpm = DDPMScheduler(betas=(1e-4, 0.02), num_timesteps=config.num_timesteps, img_shape=config.img_shape, device=config.device)
+        self.ddpm = DDPMScheduler(betas=(1e-4, 0.02), num_timesteps=config.num_timesteps, img_shape=config.img_shape, device=config.device, dtype=config.dtype)
 
         # initialize the unet
-        self.nn_model = ContextUnet(n_param=config.n_param, image_size=config.HII_DIM, dim=config.dim, stride=config.stride, use_fp16=config.use_fp16)
+        self.nn_model = ContextUnet(n_param=config.n_param, image_size=config.HII_DIM, dim=config.dim, stride=config.stride, dtype=config.dtype)
 
         # nn_model = ContextUnet(n_param=1, image_size=28)
         self.nn_model.train()
@@ -344,7 +347,7 @@ class DDPM21CM:
         if config.ema:
             self.ema = EMA(config.ema_rate)
             if config.resume and os.path.exists(config.resume):
-                self.ema_model = ContextUnet(n_param=config.n_param, image_size=config.HII_DIM, dim=config.dim, stride=config.stride, use_fp16=config.use_fp16).to(config.device)
+                self.ema_model = ContextUnet(n_param=config.n_param, image_size=config.HII_DIM, dim=config.dim, stride=config.stride, dtype=config.dtype).to(config.device)
                 self.ema_model.load_state_dict(torch.load(config.resume)['ema_unet_state_dict'])
                 print(f"resumed ema_model from {config.resume}")
             else:
@@ -433,6 +436,9 @@ class DDPM21CM:
                 # 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)
+                    print("x = x.to(self.config.device), x.dtype =", x.dtype)
+                    x = x.to(self.config.dtype)
+                    print("x = x.to(self.dtype), x.dtype =", x.dtype)
                     xt, noise, ts = self.ddpm.add_noise(x)
                     
                     if self.config.guide_w == -1:
@@ -440,6 +446,8 @@ class DDPM21CM:
                     else:
                         c = c.to(self.config.device)
                         noise_pred = self.nn_model(xt, ts, c)
+
+                    print("noise_pred = self.nn_model(xt, ts, c), noise_pred.dtype =", noise_pred.dtype)
                     
                     loss = F.mse_loss(noise, noise_pred)
                     self.accelerator.backward(loss)