zizi-training-vast.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "21859818-de73-4eb8-b6fb-20bf7a078c80",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from PIL import Image\n",
    "from torchvision import transforms\n",
    "from datasets import load_dataset\n",
    "from diffusers import UNet2DModel\n",
    "from diffusers import DDPMScheduler\n",
    "from diffusers import DDPMPipeline\n",
    "from diffusers.optimization import get_cosine_schedule_with_warmup\n",
    "\n",
    "from accelerate import Accelerator\n",
    "from tqdm.auto import tqdm\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "123f4655-5390-453e-8cdb-73b2140d69cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "\n",
    "@dataclass\n",
    "class TrainingConfig:\n",
    "    image_size = 128  # the generated image resolution\n",
    "    train_batch_size = 16\n",
    "    eval_batch_size = 16  # how many images to sample during evaluation\n",
    "    num_epochs = 50\n",
    "    gradient_accumulation_steps = 1\n",
    "    learning_rate = 1e-4\n",
    "    lr_warmup_steps = 500\n",
    "    save_image_epochs = 10\n",
    "    save_model_epochs = 10\n",
    "    mixed_precision = \"fp16\"  # `no` for float32, `fp16` for automatic mixed precision\n",
    "    output_dir = \"output/zizi-test-128\"  # the model name locally and on the HF Hub\n",
    "    overwrite_output_dir = True  # overwrite the old model when re-running the notebook\n",
    "    seed = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fbc0e594-29c7-4440-9a53-4463daa1a57a",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = TrainingConfig()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "76c27a21-a032-48bb-92c7-eff867605925",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "88a26c9409f442e6bc98b36efd24c633",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Resolving data files:   0%|          | 0/547 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset imagefolder (/root/.cache/huggingface/datasets/imagefolder/default-d61a16f7425c3955/0.0.0/37fbb85cc714a338bea574ac6c7d0b5be5aff46c1862c1989b20e0771199e93f)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "770153229a2848189a1a2147c67e2304",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dataset = load_dataset(\"imagefolder\", data_dir = \"data/zizi/meth-small\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cac9a943-c8e8-4478-bcc9-73d585c24210",
   "metadata": {},
   "outputs": [],
   "source": [
    "preprocess = transforms.Compose(\n",
    "    [\n",
    "        transforms.Resize((config.image_size, config.image_size)),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.5], [0.5])\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3e3e5413-a9ff-4d4d-8784-d0371030deb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def img_dataset_transform(examples):\n",
    "    images = [preprocess(image.convert(\"RGB\")) for image in examples[\"image\"]]\n",
    "    return {\"images\": images}\n",
    "\n",
    "def get_dataloader(dataset):\n",
    "    preprocess = transforms.Compose(\n",
    "        [\n",
    "            transforms.Resize((config.image_size, config.image_size)),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize([0.5], [0.5])\n",
    "        ]\n",
    "    )\n",
    "    dataset.set_transform(img_dataset_transform)\n",
    "    return torch.utils.data.DataLoader(dataset['train'], batch_size=config.train_batch_size, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "de15e56d-3a30-48f3-b17b-9850df66d3ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_model():\n",
    "    return UNet2DModel(\n",
    "        sample_size=config.image_size,  # the target image resolution\n",
    "        in_channels=3,  # the number of input channels, 3 for RGB images\n",
    "        out_channels=3,  # the number of output channels\n",
    "        layers_per_block=2,  # how many ResNet layers to use per UNet block\n",
    "        block_out_channels=(128, 128, 256, 256, 512, 512),  # the number of output channels for each UNet block\n",
    "        down_block_types=(\n",
    "            \"DownBlock2D\",  # a regular ResNet downsampling block\n",
    "            \"DownBlock2D\",\n",
    "            \"DownBlock2D\",\n",
    "            \"DownBlock2D\",\n",
    "            \"AttnDownBlock2D\",  # a ResNet downsampling block with spatial self-attention\n",
    "            \"DownBlock2D\",\n",
    "        ),\n",
    "        up_block_types=(\n",
    "            \"UpBlock2D\",  # a regular ResNet upsampling block\n",
    "            \"AttnUpBlock2D\",  # a ResNet upsampling block with spatial self-attention\n",
    "            \"UpBlock2D\",\n",
    "            \"UpBlock2D\",\n",
    "            \"UpBlock2D\",\n",
    "            \"UpBlock2D\",\n",
    "        ),\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d13d345d-46f2-4341-b303-05e8e7bf2bf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_grid(images, rows, cols):\n",
    "    w, h = images[0].size\n",
    "    grid = Image.new(\"RGB\", size=(cols * w, rows * h))\n",
    "    for i, image in enumerate(images):\n",
    "        grid.paste(image, box=(i % cols * w, i // cols * h))\n",
    "    return grid\n",
    "\n",
    "\n",
    "def evaluate(config, epoch, pipeline):\n",
    "    # Sample some images from random noise (this is the backward diffusion process).\n",
    "    # The default pipeline output type is `List[PIL.Image]`\n",
    "    images = pipeline(\n",
    "        batch_size=config.eval_batch_size,\n",
    "        generator=torch.manual_seed(config.seed),\n",
    "    ).images\n",
    "\n",
    "    # Make a grid out of the images\n",
    "    image_grid = make_grid(images, rows=4, cols=4)\n",
    "\n",
    "    # Save the images\n",
    "    test_dir = os.path.join(config.output_dir, \"samples\")\n",
    "    os.makedirs(test_dir, exist_ok=True)\n",
    "    image_grid.save(f\"{test_dir}/{epoch:04d}.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b9a34231-d278-4244-a42d-e9c555e50191",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataloader = get_dataloader(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "baaccf95-3ac8-45bc-9179-5acad854b304",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d4384ce4-56ef-4e65-a7c0-26a618ac4ac7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_loop(config):\n",
    "    # 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",
    "        os.makedirs(config.output_dir, exist_ok=True)\n",
    "        accelerator.init_trackers(\"train_example\")\n",
    "\n",
    "    dataset = load_dataset(\"imagefolder\", data_dir = \"data/zizi/meth-small\")\n",
    "    train_dataloader = get_dataloader(dataset)\n",
    "\n",
    "    model = get_model()\n",
    "    noise_scheduler = DDPMScheduler(num_train_timesteps=1000)\n",
    "    optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate)\n",
    "    lr_scheduler = get_cosine_schedule_with_warmup(\n",
    "        optimizer=optimizer,\n",
    "        num_warmup_steps=config.lr_warmup_steps,\n",
    "        num_training_steps=(len(train_dataloader) * config.num_epochs),\n",
    "    )\n",
    "\n",
    "    # Prepare everything\n",
    "    # There is no specific order to remember, you just need to unpack the\n",
    "    # objects in the same order you gave them to the prepare method.\n",
    "    model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(\n",
    "        model, optimizer, train_dataloader, lr_scheduler\n",
    "    )\n",
    "\n",
    "    global_step = 0\n",
    "\n",
    "    # Now you train the model\n",
    "    for epoch in range(config.num_epochs):\n",
    "        progress_bar = tqdm(total=len(train_dataloader), disable=not accelerator.is_local_main_process)\n",
    "        progress_bar.set_description(f\"Epoch {epoch}\")\n",
    "\n",
    "        for step, batch in enumerate(train_dataloader):\n",
    "            clean_images = batch[\"images\"]\n",
    "            # Sample noise to add to the images\n",
    "            noise = torch.randn(clean_images.shape).to(clean_images.device)\n",
    "            bs = clean_images.shape[0]\n",
    "\n",
    "            # Sample a random timestep for each image\n",
    "            timesteps = torch.randint(\n",
    "                0, noise_scheduler.config.num_train_timesteps, (bs,), device=clean_images.device\n",
    "            ).long()\n",
    "\n",
    "            # Add noise to the clean images according to the noise magnitude at each timestep\n",
    "            # (this is the forward diffusion process)\n",
    "            noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)\n",
    "\n",
    "            with accelerator.accumulate(model):\n",
    "                # Predict the noise residual\n",
    "                noise_pred = model(noisy_images, timesteps, return_dict=False)[0]\n",
    "                loss = F.mse_loss(noise_pred, noise)\n",
    "                accelerator.backward(loss)\n",
    "\n",
    "                accelerator.clip_grad_norm_(model.parameters(), 1.0)\n",
    "                optimizer.step()\n",
    "                lr_scheduler.step()\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "            progress_bar.update(1)\n",
    "            logs = {\"loss\": loss.detach().item(), \"lr\": lr_scheduler.get_last_lr()[0], \"step\": global_step}\n",
    "            progress_bar.set_postfix(**logs)\n",
    "            accelerator.log(logs, step=global_step)\n",
    "            global_step += 1\n",
    "\n",
    "        # After each epoch you optionally sample some demo images with evaluate() and save the model\n",
    "        if accelerator.is_main_process:\n",
    "            pipeline = DDPMPipeline(unet=accelerator.unwrap_model(model), scheduler=noise_scheduler)\n",
    "\n",
    "            if (epoch + 1) % config.save_image_epochs == 0 or epoch == config.num_epochs - 1:\n",
    "                evaluate(config, epoch, pipeline)\n",
    "\n",
    "            if (epoch + 1) % config.save_model_epochs == 0 or epoch == config.num_epochs - 1:\n",
    "                pipeline.save_pretrained(config.output_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6c57b9c3-19f0-45cc-b56f-f83a6f6ca4ad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Launching training on 2 GPUs.\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "CUDA has been initialized before the `notebook_launcher` could create a forked subprocess. This likely stems from an outside import causing issues once the `notebook_launcher()` is called. Please review your imports and test them when running the `notebook_launcher()` to identify which one is problematic.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mProcessRaisedException\u001b[0m                    Traceback (most recent call last)",
      "File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/accelerate/launchers.py:137\u001b[0m, in \u001b[0;36mnotebook_launcher\u001b[0;34m(function, args, num_processes, mixed_precision, use_port)\u001b[0m\n\u001b[1;32m    136\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 137\u001b[0m     \u001b[43mstart_processes\u001b[49m\u001b[43m(\u001b[49m\u001b[43mlauncher\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnprocs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mnum_processes\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstart_method\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mfork\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m    138\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m ProcessRaisedException \u001b[38;5;28;01mas\u001b[39;00m e:\n",
      "File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/torch/multiprocessing/spawn.py:197\u001b[0m, in \u001b[0;36mstart_processes\u001b[0;34m(fn, args, nprocs, join, daemon, start_method)\u001b[0m\n\u001b[1;32m    196\u001b[0m \u001b[38;5;66;03m# Loop on join until it returns True or raises an exception.\u001b[39;00m\n\u001b[0;32m--> 197\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[43mcontext\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mjoin\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m:\n\u001b[1;32m    198\u001b[0m     \u001b[38;5;28;01mpass\u001b[39;00m\n",
      "File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/torch/multiprocessing/spawn.py:160\u001b[0m, in \u001b[0;36mProcessContext.join\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    159\u001b[0m msg \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m original_trace\n\u001b[0;32m--> 160\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m ProcessRaisedException(msg, error_index, failed_process\u001b[38;5;241m.\u001b[39mpid)\n",
      "\u001b[0;31mProcessRaisedException\u001b[0m: \n\n-- Process 0 terminated with the following error:\nTraceback (most recent call last):\n  File \"/opt/conda/lib/python3.10/site-packages/torch/multiprocessing/spawn.py\", line 69, in _wrap\n    fn(i, *args)\n  File \"/opt/conda/lib/python3.10/site-packages/accelerate/utils/launch.py\", line 535, in __call__\n    self.launcher(*args)\n  File \"/tmp/ipykernel_1569/2873471804.py\", line 3, in train_loop\n    accelerator = Accelerator(\n  File \"/opt/conda/lib/python3.10/site-packages/accelerate/accelerator.py\", line 358, in __init__\n    self.state = AcceleratorState(\n  File \"/opt/conda/lib/python3.10/site-packages/accelerate/state.py\", line 720, in __init__\n    PartialState(cpu, **kwargs)\n  File \"/opt/conda/lib/python3.10/site-packages/accelerate/state.py\", line 198, in __init__\n    torch.cuda.set_device(self.device)\n  File \"/opt/conda/lib/python3.10/site-packages/torch/cuda/__init__.py\", line 350, in set_device\n    torch._C._cuda_setDevice(device)\n  File \"/opt/conda/lib/python3.10/site-packages/torch/cuda/__init__.py\", line 235, in _lazy_init\n    raise RuntimeError(\nRuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method\n",
      "\nThe above exception was the direct cause of the following exception:\n",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[12], line 5\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01maccelerate\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m notebook_launcher\n\u001b[1;32m      3\u001b[0m args \u001b[38;5;241m=\u001b[39m (config,)\n\u001b[0;32m----> 5\u001b[0m \u001b[43mnotebook_launcher\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtrain_loop\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnum_processes\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/accelerate/launchers.py:140\u001b[0m, in \u001b[0;36mnotebook_launcher\u001b[0;34m(function, args, num_processes, mixed_precision, use_port)\u001b[0m\n\u001b[1;32m    138\u001b[0m         \u001b[38;5;28;01mexcept\u001b[39;00m ProcessRaisedException \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m    139\u001b[0m             \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mCannot re-initialize CUDA in forked subprocess\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m e\u001b[38;5;241m.\u001b[39margs[\u001b[38;5;241m0\u001b[39m]:\n\u001b[0;32m--> 140\u001b[0m                 \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\n\u001b[1;32m    141\u001b[0m                     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mCUDA has been initialized before the `notebook_launcher` could create a forked subprocess. \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    142\u001b[0m                     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mThis likely stems from an outside import causing issues once the `notebook_launcher()` is called. \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    143\u001b[0m                     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mPlease review your imports and test them when running the `notebook_launcher()` to identify \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    144\u001b[0m                     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mwhich one is problematic.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    145\u001b[0m                 ) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01me\u001b[39;00m\n\u001b[1;32m    147\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    148\u001b[0m     \u001b[38;5;66;03m# No need for a distributed launch otherwise as it's either CPU, GPU or MPS.\u001b[39;00m\n\u001b[1;32m    149\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m is_mps_available():\n",
      "\u001b[0;31mRuntimeError\u001b[0m: CUDA has been initialized before the `notebook_launcher` could create a forked subprocess. This likely stems from an outside import causing issues once the `notebook_launcher()` is called. Please review your imports and test them when running the `notebook_launcher()` to identify which one is problematic."
     ]
    }
   ],
   "source": [
    "from accelerate import notebook_launcher\n",
    "\n",
    "args = (config,)\n",
    "\n",
    "notebook_launcher(train_loop, args, num_processes=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd3701dd-b52f-4002-aab1-16a932b385a3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}