0713-1516

diffusion.py

@@ -274,13 +274,12 @@ class TrainConfig:
     # 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
+    resume = False # if resume from the trained checkpoints
     # params_single = torch.tensor([0.2,0.80000023])
     # params = torch.tile(params_single,(n_sample,1)).to(device)
     # params =  params
     # data_dir = './data' # data directory
 
-
     mixed_precision = "fp16"
     gradient_accumulation_steps = 1
 
@@ -322,10 +321,6 @@ class DDPM21CM:
         # initialize the unet
         self.nn_model = ContextUnet(n_param=config.n_param, image_size=config.HII_DIM, dim=config.dim, stride=config.stride)
 
-        if config.resume and os.path.exists(config.resume):
-            # resume_file = os.path.join(config.output_dir, f"{config.resume}")
-            self.nn_model.load_state_dict(torch.load(config.resume)['unet_state_dict'])
-            print(f"resumed nn_model from {config.resume}")
         # nn_model = ContextUnet(n_param=1, image_size=28)
         self.nn_model.train()
         # print("self.ddpm.device =", self.ddpm.device)
@@ -333,6 +328,14 @@ class DDPM21CM:
         self.nn_model = DDP(self.nn_model, device_ids=[self.ddpm.device])
         # print("nn_model.device =", ddpm.device)
         # number of parameters to be trained
+
+        if config.resume and os.path.exists(config.resume):
+            # resume_file = os.path.join(config.output_dir, f"{config.resume}")
+            # self.nn_model.load_state_dict(torch.load(config.resume)['unet_state_dict'])
+            # print(f"resumed nn_model from {config.resume}")
+            self.nn_model.module.load_state_dict(torch.load(config.resume)['unet_state_dict'])
+            print(f"device {torch.cuda.current_device()} resumed nn_model from {config.resume}")
+
         self.number_of_params = sum(x.numel() for x in self.nn_model.parameters())
         print(f"Number of parameters for nn_model: {self.number_of_params}")
 
@@ -508,7 +511,7 @@ class DDPM21CM:
 
     def sample(self, params:torch.tensor=None, num_new_img=192, ema=False, entire=False, save=False):
         # n_sample = params.shape[0]
-        file = self.config.resume
+        # file = self.config.resume
 
         if params is None:
             params = torch.tensor([0.20000000000000018, 0.5055875000000001])
@@ -528,11 +531,11 @@ class DDPM21CM:
         # 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)
-        if ema:
-            self.nn_model.module.load_state_dict(torch.load(file)['ema_unet_state_dict'])
-        else:
-            self.nn_model.module.load_state_dict(torch.load(file)['unet_state_dict'])
-        print(f"device {torch.cuda.current_device()} resumed nn_model from {file}")
+        # if ema:
+        #     self.nn_model.module.load_state_dict(torch.load(file)['ema_unet_state_dict'])
+        # else:
+        #     self.nn_model.module.load_state_dict(torch.load(file)['unet_state_dict'])
+        # print(f"device {torch.cuda.current_device()} resumed nn_model from {file}")
         # nn_model = ContextUnet(n_param=1, image_size=28)
         # nn_model.train()
         # self.nn_model.to(self.ddpm.device)