0521-2304

diffusion.ipynb

@@ -32,7 +32,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8be92a01e78a47b792d93b35d557885d",
+       "model_id": "f925cb378800455fb1216e84a2900e58",
        "version_major": 2,
        "version_minor": 0
       },
@@ -962,7 +962,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 29,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -987,7 +987,7 @@
     "\n",
     "    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",
+    "    batch_size = 10#10#20#2#100 # 10\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",
@@ -1032,7 +1032,7 @@
    "outputs": [],
    "source": [
     "# 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",
-    "# dataloader = DataLoader(dataset, batch_size=config.train_batch_size, shuffle=True)"
+    "# dataloader = DataLoader(dataset, batch_size=config.batch_size, shuffle=True)"
    ]
   },
   {
@@ -1083,10 +1083,10 @@
     "#     # 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",
+    "#     # dataloader = DataLoader(dataset, batch_size=config.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",
+    "#     # dataloader = DataLoader(dataset, batch_size=config.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",
@@ -1268,37 +1268,15 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 25,
+   "execution_count": 31,
    "metadata": {},
    "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 200 images randomly\n",
-      "images loaded: (200, 1, 64, 512)\n",
-      "params loaded: (200, 2)\n",
-      "images rescaled to [-1.0, 1.082756519317627]\n",
-      "params rescaled to [0.0, 0.9938162632551855]\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": [
+      "Number of parameters for nn_model: 111048705\n",
       "resumed ema_model from model_state.pth\n"
      ]
     }
@@ -1306,15 +1284,14 @@
    "source": [
     "# @dataclass\n",
     "class DDPM21CM:\n",
-    "    def __init__(self, config):\n",
+    "    def __init__(self):\n",
+    "        config = TrainConfig()\n",
     "        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",
+    "        # 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.batch_size, shuffle=True)\n",
+    "        # del dataset\n",
     "        self.ddpm = DDPMScheduler(betas=(1e-4, 0.02), num_timesteps=config.num_timesteps, img_shape=config.img_shape, device=config.device)\n",
     "\n",
     "        # initialize the unet\n",
@@ -1345,10 +1322,17 @@
     "        self.lr_scheduler = get_cosine_schedule_with_warmup(\n",
     "            optimizer=self.optimizer,\n",
     "            num_warmup_steps=config.lr_warmup_steps,\n",
-    "            num_training_steps=(len(self.dataloader) * config.n_epoch),\n",
+    "            num_training_steps=(int(config.num_image/config.batch_size) * config.n_epoch),\n",
+    "            # num_training_steps=(len(self.dataloader) * config.n_epoch),\n",
     "        )\n",
     "\n",
-    "        self.accelerate(config)\n",
+    "    def load(self):\n",
+    "        dataset = Dataset4h5(self.config.dataset_name, num_image=self.config.num_image, HII_DIM=self.config.HII_DIM, num_redshift=self.config.num_redshift, drop_prob=self.config.drop_prob, dim=self.config.dim)\n",
+    "        # self.shape_loaded = dataset.images.shape\n",
+    "        # print(\"shape_loaded =\", self.shape_loaded)\n",
+    "        self.dataloader = DataLoader(dataset, batch_size=self.config.batch_size, shuffle=True)\n",
+    "        # del dataset\n",
+    "        self.accelerate(self.config)\n",
     "\n",
     "    def accelerate(self, config):\n",
     "        self.accelerator = Accelerator(\n",
@@ -1376,6 +1360,7 @@
     "        ## training loop ##\n",
     "        ###################\n",
     "        # plot_unet = True\n",
+    "        self.load()\n",
     "        global_step = 0\n",
     "        for ep in range(self.config.n_epoch):\n",
     "            self.ddpm.train()\n",
@@ -1446,26 +1431,54 @@
     "        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, params=params, shape=shape, device=self.config.device, guide_w=self.config.guide_w)\n",
+    "        x_last, x_entire = self.ddpm.sample(model, params=params, device=self.config.device, 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",
+    "ddpm21cm = DDPM21CM()\n",
     "# print(\"device =\", config.device)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 26,
+   "execution_count": 32,
    "metadata": {},
    "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 200 images randomly\n",
+      "images loaded: (200, 1, 64, 512)\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": [
+      "params loaded: (200, 2)\n",
+      "images rescaled to [-1.0, 1.095154047012329]\n",
+      "params rescaled to [0.0, 0.997810682944812]\n"
+     ]
+    },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7a0b627f28ef409f8504113bc3af36e3",
+       "model_id": "e6947686dc1344b4b446910cca5326dc",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1479,7 +1492,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "62f09cd440a84841b336ab15e76e2fe6",
+       "model_id": "90757bce414c44778e816eb75fbeb106",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1493,7 +1506,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9db24e29de0c47328f1aba68db61bbae",
+       "model_id": "cc51568c1fb94c82a96726973b103b36",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1507,7 +1520,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ee59d1a664d04a2b90a7a448a816ed10",
+       "model_id": "62c75108d8344ec69b34873181e3d308",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1521,7 +1534,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8690c736f7eb4a23925b450c05659575",
+       "model_id": "6bece94eb3cd45c7933bae67ad774f06",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1535,7 +1548,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7dc014a33bfd43408e0aafc208bb403e",
+       "model_id": "d3c093174fab4863af65f477f8baee8f",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1549,7 +1562,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6715e5cccc6d480397f76bcea34f94e5",
+       "model_id": "f2565e70177a402ba90215f4d80f717e",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1563,7 +1576,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b7410efd4a5d4efdb9b8be38ba1c2fcb",
+       "model_id": "f11d5a63fb1040c2a4366207e53476b7",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1577,7 +1590,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "3b6c0478c9ff4a99b7f79ba4422dbd7d",
+       "model_id": "4641851ffed541edb4bfd99d156bf000",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1591,7 +1604,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d26f49f5a9804d84b6b6a531a56eb03a",
+       "model_id": "9338bea79f5d49c4b79a2327b9833954",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1609,28 +1622,20 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 29,
    "metadata": {},
    "outputs": [
     {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "params.shape[0], len(params) 2 2\n"
-     ]
-    },
-    {
-     "ename": "AttributeError",
-     "evalue": "'DDPMScheduler' object has no attribute 'shape'",
+     "ename": "IndexError",
+     "evalue": "tuple index out of range",
      "output_type": "error",
      "traceback": [
       "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[116], 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[115], line 143\u001b[0m, in \u001b[0;36mDDPM21CM.sample\u001b[0;34m(self, file, params, ema, entire)\u001b[0m\n\u001b[1;32m    140\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[1;32m    141\u001b[0m \u001b[39m# params = torch.tile(params, (n_sample,1)).to(device)\u001b[39;00m\n\u001b[0;32m--> 143\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, params\u001b[39m=\u001b[39;49mparams, device\u001b[39m=\u001b[39;49m\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mconfig\u001b[39m.\u001b[39;49mdevice, 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    145\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    146\u001b[0m \u001b[39mif\u001b[39;00m entire:\n",
-      "Cell \u001b[0;32mIn[90], line 40\u001b[0m, in \u001b[0;36mDDPMScheduler.sample\u001b[0;34m(self, nn_model, params, device, guide_w)\u001b[0m\n\u001b[1;32m     38\u001b[0m n_sample \u001b[39m=\u001b[39m params\u001b[39m.\u001b[39mshape[\u001b[39m0\u001b[39m]\n\u001b[1;32m     39\u001b[0m \u001b[39mprint\u001b[39m(\u001b[39m\"\u001b[39m\u001b[39mparams.shape[0], len(params)\u001b[39m\u001b[39m\"\u001b[39m, params\u001b[39m.\u001b[39mshape[\u001b[39m0\u001b[39m], \u001b[39mlen\u001b[39m(params))\n\u001b[0;32m---> 40\u001b[0m x_i \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mrandn(n_sample, \u001b[39m*\u001b[39m\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mshape[\u001b[39m1\u001b[39m:])\u001b[39m.\u001b[39mto(device)\n\u001b[1;32m     41\u001b[0m \u001b[39m# print(\"x_i.shape =\", x_i.shape)\u001b[39;00m\n\u001b[1;32m     42\u001b[0m \u001b[39mif\u001b[39;00m guide_w \u001b[39m!=\u001b[39m \u001b[39m-\u001b[39m\u001b[39m1\u001b[39m:\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:1207\u001b[0m, in \u001b[0;36mModule.__getattr__\u001b[0;34m(self, name)\u001b[0m\n\u001b[1;32m   1205\u001b[0m     \u001b[39mif\u001b[39;00m name \u001b[39min\u001b[39;00m modules:\n\u001b[1;32m   1206\u001b[0m         \u001b[39mreturn\u001b[39;00m modules[name]\n\u001b[0;32m-> 1207\u001b[0m \u001b[39mraise\u001b[39;00m \u001b[39mAttributeError\u001b[39;00m(\u001b[39m\"\u001b[39m\u001b[39m'\u001b[39m\u001b[39m{}\u001b[39;00m\u001b[39m'\u001b[39m\u001b[39m object has no attribute \u001b[39m\u001b[39m'\u001b[39m\u001b[39m{}\u001b[39;00m\u001b[39m'\u001b[39m\u001b[39m\"\u001b[39m\u001b[39m.\u001b[39mformat(\n\u001b[1;32m   1208\u001b[0m     \u001b[39mtype\u001b[39m(\u001b[39mself\u001b[39m)\u001b[39m.\u001b[39m\u001b[39m__name__\u001b[39m, name))\n",
-      "\u001b[0;31mAttributeError\u001b[0m: 'DDPMScheduler' object has no attribute 'shape'"
+      "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[29], 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[28], line 143\u001b[0m, in \u001b[0;36mDDPM21CM.sample\u001b[0;34m(self, file, params, ema, entire)\u001b[0m\n\u001b[1;32m    140\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[1;32m    141\u001b[0m \u001b[39m# params = torch.tile(params, (n_sample,1)).to(device)\u001b[39;00m\n\u001b[0;32m--> 143\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, params\u001b[39m=\u001b[39;49mparams, device\u001b[39m=\u001b[39;49m\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mconfig\u001b[39m.\u001b[39;49mdevice, 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    145\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    146\u001b[0m \u001b[39mif\u001b[39;00m entire:\n",
+      "Cell \u001b[0;32mIn[7], line 45\u001b[0m, in \u001b[0;36mDDPMScheduler.sample\u001b[0;34m(self, nn_model, params, device, guide_w)\u001b[0m\n\u001b[1;32m     43\u001b[0m \u001b[39mif\u001b[39;00m guide_w \u001b[39m!=\u001b[39m \u001b[39m-\u001b[39m\u001b[39m1\u001b[39m:\n\u001b[1;32m     44\u001b[0m     c_i \u001b[39m=\u001b[39m params\n\u001b[0;32m---> 45\u001b[0m     uncond_tokens \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mzeros(\u001b[39mint\u001b[39m(n_sample), params\u001b[39m.\u001b[39;49mshape[\u001b[39m1\u001b[39;49m])\u001b[39m.\u001b[39mto(device)\n\u001b[1;32m     46\u001b[0m     \u001b[39m# uncond_tokens = torch.tensor(np.float32(np.array([0,0]))).to(device)\u001b[39;00m\n\u001b[1;32m     47\u001b[0m     \u001b[39m# uncond_tokens = uncond_tokens.repeat(int(n_sample),1)\u001b[39;00m\n\u001b[1;32m     48\u001b[0m     c_i \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mcat((c_i, uncond_tokens), \u001b[39m0\u001b[39m)\n",
+      "\u001b[0;31mIndexError\u001b[0m: tuple index out of range"
      ]
     }
    ],
@@ -2030,7 +2035,7 @@
     "@dataclass\n",
     "class TrainingConfig:\n",
     "    image_size = 128\n",
-    "    train_batch_size = 16\n",
+    "    batch_size = 16\n",
     "    eval_batch_size = 16\n",
     "    num_epochs = 50\n",
     "    gradient_accumulation_steps = 1\n",
@@ -2067,7 +2072,7 @@
     "    return {\"images\": images}\n",
     "\n",
     "dataset.set_transform(transform)\n",
-    "dataloader = torch.utils.data.DataLoader(dataset, batch_size=config.train_batch_size, shuffle=True)\n",
+    "dataloader = torch.utils.data.DataLoader(dataset, batch_size=config.batch_size, shuffle=True)\n",
     "\n",
     "model = UNet2DModel(\n",
     "    sample_size = config.image_size,\n",
@@ -2268,7 +2273,7 @@
     "class TrainingConfig:\n",
     "    num_images = 600\n",
     "    image_size = [64,512]\n",
-    "    train_batch_size = 10\n",
+    "    batch_size = 10\n",
     "    eval_batch_size = 24\n",
     "    num_epochs = 20\n",
     "    gradient_accumulation_steps = 1\n",
@@ -2547,7 +2552,7 @@
    "outputs": [],
    "source": [
     "# dataset.set_transform(transform)\n",
-    "dataloader = torch.utils.data.DataLoader(dataset, batch_size=config.train_batch_size, shuffle=True)"
+    "dataloader = torch.utils.data.DataLoader(dataset, batch_size=config.batch_size, shuffle=True)"
    ]
   },
   {