StyleGAN2_PyTorch_Train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jG7ZEc_982io"
   },
   "source": [
    "# StyleGAN2-ADA-PyTorch\n",
    "\n",
    "**Notes**\n",
    "* Training and Inference sections should be fairly stable. I’ll slowly add new features but it should work for most mainstream use cases.\n",
    "* Advanced Features are being documented toward the bottom of this notebook\n",
    "\n",
    "---\n",
    "\n",
    "If you find this notebook useful, consider signing up for my [Patreon](https://www.patreon.com/bustbright) or [YouTube channel](https://www.youtube.com/channel/UCaZuPdmZ380SFUMKHVsv_AA/join). You can also send me a one-time payment on [Venmo](https://venmo.com/Derrick-Schultz)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Vj4PG4_i9Alt"
   },
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qGEXPcFJ9UTY"
   },
   "source": [
    "Let’s start by checking to see what GPU we’ve been assigned. Ideally we get a V100, but a P100 is fine too. Other GPUs may lead to issues."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 7,
     "status": "ok",
     "timestamp": 1685247953230,
     "user": {
      "displayName": "한상진",
      "userId": "05594748000998016422"
     },
     "user_tz": -540
    },
    "id": "7VVICTCvd4mc",
    "outputId": "d7803540-5d14-4464-b12a-2b88f9650e03"
   },
   "outputs": [],
   "source": [
    "!nvidia-smi -L"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rSV_HEoD9dxo"
   },
   "source": [
    "Next let’s connect our Google Drive account. This is optional but highly recommended."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 18728,
     "status": "ok",
     "timestamp": 1685248091584,
     "user": {
      "displayName": "한상진",
      "userId": "05594748000998016422"
     },
     "user_tz": -540
    },
    "id": "IuVPuJmbigRs",
    "outputId": "bdc9a494-62c6-46d8-e757-feef04a200df"
   },
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 11594,
     "status": "ok",
     "timestamp": 1685248103175,
     "user": {
      "displayName": "한상진",
      "userId": "05594748000998016422"
     },
     "user_tz": -540
    },
    "id": "mkBET-Cq-otL",
    "outputId": "804385d4-ff8a-4314-9cfd-29f31ceee74f"
   },
   "outputs": [],
   "source": [
    "!sudo apt install python3.7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 88742,
     "status": "ok",
     "timestamp": 1685248191914,
     "user": {
      "displayName": "한상진",
      "userId": "05594748000998016422"
     },
     "user_tz": -540
    },
    "id": "Y0C04YAx-vs-",
    "outputId": "0f246c17-4abe-4bb3-c801-5290d8613a10"
   },
   "outputs": [],
   "source": [
    "!sudo update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.7 1\n",
    "!sudo update-alternatives --config python3\n",
    "!sudo apt install python3-pip\n",
    "!sudo apt install python3.7-distutils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 620,
     "status": "ok",
     "timestamp": 1685248192528,
     "user": {
      "displayName": "한상진",
      "userId": "05594748000998016422"
     },
     "user_tz": -540
    },
    "id": "RNIt1nuK_MvD",
    "outputId": "4afe25c4-5879-4e29-889a-4f35742e07f0"
   },
   "outputs": [],
   "source": [
    "!python --version\n",
    "!pip --version"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nTjVmfSK9CYa"
   },
   "source": [
    "## Install repo\n",
    "\n",
    "The next cell will install the StyleGAN repository in Google Drive. If you have already installed it it will just move into that folder. If you don’t have Google Drive connected it will just install the necessary code in Colab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 752
    },
    "executionInfo": {
     "elapsed": 9251,
     "status": "ok",
     "timestamp": 1685248201778,
     "user": {
      "displayName": "한상진",
      "userId": "05594748000998016422"
     },
     "user_tz": -540
    },
    "id": "B8ADVNpBh8Ox",
    "outputId": "e64ec40e-6909-406f-8e52-7cf0a019a73f"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "!pip install gdown --upgrade\n",
    "\n",
    "if os.path.isdir(\"/content/drive/MyDrive/colab-sg2-ada-pytorch\"):\n",
    "    %cd \"/content/drive/MyDrive/colab-sg2-ada-pytorch/stylegan2-ada-pytorch\"\n",
    "elif os.path.isdir(\"/content/drive/\"):\n",
    "    #install script\n",
    "    %cd \"/content/drive/MyDrive/\"\n",
    "    !mkdir colab-sg2-ada-pytorch\n",
    "    %cd colab-sg2-ada-pytorch\n",
    "    !git clone https://github.com/dvschultz/stylegan2-ada-pytorch\n",
    "    %cd stylegan2-ada-pytorch\n",
    "    !mkdir downloads\n",
    "    !mkdir datasets\n",
    "    !mkdir pretrained\n",
    "    !gdown --id 1-5xZkD8ajXw1DdopTkH_rAoCsD72LhKU -O /content/drive/MyDrive/colab-sg2-ada-pytorch/stylegan2-ada-pytorch/pretrained/wikiart.pkl\n",
    "else:\n",
    "    !git clone https://github.com/dvschultz/stylegan2-ada-pytorch\n",
    "    %cd stylegan2-ada-pytorch\n",
    "    !mkdir downloads\n",
    "    !mkdir datasets\n",
    "    !mkdir pretrained\n",
    "    %cd pretrained\n",
    "    !gdown --id 1-5xZkD8ajXw1DdopTkH_rAoCsD72LhKU\n",
    "    %cd ../"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "executionInfo": {
     "elapsed": 329290,
     "status": "ok",
     "timestamp": 1685248531057,
     "user": {
      "displayName": "한상진",
      "userId": "05594748000998016422"
     },
     "user_tz": -540
    },
    "id": "jBeDEGqEbmLy",
    "outputId": "a88ef3d0-72f5-44d3-fccc-555707135b73"
   },
   "outputs": [],
   "source": [
    "#Uninstall new JAX\n",
    "!pip uninstall jax jaxlib -y\n",
    "#GPU frontend\n",
    "!pip install \"jax[cuda11_cudnn805]==0.3.10\" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
    "#CPU frontend\n",
    "#!pip install jax[cpu]==0.3.10\n",
    "#Downgrade Pytorch\n",
    "!pip uninstall torch torchvision -y\n",
    "!pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html\n",
    "!pip install timm==0.4.12 ftfy==6.1.1 ninja==1.10.2 opensimplex"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9jMmUpn4DWRe"
   },
   "source": [
    "You probably don’t need to run this, but this will update your repo to the latest and greatest."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 73634,
     "status": "ok",
     "timestamp": 1684773801219,
     "user": {
      "displayName": "구태형",
      "userId": "12113658638621684006"
     },
     "user_tz": -540
    },
    "id": "uV9bdvzeDRPd",
    "outputId": "a7888a66-1b91-4ff8-8beb-c9be854b43aa"
   },
   "outputs": [],
   "source": [
    "%cd \"/content/drive/My Drive/colab-sg2-ada-pytorch/stylegan2-ada-pytorch\"\n",
    "!git config --global user.name \"test\"\n",
    "!git config --global user.email \"test@test.com\"\n",
    "!git fetch origin\n",
    "!git pull\n",
    "!git stash\n",
    "!git checkout origin/main -- train.py generate.py legacy.py closed_form_factorization.py flesh_digression.py apply_factor.py README.md calc_metrics.py training/stylegan2_multi.py training/training_loop.py util/utilgan.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cZkcJ58P97Ls"
   },
   "source": [
    "## Dataset Preparation\n",
    "\n",
    "Upload a .zip of square images to the `datasets` folder. Previously you had to convert your model to .tfrecords. That’s no longer needed :)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5B-h6FpB9FaK"
   },
   "source": [
    "## Train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 342,
     "status": "ok",
     "timestamp": 1685247474035,
     "user": {
      "displayName": "한상진",
      "userId": "05594748000998016422"
     },
     "user_tz": -540
    },
    "id": "7IwZrIfzkT2e",
    "outputId": "a39f0870-8f1a-48c6-facf-0591ec3dd803"
   },
   "outputs": [],
   "source": [
    "!python stylegan2-ada-pytorch/train.py \\\n",
    " --outdir={output_folder} \\\n",
    " --snap={snap} \\\n",
    " --metrics=none \\\n",
    " --data={'/content/drive/MyDrive/face_img.zip'} \\\n",
    " --cfg=auto_norp \\\n",
    " --cifar_tune=1 \\\n",
    " --gamma={gamma} \\\n",
    " --kimg={kimg} \\\n",
    " --batch={mini_batch} \\\n",
    " --cfg_map=8 \\\n",
    " --augpipe=bg \\\n",
    " --freezed=10 \\\n",
    " --resume=ffhq_1024 \\\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bNc-3wTO-MUd"
   },
   "source": [
    "Below are a series of variables you need to set to run the training. You probably won’t need to touch most of them.\n",
    "\n",
    "* `dataset_path`: this is the path to your .zip file\n",
    "* `resume_from`: if you’re starting a new dataset I recommend `'ffhq1024'` or `'./pretrained/wikiart.pkl'`\n",
    "* `mirror_x` and `mirror_y`: Allow the dataset to use horizontal or vertical mirroring."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JV0W6yxP-UIn"
   },
   "outputs": [],
   "source": [
    "#required: definitely edit these!\n",
    "dataset_path = '/content/drive/MyDrive/face_img.zip'\n",
    "resume_from = '/content/drive/MyDrive/ffhq.pkl'\n",
    "aug_strength = 0.328\n",
    "train_count = 10\n",
    "mirror_x = True\n",
    "#mirror_y = False\n",
    "\n",
    "#optional: you might not need to edit these\n",
    "gamma_value = 50.0\n",
    "augs = 'bg'\n",
    "config = 'auto'\n",
    "snapshot_count = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "executionInfo": {
     "elapsed": 24678,
     "status": "ok",
     "timestamp": 1685248555725,
     "user": {
      "displayName": "한상진",
      "userId": "05594748000998016422"
     },
     "user_tz": -540
    },
    "id": "y75dVHbYKUqg",
    "outputId": "87947b1e-0a07-4926-b42f-5e92bc118d4d"
   },
   "outputs": [],
   "source": [
    "!pip install click requests tqdm pyspng psutil tensorboard ninja imageio-ffmpeg==0.4.3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 9452,
     "status": "ok",
     "timestamp": 1685248565171,
     "user": {
      "displayName": "한상진",
      "userId": "05594748000998016422"
     },
     "user_tz": -540
    },
    "id": "heS20GQyeLOn",
    "outputId": "7b5b185c-52b2-4647-928c-c4ba7b147256"
   },
   "outputs": [],
   "source": [
    "!sudo apt install python3.7-dev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true,
     "base_uri": "https://localhost:8080/"
    },
    "id": "EL-M7WnnfMDI",
    "outputId": "788d0da2-194a-4123-e1e9-b57b6ec09c8b"
   },
   "outputs": [],
   "source": [
    "!python train.py --gpus=1 --cfg=$config --metrics=None --outdir=./results2 --data=$dataset_path --snap=$snapshot_count --resume=$resume_from --augpipe=$augs --initstrength=$aug_strength --gamma=$gamma_value --mirror=$mirror_x --mirrory=False --nkimg=$train_count --freezed=4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zAuYTVs9H1c7"
   },
   "outputs": [],
   "source": [
    "#required: definitely edit these!\n",
    "dataset_path = '/content/drive/MyDrive/resized.zip'\n",
    "resume_from = '/content/drive/MyDrive/ffhq_1024.pkl'\n",
    "aug_strength = 0.328\n",
    "train_count = 100\n",
    "mirror_x = False\n",
    "#mirror_y = False\n",
    "\n",
    "#optional: you might not need to edit these\n",
    "gamma_value = 50.0\n",
    "augs = 'bg'\n",
    "config = 'paper256'\n",
    "snapshot_count = 4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RgvSvfyi_R_-"
   },
   "source": [
    "### Resume Training\n",
    "\n",
    "Once Colab has shutdown, you’ll need to resume your training. Reset the variables above, particularly the `resume_from` and `aug_strength` settings.\n",
    "\n",
    "1. Point `resume_from` to the last .pkl you trained (you’ll find these in the `results` folder)\n",
    "2. Update `aug_strength` to match the augment value of the last pkl file. Often you’ll see this in the console, but you may need to look at the `log.txt`. Updating this makes sure training stays as stable as possible.\n",
    "3. You may want to update `train_count` to keep track of your training progress.\n",
    "\n",
    "Once all of this has been reset, run that variable cell and the training command cell after it."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuClass": "premium",
   "machine_shape": "hm",
   "provenance": [
    {
     "file_id": "1ldNI7bHjZgvvL8RaqYgdT4oM0mVD6Hkt",
     "timestamp": 1687325633925
    }
   ]
  },
  "gpuClass": "premium",
  "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}