0521-1651

diffusion.ipynb

@@ -1053,203 +1053,235 @@
    "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, test_param=config.test_param, 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, test_param=config.test_param, 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": null,
+   "execution_count": 67,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "dataset content: <KeysViewHDF5 ['brightness_temp', 'density', 'kwargs', 'params', 'redshifts_distances', 'seeds', 'xH_box']>\n",
+      "51200 images can be loaded\n",
+      "field.shape = (64, 64, 514)\n",
+      "params keys = [b'ION_Tvir_MIN', b'HII_EFF_FACTOR']\n",
+      "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",
+      "resumed nn_model from model_state.pth\n",
+      "Number of parameters for nn_model: 111048705\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Detected kernel version 3.10.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "resumed ema_model from model_state.pth\n"
+     ]
+    }
+   ],
    "source": [
     "# @dataclass\n",
     "class DDPM21CM:\n",
@@ -1257,6 +1289,8 @@
     "        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",
+    "        # print(\"shape_loaded =\", self.shape_loaded)\n",
     "        self.dataloader = DataLoader(dataset, batch_size=config.train_batch_size, shuffle=True)\n",
     "        del dataset\n",
     "\n",
@@ -1386,8 +1420,15 @@
     "                        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):\n",
-    "        pass\n",
+    "    def sample(self, file, n_sample=12, ema=False, entire=False):\n",
+    "        model = self.ema_model if ema else self.nn_model\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",
+    "\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",
+    "            np.save(os.path.join(self.config.output_dir, f\"{self.config.run_name}{'ema' if ema else None}_entire.npy\"), x_last)\n",
+    "\n",
     "\n",
     "ddpm21cm = DDPM21CM(TrainConfig())\n",
     "# print(\"device =\", config.device)"
@@ -1395,13 +1436,13 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 68,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b900cc8b147649d89832371147da66f1",
+       "model_id": "4874b29272224f1aa6bcbead8dc5d11f",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1413,23 +1454,75 @@
      "output_type": "display_data"
     },
     {
-     "ename": "AttributeError",
-     "evalue": "'TrainConfig' object has no attribute 'save_freq'",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[56], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m ddpm21cm\u001b[39m.\u001b[39;49mtrain()\n",
-      "Cell \u001b[0;32mIn[55], line 111\u001b[0m, in \u001b[0;36mDDPM21CM.train\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    108\u001b[0m     global_step \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m \u001b[39m1\u001b[39m\n\u001b[1;32m    110\u001b[0m \u001b[39m# if ep == config.n_epoch-1 or (ep+1)*config.save_freq==1:\u001b[39;00m\n\u001b[0;32m--> 111\u001b[0m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49msave(ep)\n",
-      "Cell \u001b[0;32mIn[55], line 116\u001b[0m, in \u001b[0;36mDDPM21CM.save\u001b[0;34m(self, ep)\u001b[0m\n\u001b[1;32m    113\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39msave\u001b[39m(\u001b[39mself\u001b[39m, ep):\n\u001b[1;32m    114\u001b[0m     \u001b[39m# save model\u001b[39;00m\n\u001b[1;32m    115\u001b[0m     \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39maccelerator\u001b[39m.\u001b[39mis_main_process:\n\u001b[0;32m--> 116\u001b[0m         \u001b[39mif\u001b[39;00m ep \u001b[39m==\u001b[39m config\u001b[39m.\u001b[39mn_epoch\u001b[39m-\u001b[39m\u001b[39m1\u001b[39m \u001b[39mor\u001b[39;00m (ep\u001b[39m+\u001b[39m\u001b[39m1\u001b[39m)\u001b[39m*\u001b[39mconfig\u001b[39m.\u001b[39;49msave_freq\u001b[39m==\u001b[39m\u001b[39m1\u001b[39m:\n\u001b[1;32m    117\u001b[0m             nn_model\u001b[39m.\u001b[39meval()\n\u001b[1;32m    118\u001b[0m             \u001b[39mwith\u001b[39;00m torch\u001b[39m.\u001b[39mno_grad():\n",
-      "\u001b[0;31mAttributeError\u001b[0m: 'TrainConfig' object has no attribute 'save_freq'"
-     ]
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "1caa087a9a7b4252a633054990ff76d8",
+       "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": "81644fce614d49f5aa63e291ec458ccf",
+       "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": "a011d7ece27e42128a6ec51227313e60",
+       "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": "6d2f56d5d27443e589a8cca5d45892e3",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
     }
    ],
    "source": [
     "ddpm21cm.train()"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ddpm21cm.sample()"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 45,