{ "nbformat": 4, "nbformat_minor": 0, "metadata": { "colab": { "name": "Stylegan2-ada Custom Training.ipynb", "provenance": [], "private_outputs": true, "collapsed_sections": [], "toc_visible": true, "machine_shape": "hm", "mount_file_id": "1DceSJrxvS1g3cm7ZpdKPlCrX_--xfm_i", "authorship_tag": "ABX9TyPpCmsx1K83ZEugZkpg4uLQ", "include_colab_link": true }, "kernelspec": { "name": "python3", "display_name": "Python 3" }, "accelerator": "GPU" }, "cells": [ { "cell_type": "markdown", "metadata": { "id": "view-in-github", "colab_type": "text" }, "source": [ "\"Open" ] }, { "cell_type": "markdown", "metadata": { "id": "cPI5E5y0pujD" }, "source": [ "# Custom Training StyleGan2-ADA" ] }, { "cell_type": "markdown", "metadata": { "id": "SI_i1MwgpzOD" }, "source": [ "StyleGAN2-ADA only work with Tensorflow 1. Run the next cell before anything else to make sure we’re using TF1 and not TF2." ] }, { "cell_type": "code", "metadata": { "id": "iKYAU7Wub3WW" }, "source": [ "%tensorflow_version 1.x" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "51ei6d5kxVDm" }, "source": [ "!nvidia-smi" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "id": "5YcUMPQp6ipP" }, "source": [ "## Install Repo to Google Drive\n", "\n", "Colab is a little funky with training. I’ve found the best way to do this is to install the repo directly into your Google Drive folder.\n", "\n", "First, mount your Drive to the Colab notebook: " ] }, { "cell_type": "code", "metadata": { "id": "pxxYlEKI9Gis" }, "source": [ "from google.colab import drive\n", "drive.mount('/content/drive')" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "id": "epV6TDzAjox1" }, "source": [ "Next, run this cell. If you’re already installed the repo, it will skip the installation process and change into the repo’s directory. If you haven’t installed it, it will install all the files necessary." ] }, { "cell_type": "code", "metadata": { "id": "8HX77jscX2zV" }, "source": [ "import os\n", "if os.path.isdir(\"/content/drive/My Drive/colab-sg2-ada\"):\n", " %cd \"/content/drive/My Drive/colab-sg2-ada/stylegan2-ada\"\n", "else:\n", " #install script\n", " %cd \"/content/drive/My Drive/\"\n", " !mkdir colab-sg2-ada\n", " %cd colab-sg2-ada\n", " !git clone https://github.com/dvschultz/stylegan2-ada\n", " %cd stylegan2-ada\n", " !mkdir downloads\n", " !mkdir datasets" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "G7oj_kBaemol" }, "source": [ "%cd \"/content/drive/My Drive/colab-sg2-ada/stylegan2-ada\"\n", "!git config --global user.name \"test\"\n", "!git config --global user.email \"test@test.com\"\n", "!git fetch origin\n", "!git checkout origin/main -- train.py" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "id": "HeS9tDvt61VG" }, "source": [ "## Convert dataset to .tfrecords" ] }, { "cell_type": "markdown", "metadata": { "id": "_Q58MJbckLUc" }, "source": [ "**Note: You only need to do this once per dataset. If you have already run this and are returning to conntinue training, skip these cells.**\n", "\n", "Next we need to convert our image dataset to a format that StyleGAN2-ADA can read from. There are two options here. You can upload your dataset directly to Colab (as a zipped file), or you can upload it to Drive directly and read it from there." ] }, { "cell_type": "code", "metadata": { "id": "8JUP51nJdEjz" }, "source": [ "#if you manually uploaded your dataset to Colab, unzip it\n", "zip_path = \"/content/CAT1024.zip\"\n", "\n", "!unzip {zip_path} -d /content/" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "id": "D0QH0nzjlbEE" }, "source": [ "Now that your image dataset is uploaded, we need to convert it to the `.tfrecords` format.\n", "\n", "Depending on the resolution of your images and how many you have, this can take a while." ] }, { "cell_type": "code", "metadata": { "id": "T-BZHhBe7AvO" }, "source": [ "#update this to the path to your image folder\n", "dataset_path = \"/content/CAT1024\"\n", "#give your dataset a name\n", "dataset_name = \"cat1024\"\n", "\n", "#you don't need to edit anything here\n", "!python dataset_tool.py create_from_images ./datasets/{dataset_name} {dataset_path}" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "id": "8DvTupHzP2s_" }, "source": [ "## Train a custom model\n", "\n", "We’re ready to start training! There are numerous arguments to training, what’s listed below are the most popular options. To see all the options, run the following cell." ] }, { "cell_type": "code", "metadata": { "id": "Fxu7CA0Qb1Yd" }, "source": [ "!python train.py --help" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": { "id": "jOftFoyiDU3s" }, "source": [ "#this name must EXACTLY match the dataset name you used when creating the .tfrecords file\n", "dataset_name = \"dante1024\"\n", "#how often should the model generate samples and a .pkl file\n", "snapshot_count = 4\n", "#should the images be mirrored left to right?\n", "mirrored = True\n", "#should the images be mirrored top to bottom?\n", "mirroredY = True\n", "#metrics? \n", "metric_list = None\n", "#augments\n", "augs = \"bg\"\n", "\n", "#\n", "# this is the most important cell to update\n", "#\n", "# running it for the first time? set it to ffhq(+resolution)\n", "# resuming? get the path to your latest .pkl file and use that\n", "resume_from = \"/content/drive/My\\ Drive/colab-sg2-ada2/stylegan2-ada/results/00008-dante1024-mirror-mirrory-11gb-gpu-bg-resumecustom/network-snapshot-000160.pkl\"\n", "\n", "#don't edit this unless you know what you're doing :)\n", "!python train.py --outdir ./results --snap={snapshot_count} --cfg=11gb-gpu --data=./datasets/{dataset_name} --augpipe={augs} --mirror={mirrored} --mirrory={mirroredY} --metrics={metric_list} --resume={resume_from} --augpipe=\"bg\"" ], "execution_count": null, "outputs": [] }, { "cell_type": "markdown", "metadata": { "id": "lohotw1FqC54" }, "source": [ "### While it’s training...\n", "**Once the above cell is running you should be training!**\n", "\n", "Don’t close this tab! Colab needs to be open and running in order to continue training. Every ~15min or so a new line should get added to your output, indicated its still training. Depending on you `snapshot_count` setting you should see the results folder in your Google drive folder fill with both samples (`fakesXXXXXx.jpg`) and model weights (`network-snapshot-XXXXXX.pkl`). The samples are worth looking at while it trains but don’t get too worried about each individual sample.\n", "\n", "If you chose a metric, you will also see scores for each snapshot. Don’t obsess over these! they are a guide, it can go up or down slightly for each snapshot. What you want to see is a gradual lowering of the score over time.\n", "\n", "Once Colab shuts off, you can Reconnect the notebook and re-run every cell from top to bottom. Make sure you update the `resume_from` path to continue training from the latest model." ] }, { "cell_type": "code", "metadata": { "id": "UZR4ireZE6JF" }, "source": [ "" ], "execution_count": null, "outputs": [] } ] }