0712-1603

diffusion.py

@@ -506,9 +506,10 @@ class DDPM21CM:
         value = value * (to[1]-to[0]) + to[0]
         return value 
 
-    def sample(self, file, params:torch.tensor=None, num_new_img=192, ema=False, entire=False, save=False):
+    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
+
         if params is None:
             params = torch.tensor([0.20000000000000018, 0.5055875000000001])
             params_backup = params.numpy().copy()
@@ -604,12 +605,12 @@ if __name__ == "__main__":
 def generate_samples(ddpm21cm, num_new_img, max_num_img_per_gpu, rank, world_size):
     samples = []
     for _ in range(num_new_img // max_num_img_per_gpu):    
-        sample = ddpm21cm.sample(filename, params=torch.tensor([4.4, 131.341]), num_new_img=max_num_img_per_gpu)
+        sample = ddpm21cm.sample(params=torch.tensor([4.4, 131.341]), num_new_img=max_num_img_per_gpu)
         samples.append(sample)
-        # ddpm21cm.sample(filename, params=torch.tensor((5.6, 19.037)), num_new_img=max_num_img_per_gpu)
-        # ddpm21cm.sample(filename, params=torch.tensor((4.699, 30)), num_new_img=max_num_img_per_gpu)
-        # ddpm21cm.sample(filename, params=torch.tensor((5.477, 200)), num_new_img=max_num_img_per_gpu)
-        # ddpm21cm.sample(filename, params=torch.tensor((4.8, 131.341)), num_new_img=max_num_img_per_gpu)
+        # ddpm21cm.sample(params=torch.tensor((5.6, 19.037)), num_new_img=max_num_img_per_gpu)
+        # ddpm21cm.sample(params=torch.tensor((4.699, 30)), num_new_img=max_num_img_per_gpu)
+        # ddpm21cm.sample(params=torch.tensor((5.477, 200)), num_new_img=max_num_img_per_gpu)
+        # ddpm21cm.sample(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)]
@@ -648,8 +649,8 @@ if __name__ == "__main__":
     # print("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
+        config.resume = f"./outputs/model_state-N{num_image}-epoch9-device0"
 
         # print("ddpm21cm = DDPM21CM(config)")
         manager = mp.Manager()