0521-1738

diffusion.ipynb

@@ -234,7 +234,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 90,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -256,30 +256,32 @@
     "        self.bar_alpha_t = torch.cumprod(self.alpha_t, dim=0)\n",
     "\n",
     "    def add_noise(self, clean_images):\n",
-    "        shape = clean_images.shape\n",
-    "        expand = torch.ones(len(shape)-1, dtype=int)\n",
+    "        self.shape = clean_images.shape\n",
+    "        expand = torch.ones(len(self.shape)-1, dtype=int)\n",
     "        # ts_expand = ts.view(ts.shape[0], *expand.tolist())\n",
     "        # expand = [1 for i in range(len(shape)-1)]\n",
     "\n",
     "        noise = torch.randn_like(clean_images).to(self.device)\n",
-    "        ts = torch.randint(0, self.num_timesteps, (shape[0],)).to(self.device)\n",
+    "        ts = torch.randint(0, self.num_timesteps, (self.shape[0],)).to(self.device)\n",
     "                \n",
     "        # test_expand = test.view(test.shape[0],*expand)\n",
     "        # extend_dim = [None for i in range(shape.dim()-1)]\n",
     "        noisy_images = (\n",
-    "            clean_images * torch.sqrt(self.bar_alpha_t[ts]).view(shape[0], *expand.tolist())\n",
-    "            + noise * torch.sqrt(1-self.bar_alpha_t[ts]).view(shape[0], *expand.tolist())\n",
+    "            clean_images * torch.sqrt(self.bar_alpha_t[ts]).view(self.shape[0], *expand.tolist())\n",
+    "            + noise * torch.sqrt(1-self.bar_alpha_t[ts]).view(self.shape[0], *expand.tolist())\n",
     "            )\n",
     "        # print(x_t.shape)\n",
     "\n",
     "        return noisy_images, noise, ts\n",
     "\n",
-    "    def sample(self, nn_model, n_sample, shape, device, test_param, guide_w = 0):\n",
-    "        x_i = torch.randn(n_sample, *shape).to(device)\n",
+    "    def sample(self, nn_model, params, device, guide_w = 0):\n",
+    "        n_sample = params.shape[0]\n",
+    "        print(\"params.shape[0], len(params)\", params.shape[0], len(params))\n",
+    "        x_i = torch.randn(n_sample, *self.shape[1:]).to(device)\n",
     "        # print(\"x_i.shape =\", x_i.shape)\n",
     "        if guide_w != -1:\n",
-    "            c_i = test_param\n",
-    "            uncond_tokens = torch.zeros(int(n_sample), test_param.shape[1]).to(device)\n",
+    "            c_i = params\n",
+    "            uncond_tokens = torch.zeros(int(n_sample), params.shape[1]).to(device)\n",
     "            # uncond_tokens = torch.tensor(np.float32(np.array([0,0]))).to(device)\n",
     "            # uncond_tokens = uncond_tokens.repeat(int(n_sample),1)\n",
     "            c_i = torch.cat((c_i, uncond_tokens), 0)\n",
@@ -295,7 +297,7 @@
     "            t_is = torch.tensor([i]).to(device)\n",
     "            t_is = t_is.repeat(n_sample)\n",
     "\n",
-    "            z = torch.randn(n_sample, *shape).to(device) if i > 0 else 0\n",
+    "            z = torch.randn(n_sample, *self.shape[1:]).to(device) if i > 0 else 0\n",
     "\n",
     "            if guide_w == -1:\n",
     "                # eps = nn_model(x_i, t_is, return_dict=False)[0]\n",
@@ -303,7 +305,7 @@
     "                # 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\n",
     "            else:\n",
     "                # double batch\n",
-    "                x_i = x_i.repeat(2, *torch.ones(len(shape), dtype=int).tolist())\n",
+    "                x_i = x_i.repeat(2, *torch.ones(len(self.shape[1:]), dtype=int).tolist())\n",
     "                t_is = t_is.repeat(2)\n",
     "\n",
     "                # split predictions and compute weighting\n",
@@ -336,7 +338,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 91,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -378,7 +380,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 92,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -403,7 +405,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 93,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -432,7 +434,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 94,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -447,7 +449,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 95,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -461,7 +463,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 96,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -479,7 +481,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 97,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -560,7 +562,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 98,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -593,7 +595,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 99,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -642,7 +644,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 100,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -671,7 +673,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 101,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -913,7 +915,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 102,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -943,7 +945,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 57,
+   "execution_count": 103,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -967,7 +969,7 @@
     "    n_epoch = 10#2#5#25 # 120\n",
     "    num_timesteps = 1000#1000 # 1000, 500; DDPM time steps\n",
     "    train_batch_size = 10#10#20#2#100 # 10\n",
-    "    n_sample = 24 # 64, the number of samples in sampling process\n",
+    "    # n_sample = 24 # 64, the number of samples in sampling process\n",
     "    n_param = 2\n",
     "    guide_w = 0#-1#0#-1#0.1#[0,0.1] #[0,0.5,2] strength of generative guidance\n",
     "    drop_prob = 0.28 # only takes effect when guide_w != -1\n",
@@ -987,9 +989,9 @@
     "    # cond = True # if training using the conditional information\n",
     "    # lr_decay = False #True# if using the learning rate decay\n",
     "    resume = 'model_state.pth' # if resume from the trained checkpoints\n",
-    "    test_param_single = torch.tensor([0.2,0.80000023])\n",
-    "    test_param = torch.tile(test_param_single,(n_sample,1)).to(device)\n",
-    "    # test_param =  test_param\n",
+    "    # params_single = torch.tensor([0.2,0.80000023])\n",
+    "    # params = torch.tile(params_single,(n_sample,1)).to(device)\n",
+    "    # params =  params\n",
     "    # data_dir = './data' # data directory\n",
     "\n",
     "    output_dir = \"./outputs/\"\n",
@@ -1006,7 +1008,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 58,
+   "execution_count": 104,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1016,7 +1018,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 59,
+   "execution_count": 105,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1025,7 +1027,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 60,
+   "execution_count": 106,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1049,205 +1051,205 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 61,
+   "execution_count": 107,
    "metadata": {},
    "outputs": [],
    "source": [
-    "def train_loop(config, nn_model, ddpm, optimizer, dataloader, lr_scheduler):    \n",
-    "    ########################\n",
-    "    ## ready for training ##\n",
-    "    ########################\n",
-    "    # initialize the dataset\n",
-    "    # num_image = 600\n",
-    "    # HII_DIM = 64\n",
-    "    # num_redshift = 64#512#128\n",
-    "    # dataset = Dataset4h5(config.dataset_name, num_image=config.num_image, HII_DIM=config.HII_DIM, num_redshift=config.num_redshift, drop_prob=config.drop_prob)\n",
-    "    # dataloader = DataLoader(dataset, batch_size=config.train_batch_size, shuffle=True)\n",
-    "    # tf = transforms.Compose([transforms.ToTensor()]) # mnist is already normalised 0 to 1\n",
-    "    # dataset = MNIST(\"./data\", train=True, download=True, transform=tf)\n",
-    "    # dataloader = DataLoader(dataset, batch_size=config.train_batch_size, shuffle=True, num_workers=5)    # Initialize accelerator and tensorboard logging\n",
-    "    accelerator = Accelerator(\n",
-    "        mixed_precision=config.mixed_precision,\n",
-    "        gradient_accumulation_steps=config.gradient_accumulation_steps,\n",
-    "        log_with=\"tensorboard\",\n",
-    "        project_dir=os.path.join(config.output_dir, \"logs\"),\n",
-    "    )\n",
-    "    if accelerator.is_main_process:\n",
-    "        if config.output_dir is not None:\n",
-    "            os.makedirs(config.output_dir, exist_ok=True)\n",
-    "        if config.push_to_hub:\n",
-    "            repo_id = create_repo(\n",
-    "                repo_id=config.hub_model_id or Path(config.output_dir).name, exist_ok=True\n",
-    "            ).repo_id\n",
-    "        accelerator.init_trackers(f\"{config.date}\")\n",
-    "\n",
-    "    nn_model, optimizer, dataloader, lr_scheduler = accelerator.prepare(\n",
-    "        nn_model, optimizer, dataloader, lr_scheduler)\n",
+    "# def train_loop(config, nn_model, ddpm, optimizer, dataloader, lr_scheduler):    \n",
+    "#     ########################\n",
+    "#     ## ready for training ##\n",
+    "#     ########################\n",
+    "#     # initialize the dataset\n",
+    "#     # num_image = 600\n",
+    "#     # HII_DIM = 64\n",
+    "#     # num_redshift = 64#512#128\n",
+    "#     # dataset = Dataset4h5(config.dataset_name, num_image=config.num_image, HII_DIM=config.HII_DIM, num_redshift=config.num_redshift, drop_prob=config.drop_prob)\n",
+    "#     # dataloader = DataLoader(dataset, batch_size=config.train_batch_size, shuffle=True)\n",
+    "#     # tf = transforms.Compose([transforms.ToTensor()]) # mnist is already normalised 0 to 1\n",
+    "#     # dataset = MNIST(\"./data\", train=True, download=True, transform=tf)\n",
+    "#     # dataloader = DataLoader(dataset, batch_size=config.train_batch_size, shuffle=True, num_workers=5)    # Initialize accelerator and tensorboard logging\n",
+    "#     accelerator = Accelerator(\n",
+    "#         mixed_precision=config.mixed_precision,\n",
+    "#         gradient_accumulation_steps=config.gradient_accumulation_steps,\n",
+    "#         log_with=\"tensorboard\",\n",
+    "#         project_dir=os.path.join(config.output_dir, \"logs\"),\n",
+    "#     )\n",
+    "#     if accelerator.is_main_process:\n",
+    "#         if config.output_dir is not None:\n",
+    "#             os.makedirs(config.output_dir, exist_ok=True)\n",
+    "#         if config.push_to_hub:\n",
+    "#             repo_id = create_repo(\n",
+    "#                 repo_id=config.hub_model_id or Path(config.output_dir).name, exist_ok=True\n",
+    "#             ).repo_id\n",
+    "#         accelerator.init_trackers(f\"{config.date}\")\n",
+    "\n",
+    "#     nn_model, optimizer, dataloader, lr_scheduler = accelerator.prepare(\n",
+    "#         nn_model, optimizer, dataloader, lr_scheduler)\n",
     "        \n",
-    "    # initialize the DDPM\n",
-    "    # logger = SummaryWriter(os.path.join(\"runs\", config.run_name)) # To log\n",
-    "\n",
-    "    # ddpm = DDPMScheduler(betas=(1e-4, 0.02), num_timesteps=config.num_timesteps, device=config.device)\n",
-    "\n",
-    "    # # initialize the unet\n",
-    "    # nn_model = ContextUnet(n_param=config.n_param, image_size=config.HII_DIM)\n",
-    "    # # nn_model = ContextUnet(n_param=1, image_size=28)\n",
-    "    # nn_model.train()\n",
-    "    # nn_model.to(ddpm.device)\n",
-    "\n",
-    "    # parameters to be optimized\n",
-    "    # params_to_optimize = [\n",
-    "    #     {'params': nn_model.parameters()}\n",
-    "    # ]\n",
-    "\n",
-    "    # number of parameters to be trained\n",
-    "    number_of_params = sum(x.numel() for x in nn_model.parameters())\n",
-    "    print(f\"Number of parameters for unet: {number_of_params}\")\n",
-    "\n",
-    "    # # optionally load a model\n",
-    "    # if config.resume:\n",
-    "    #     ddpm.load_state_dict(torch.load(os.path.join(config.save_dir, f\"train-{ep}xscale_test_{run_name}.npy\")))\n",
-    "\n",
-    "    # define the loss function\n",
-    "    loss_mse = nn.MSELoss()\n",
-    "\n",
-    "\n",
-    "    # initialize optimizer\n",
-    "    # optim = torch.optim.Adam(params_to_optimize, lr=config.lrate)\n",
-    "\n",
-    "    # whether to use ema\n",
-    "    if config.ema:\n",
-    "        ema = EMA(config.ema_rate)\n",
-    "        if config.resume:\n",
-    "            print(\"resuming ema_model\")\n",
-    "            # ema_model = DDPMScheduler(betas=(1e-4, 0.02), num_timesteps=config.num_timesteps, device=config.device)\n",
-    "            ema_model = ContextUnet(n_param=config.n_param, image_size=config.HII_DIM).to(config.device)\n",
-    "            # print(\"ema_model.device =\", ema_model.device)\n",
-    "            ema_model.load_state_dict(torch.load(os.path.join(config.output_dir, f\"model_state.pth\"))['ema_unet_state_dict'])\n",
-    "            # ema_model.load_state_dict(torch.load(os.path.join(config.output_dir, f\"train-{ep}xscale_test_{config.run_name}_ema.npy\")))\n",
-    "        else:\n",
-    "            ema_model = copy.deepcopy(nn_model).eval().requires_grad_(False)\n",
-    "\n",
-    "    ###################      \n",
-    "    ## training loop ##\n",
-    "    ###################\n",
-    "    # plot_unet = True\n",
-    "    global_step = 0\n",
-    "    for ep in range(config.n_epoch):\n",
-    "        # print(f'epoch {ep}')\n",
-    "        # print(\"ddpm.train()\")\n",
-    "        ddpm.train()\n",
-    "        # linear lrate decay\n",
-    "        # if config.lr_decay:\n",
-    "        #     optim.param_groups[0]['lr'] = config.lrate*(1-ep/config.n_epoch)\n",
-    "\n",
-    "        # data loader with progress bar\n",
-    "        pbar_train = tqdm(total=len(dataloader), disable=not accelerator.is_local_main_process)\n",
-    "        pbar_train.set_description(f\"Epoch {ep}\")\n",
-    "        for i, (x, c) in enumerate(dataloader):\n",
-    "            # global_step = ep * len(dataloader) + i\n",
-    "            with accelerator.accumulate(nn_model):\n",
-    "                # optim.zero_grad()\n",
-    "                x = x.to(config.device)\n",
-    "                xt, noise, ts = ddpm.add_noise(x)\n",
-    "\n",
-    "            # noise = torch.randn(x.shape, device=x.device)\n",
-    "            # ts = torch.randint(0, num_timesteps, (x.shape[0],), device=x.device, dtype=torch.int64)\n",
-    "            # xt = ddpm.add_noise(x, noise, ts)\n",
+    "#     # initialize the DDPM\n",
+    "#     # logger = SummaryWriter(os.path.join(\"runs\", config.run_name)) # To log\n",
+    "\n",
+    "#     # ddpm = DDPMScheduler(betas=(1e-4, 0.02), num_timesteps=config.num_timesteps, device=config.device)\n",
+    "\n",
+    "#     # # initialize the unet\n",
+    "#     # nn_model = ContextUnet(n_param=config.n_param, image_size=config.HII_DIM)\n",
+    "#     # # nn_model = ContextUnet(n_param=1, image_size=28)\n",
+    "#     # nn_model.train()\n",
+    "#     # nn_model.to(ddpm.device)\n",
+    "\n",
+    "#     # parameters to be optimized\n",
+    "#     # params_to_optimize = [\n",
+    "#     #     {'params': nn_model.parameters()}\n",
+    "#     # ]\n",
+    "\n",
+    "#     # number of parameters to be trained\n",
+    "#     number_of_params = sum(x.numel() for x in nn_model.parameters())\n",
+    "#     print(f\"Number of parameters for unet: {number_of_params}\")\n",
+    "\n",
+    "#     # # optionally load a model\n",
+    "#     # if config.resume:\n",
+    "#     #     ddpm.load_state_dict(torch.load(os.path.join(config.save_dir, f\"train-{ep}xscale_test_{run_name}.npy\")))\n",
+    "\n",
+    "#     # define the loss function\n",
+    "#     loss_mse = nn.MSELoss()\n",
+    "\n",
+    "\n",
+    "#     # initialize optimizer\n",
+    "#     # optim = torch.optim.Adam(params_to_optimize, lr=config.lrate)\n",
+    "\n",
+    "#     # whether to use ema\n",
+    "#     if config.ema:\n",
+    "#         ema = EMA(config.ema_rate)\n",
+    "#         if config.resume:\n",
+    "#             print(\"resuming ema_model\")\n",
+    "#             # ema_model = DDPMScheduler(betas=(1e-4, 0.02), num_timesteps=config.num_timesteps, device=config.device)\n",
+    "#             ema_model = ContextUnet(n_param=config.n_param, image_size=config.HII_DIM).to(config.device)\n",
+    "#             # print(\"ema_model.device =\", ema_model.device)\n",
+    "#             ema_model.load_state_dict(torch.load(os.path.join(config.output_dir, f\"model_state.pth\"))['ema_unet_state_dict'])\n",
+    "#             # ema_model.load_state_dict(torch.load(os.path.join(config.output_dir, f\"train-{ep}xscale_test_{config.run_name}_ema.npy\")))\n",
+    "#         else:\n",
+    "#             ema_model = copy.deepcopy(nn_model).eval().requires_grad_(False)\n",
+    "\n",
+    "#     ###################      \n",
+    "#     ## training loop ##\n",
+    "#     ###################\n",
+    "#     # plot_unet = True\n",
+    "#     global_step = 0\n",
+    "#     for ep in range(config.n_epoch):\n",
+    "#         # print(f'epoch {ep}')\n",
+    "#         # print(\"ddpm.train()\")\n",
+    "#         ddpm.train()\n",
+    "#         # linear lrate decay\n",
+    "#         # if config.lr_decay:\n",
+    "#         #     optim.param_groups[0]['lr'] = config.lrate*(1-ep/config.n_epoch)\n",
+    "\n",
+    "#         # data loader with progress bar\n",
+    "#         pbar_train = tqdm(total=len(dataloader), disable=not accelerator.is_local_main_process)\n",
+    "#         pbar_train.set_description(f\"Epoch {ep}\")\n",
+    "#         for i, (x, c) in enumerate(dataloader):\n",
+    "#             # global_step = ep * len(dataloader) + i\n",
+    "#             with accelerator.accumulate(nn_model):\n",
+    "#                 # optim.zero_grad()\n",
+    "#                 x = x.to(config.device)\n",
+    "#                 xt, noise, ts = ddpm.add_noise(x)\n",
+    "\n",
+    "#             # noise = torch.randn(x.shape, device=x.device)\n",
+    "#             # ts = torch.randint(0, num_timesteps, (x.shape[0],), device=x.device, dtype=torch.int64)\n",
+    "#             # xt = ddpm.add_noise(x, noise, ts)\n",
     "                \n",
-    "                if config.guide_w == -1:\n",
-    "                    # noise_pred = nn_model(xt, ts, return_dict=False)[0]\n",
-    "                    noise_pred = nn_model(xt, ts)\n",
-    "                else:\n",
-    "                    c = c.to(config.device)\n",
-    "                    noise_pred = nn_model(xt, ts, c)\n",
+    "#                 if config.guide_w == -1:\n",
+    "#                     # noise_pred = nn_model(xt, ts, return_dict=False)[0]\n",
+    "#                     noise_pred = nn_model(xt, ts)\n",
+    "#                 else:\n",
+    "#                     c = c.to(config.device)\n",
+    "#                     noise_pred = nn_model(xt, ts, c)\n",
     "                \n",
-    "                loss = loss_mse(noise, noise_pred)\n",
-    "                accelerator.backward(loss)\n",
-    "                # loss.backward()\n",
-    "                # optim.step()\n",
-    "                accelerator.clip_grad_norm_(nn_model.parameters(), 1)\n",
-    "                optimizer.step()\n",
-    "                lr_scheduler.step()\n",
-    "                optimizer.zero_grad()\n",
-    "\n",
-    "            # ema update\n",
-    "            if config.ema:\n",
-    "                ema.step_ema(ema_model, nn_model)\n",
-    "\n",
-    "            # pbar.set_description(f\"epoch {ep} loss {loss.item():.4f}\")\n",
-    "            pbar_train.update(1)\n",
-    "            logs = dict(\n",
-    "                loss=loss.detach().item(),\n",
-    "                lr=optimizer.param_groups[0]['lr'],\n",
-    "                step=global_step\n",
-    "            )\n",
-    "            pbar_train.set_postfix(**logs)\n",
-    "\n",
-    "            # logging loss\n",
-    "            # logger.add_scalar(\"MSE\", loss.item(), global_step=global_step)\n",
-    "            accelerator.log(logs, step=global_step)\n",
-    "            global_step += 1\n",
-    "\n",
-    "\n",
-    "        if accelerator.is_main_process:\n",
-    "            # sample the image\n",
-    "            if ep == config.n_epoch-1 or (ep+1)*config.save_freq==1:\n",
-    "                nn_model.eval()\n",
-    "                with torch.no_grad():\n",
-    "                    # save model\n",
-    "                    if config.push_to_hub:\n",
-    "                        upload_folder(\n",
-    "                            repo_id = repo_id,\n",
-    "                            folder_path = \".\",#config.output_dir,\n",
-    "                            commit_message = f\"{config.date}\",\n",
-    "                            ignore_patterns = [\"step_*\", \"epoch_*\", \"*.npy\"],\n",
-    "                            )\n",
-    "                    if config.save_model:\n",
-    "                        model_state = {\n",
-    "                            'epoch': ep,\n",
-    "                            'unet_state_dict': nn_model.state_dict(),\n",
-    "                            'ema_unet_state_dict': ema_model.state_dict(),\n",
-    "                            }\n",
-    "                        torch.save(model_state, config.output_dir + f\"model_state.pth\")\n",
-    "                        print('saved model at ' + config.output_dir + f\"model_state.pth\")\n",
-    "                        # print('saved model at ' + config.save_dir + f\"model_epoch_{ep}_test_{config.run_name}.pth\")\n",
-    "\n",
-    "                    # loop over the guidance scale\n",
-    "                    # for w in config.ws_test: \n",
+    "#                 loss = loss_mse(noise, noise_pred)\n",
+    "#                 accelerator.backward(loss)\n",
+    "#                 # loss.backward()\n",
+    "#                 # optim.step()\n",
+    "#                 accelerator.clip_grad_norm_(nn_model.parameters(), 1)\n",
+    "#                 optimizer.step()\n",
+    "#                 lr_scheduler.step()\n",
+    "#                 optimizer.zero_grad()\n",
+    "\n",
+    "#             # ema update\n",
+    "#             if config.ema:\n",
+    "#                 ema.step_ema(ema_model, nn_model)\n",
+    "\n",
+    "#             # pbar.set_description(f\"epoch {ep} loss {loss.item():.4f}\")\n",
+    "#             pbar_train.update(1)\n",
+    "#             logs = dict(\n",
+    "#                 loss=loss.detach().item(),\n",
+    "#                 lr=optimizer.param_groups[0]['lr'],\n",
+    "#                 step=global_step\n",
+    "#             )\n",
+    "#             pbar_train.set_postfix(**logs)\n",
+    "\n",
+    "#             # logging loss\n",
+    "#             # logger.add_scalar(\"MSE\", loss.item(), global_step=global_step)\n",
+    "#             accelerator.log(logs, step=global_step)\n",
+    "#             global_step += 1\n",
+    "\n",
+    "\n",
+    "#         if accelerator.is_main_process:\n",
+    "#             # sample the image\n",
+    "#             if ep == config.n_epoch-1 or (ep+1)*config.save_freq==1:\n",
+    "#                 nn_model.eval()\n",
+    "#                 with torch.no_grad():\n",
+    "#                     # save model\n",
+    "#                     if config.push_to_hub:\n",
+    "#                         upload_folder(\n",
+    "#                             repo_id = repo_id,\n",
+    "#                             folder_path = \".\",#config.output_dir,\n",
+    "#                             commit_message = f\"{config.date}\",\n",
+    "#                             ignore_patterns = [\"step_*\", \"epoch_*\", \"*.npy\"],\n",
+    "#                             )\n",
+    "#                     if config.save_model:\n",
+    "#                         model_state = {\n",
+    "#                             'epoch': ep,\n",
+    "#                             'unet_state_dict': nn_model.state_dict(),\n",
+    "#                             'ema_unet_state_dict': ema_model.state_dict(),\n",
+    "#                             }\n",
+    "#                         torch.save(model_state, config.output_dir + f\"model_state.pth\")\n",
+    "#                         print('saved model at ' + config.output_dir + f\"model_state.pth\")\n",
+    "#                         # print('saved model at ' + config.save_dir + f\"model_epoch_{ep}_test_{config.run_name}.pth\")\n",
+    "\n",
+    "#                     # loop over the guidance scale\n",
+    "#                     # for w in config.ws_test: \n",
     "                    \n",
-    "                    # pipeline = DDPMPipeline(unet=nn_model, scheduler=ddpm)\n",
-    "                    # evaluate(config, ep, pipeline)\n",
+    "#                     # pipeline = DDPMPipeline(unet=nn_model, scheduler=ddpm)\n",
+    "#                     # evaluate(config, ep, pipeline)\n",
     "\n",
-    "                    # only output the image x0, omit the stored intermediate steps, OTHERWISE, uncomment \n",
-    "                    # line 142, 143 and output 'x_last, x_store = ' here.\n",
+    "#                     # only output the image x0, omit the stored intermediate steps, OTHERWISE, uncomment \n",
+    "#                     # line 142, 143 and output 'x_last, x_store = ' here.\n",
     "\n",
-    "                    # x_last_tot = []\n",
-    "                    x_last, x_entire = ddpm.sample(nn_model,config.n_sample, x.shape[1:], config.device, test_param=config.test_param, guide_w=config.guide_w)\n",
+    "#                     # x_last_tot = []\n",
+    "#                     x_last, x_entire = ddpm.sample(nn_model,config.n_sample, x.shape[1:], config.device, params=config.params, guide_w=config.guide_w)\n",
     "\n",
-    "                    # sample_save_dir = os.path.join(config.save_dir, f\"{config.run_name}.npy\")\n",
-    "                    np.save(os.path.join(config.output_dir, f\"{config.run_name}.npy\"), x_last)\n",
-    "                    # np.save(os.path.join(config.save_dir, f\"{config.run_name}_entire.npy\"), x_entire)\n",
-    "                    # print(f\"saved to {config.save_dir}\")\n",
+    "#                     # sample_save_dir = os.path.join(config.save_dir, f\"{config.run_name}.npy\")\n",
+    "#                     np.save(os.path.join(config.output_dir, f\"{config.run_name}.npy\"), x_last)\n",
+    "#                     # np.save(os.path.join(config.save_dir, f\"{config.run_name}_entire.npy\"), x_entire)\n",
+    "#                     # print(f\"saved to {config.save_dir}\")\n",
     "\n",
-    "                    if config.ema:\n",
-    "                        # x_last_tot_ema = []\n",
-    "                        x_last_ema, x_entire_ema = ddpm.sample(ema_model,config.n_sample, x.shape[1:], config.device, test_param=config.test_param, guide_w=config.guide_w)\n",
+    "#                     if config.ema:\n",
+    "#                         # x_last_tot_ema = []\n",
+    "#                         x_last_ema, x_entire_ema = ddpm.sample(ema_model,config.n_sample, x.shape[1:], config.device, params=config.params, guide_w=config.guide_w)\n",
     "\n",
-    "                        np.save(os.path.join(config.output_dir, f\"{config.run_name}_ema.npy\"), x_last_ema)\n",
-    "                        # np.save(os.path.join(config.save_dir, f\"{config.run_name}_ema_entire.npy\"), x_entire_ema)\n",
-    "                        # print(f\"saved to {config.save_dir}\")\n",
+    "#                         np.save(os.path.join(config.output_dir, f\"{config.run_name}_ema.npy\"), x_last_ema)\n",
+    "#                         # np.save(os.path.join(config.save_dir, f\"{config.run_name}_ema_entire.npy\"), x_entire_ema)\n",
+    "#                         # print(f\"saved to {config.save_dir}\")\n",
     "\n",
-    "                    # x_last_tot.append(np.array(x_last.cpu()))\n",
-    "                    # x_last_tot=np.array(x_last_tot)\n",
-    "                    # x_last_tot_ema.append(np.array(x_last_ema.cpu()))\n",
-    "                    # x_last_tot_ema=np.array(x_last_tot_ema)\n",
+    "#                     # x_last_tot.append(np.array(x_last.cpu()))\n",
+    "#                     # x_last_tot=np.array(x_last_tot)\n",
+    "#                     # x_last_tot_ema.append(np.array(x_last_ema.cpu()))\n",
+    "#                     # x_last_tot_ema=np.array(x_last_tot_ema)\n",
     "\n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 67,
+   "execution_count": 109,
    "metadata": {},
    "outputs": [
     {
@@ -1261,8 +1263,8 @@
       "loading 20 images randomly\n",
       "images loaded: (20, 1, 64, 512)\n",
       "params loaded: (20, 2)\n",
-      "images rescaled to [-1.0, 0.8148523569107056]\n",
-      "params rescaled to [0.0, 0.9062639012309924]\n",
+      "images rescaled to [-1.0, 0.946401834487915]\n",
+      "params rescaled to [0.0, 0.9683106587269014]\n",
       "resumed nn_model from model_state.pth\n",
       "Number of parameters for nn_model: 111048705\n"
      ]
@@ -1289,7 +1291,7 @@
     "        self.config = config\n",
     "\n",
     "        dataset = Dataset4h5(config.dataset_name, num_image=config.num_image, HII_DIM=config.HII_DIM, num_redshift=config.num_redshift, drop_prob=config.drop_prob, dim=config.dim)\n",
-    "        self.shape_loaded = dataset.images.shape\n",
+    "        # self.shape_loaded = dataset.images.shape\n",
     "        # print(\"shape_loaded =\", self.shape_loaded)\n",
     "        self.dataloader = DataLoader(dataset, batch_size=config.train_batch_size, shuffle=True)\n",
     "        del dataset\n",
@@ -1420,10 +1422,12 @@
     "                        print('saved model at ' + self.config.output_dir + f\"model_state_{ep:02d}.pth\")\n",
     "                        # print('saved model at ' + config.save_dir + f\"model_epoch_{ep}_test_{config.run_name}.pth\")\n",
     "\n",
-    "    def sample(self, file, n_sample=12, ema=False, entire=False):\n",
+    "    def sample(self, file, params=[0.2,0.8], ema=False, entire=False):\n",
+    "        n_sample = params.shape[0]\n",
     "        model = self.ema_model if ema else self.nn_model\n",
+    "        # params = torch.tile(params, (n_sample,1)).to(device)\n",
     "\n",
-    "        x_last, x_entire = self.ddpm.sample(model, n_sample, self.shape_loaded[1:], self.config.device, test_param=self.config.test_param, guide_w=self.config.guide_w)\n",
+    "        x_last, x_entire = self.ddpm.sample(model, n_sample, self.config.device, params=params, guide_w=self.config.guide_w)\n",
     "\n",
     "        np.save(os.path.join(self.config.output_dir, f\"{self.config.run_name}{'ema' if ema else None}.npy\"), x_last)\n",
     "        if entire:\n",
@@ -1436,13 +1440,27 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 68,
+   "execution_count": 110,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "4874b29272224f1aa6bcbead8dc5d11f",
+       "model_id": "7b26c7444a1144728f0299db5d5683b1",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e3b8c1a18460443d986282f284bf0b42",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1456,7 +1474,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1caa087a9a7b4252a633054990ff76d8",
+       "model_id": "84d558fcea2f4b998f311da7452a2c93",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1470,7 +1488,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "81644fce614d49f5aa63e291ec458ccf",
+       "model_id": "364208bd1fa04859baf233ed30451638",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1484,7 +1502,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a011d7ece27e42128a6ec51227313e60",
+       "model_id": "e8aa693ca9b444d7b6adb77531f8b718",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1498,7 +1516,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6d2f56d5d27443e589a8cca5d45892e3",
+       "model_id": "aab44a74df8c4a7a8baf975b8ec50b8b",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1516,11 +1534,44 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 69,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "89a5be983ade43d89a1be3d977750a40",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/1000 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "ename": "RuntimeError",
+     "evalue": "The size of tensor a (24) must match the size of tensor b (36) at non-singleton dimension 0",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[69], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m ddpm21cm\u001b[39m.\u001b[39;49msample(\u001b[39m\"\u001b[39;49m\u001b[39m./outputs/model_state_09.pth\u001b[39;49m\u001b[39m\"\u001b[39;49m)\n",
+      "Cell \u001b[0;32mIn[67], line 141\u001b[0m, in \u001b[0;36mDDPM21CM.sample\u001b[0;34m(self, file, n_sample, ema, entire)\u001b[0m\n\u001b[1;32m    138\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39msample\u001b[39m(\u001b[39mself\u001b[39m, file, n_sample\u001b[39m=\u001b[39m\u001b[39m12\u001b[39m, ema\u001b[39m=\u001b[39m\u001b[39mFalse\u001b[39;00m, entire\u001b[39m=\u001b[39m\u001b[39mFalse\u001b[39;00m):\n\u001b[1;32m    139\u001b[0m     model \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mema_model \u001b[39mif\u001b[39;00m ema \u001b[39melse\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mnn_model\n\u001b[0;32m--> 141\u001b[0m     x_last, x_entire \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mddpm\u001b[39m.\u001b[39;49msample(model, n_sample, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mshape_loaded[\u001b[39m1\u001b[39;49m:], \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mconfig\u001b[39m.\u001b[39;49mdevice, test_param\u001b[39m=\u001b[39;49m\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mconfig\u001b[39m.\u001b[39;49mtest_param, guide_w\u001b[39m=\u001b[39;49m\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mconfig\u001b[39m.\u001b[39;49mguide_w)\n\u001b[1;32m    143\u001b[0m     np\u001b[39m.\u001b[39msave(os\u001b[39m.\u001b[39mpath\u001b[39m.\u001b[39mjoin(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mconfig\u001b[39m.\u001b[39moutput_dir, \u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39m{\u001b[39;00m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mconfig\u001b[39m.\u001b[39mrun_name\u001b[39m}\u001b[39;00m\u001b[39m{\u001b[39;00m\u001b[39m'\u001b[39m\u001b[39mema\u001b[39m\u001b[39m'\u001b[39m\u001b[39m \u001b[39m\u001b[39mif\u001b[39;00m\u001b[39m \u001b[39mema\u001b[39m \u001b[39m\u001b[39melse\u001b[39;00m\u001b[39m \u001b[39m\u001b[39mNone\u001b[39;00m\u001b[39m}\u001b[39;00m\u001b[39m.npy\u001b[39m\u001b[39m\"\u001b[39m), x_last)\n\u001b[1;32m    144\u001b[0m     \u001b[39mif\u001b[39;00m entire:\n",
+      "Cell \u001b[0;32mIn[7], line 70\u001b[0m, in \u001b[0;36mDDPMScheduler.sample\u001b[0;34m(self, nn_model, n_sample, shape, device, test_param, guide_w)\u001b[0m\n\u001b[1;32m     67\u001b[0m t_is \u001b[39m=\u001b[39m t_is\u001b[39m.\u001b[39mrepeat(\u001b[39m2\u001b[39m)\n\u001b[1;32m     69\u001b[0m \u001b[39m# split predictions and compute weighting\u001b[39;00m\n\u001b[0;32m---> 70\u001b[0m eps \u001b[39m=\u001b[39m nn_model(x_i, t_is, c_i)\n\u001b[1;32m     71\u001b[0m eps1 \u001b[39m=\u001b[39m eps[:n_sample]\n\u001b[1;32m     72\u001b[0m eps2 \u001b[39m=\u001b[39m eps[n_sample:]\n",
+      "File \u001b[0;32m/usr/local/pace-apps/manual/packages/pytorch/1.12.0/lib/python3.9/site-packages/torch/nn/modules/module.py:1130\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1126\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1127\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1128\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49m\u001b[39minput\u001b[39;49m, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1131\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1132\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
+      "File \u001b[0;32m~/.conda/envs/diffusers/lib/python3.9/site-packages/accelerate/utils/operations.py:822\u001b[0m, in \u001b[0;36mconvert_outputs_to_fp32.<locals>.forward\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    821\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mforward\u001b[39m(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs):\n\u001b[0;32m--> 822\u001b[0m     \u001b[39mreturn\u001b[39;00m model_forward(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
+      "File \u001b[0;32m~/.conda/envs/diffusers/lib/python3.9/site-packages/accelerate/utils/operations.py:810\u001b[0m, in \u001b[0;36mConvertOutputsToFp32.__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    809\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m__call__\u001b[39m(\u001b[39mself\u001b[39m, \u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs):\n\u001b[0;32m--> 810\u001b[0m     \u001b[39mreturn\u001b[39;00m convert_to_fp32(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mmodel_forward(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs))\n",
+      "File \u001b[0;32m/usr/local/pace-apps/manual/packages/pytorch/1.12.0/lib/python3.9/site-packages/torch/amp/autocast_mode.py:12\u001b[0m, in \u001b[0;36mautocast_decorator.<locals>.decorate_autocast\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m      9\u001b[0m \u001b[39m@functools\u001b[39m\u001b[39m.\u001b[39mwraps(func)\n\u001b[1;32m     10\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mdecorate_autocast\u001b[39m(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs):\n\u001b[1;32m     11\u001b[0m     \u001b[39mwith\u001b[39;00m autocast_instance:\n\u001b[0;32m---> 12\u001b[0m         \u001b[39mreturn\u001b[39;00m func(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
+      "Cell \u001b[0;32mIn[18], line 211\u001b[0m, in \u001b[0;36mContextUnet.forward\u001b[0;34m(self, x, timesteps, y)\u001b[0m\n\u001b[1;32m    209\u001b[0m \u001b[39mif\u001b[39;00m y \u001b[39m!=\u001b[39m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m    210\u001b[0m     text_outputs \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mtoken_embedding(y\u001b[39m.\u001b[39mfloat())\n\u001b[0;32m--> 211\u001b[0m     emb \u001b[39m=\u001b[39m emb \u001b[39m+\u001b[39;49m text_outputs\u001b[39m.\u001b[39;49mto(emb)\n\u001b[1;32m    213\u001b[0m h \u001b[39m=\u001b[39m x\u001b[39m.\u001b[39mtype(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdtype)\n\u001b[1;32m    214\u001b[0m \u001b[39m# print(\"0,h.shape =\", h.shape)\u001b[39;00m\n",
+      "\u001b[0;31mRuntimeError\u001b[0m: The size of tensor a (24) must match the size of tensor b (36) at non-singleton dimension 0"
+     ]
+    }
+   ],
    "source": [
-    "ddpm21cm.sample()"
+    "ddpm21cm.sample(\"./outputs/model_state_09.pth\")"
    ]
   },
   {
@@ -1718,7 +1769,7 @@
     "\n",
     "n_sample = 20\n",
     "with torch.no_grad():\n",
-    "    x_last_ema, x_ema_entire = ddpm.sample(nn_model, n_sample, (1,config.HII_DIM, config.num_redshift), config.device, test_param = torch.tile(config.test_param_single,(n_sample,1)).to(config.device), guide_w=config.guide_w)\n",
+    "    x_last_ema, x_ema_entire = ddpm.sample(nn_model, n_sample, (1,config.HII_DIM, config.num_redshift), config.device, params = torch.tile(config.params_single,(n_sample,1)).to(config.device), guide_w=config.guide_w)\n",
     "\n",
     "np.save(os.path.join(config.output_dir, f\"{config.run_name}_ema.npy\"), x_last_ema)"
    ]