To start, we need to install Lightning, but first, I will create and activate a new Python 3.8 environment for my new app. Check out <\/span>this Lightning Bits episode<\/u><\/span><\/a> to learn more about why this is important. To manage Python environments, I use <\/span>conda<\/u><\/span><\/a>.<\/span><\/p>\n While my conda environment is being created, let\u2019s review what we know about GANs.\u00a0<\/span><\/p>\n A Generative Adversarial Network is a machine learning (ML) model in which two neural networks compete with each other to become more accurate in their predictions. When implementing GANs, we need two networks: generator and discriminator. Generator is a neural network tasked with creating something out of random noise (also called seed). Discriminator is a neural network that will learn to differentiate between real and generated output. The two networks are being trained alongside each other, with discriminator\u2019s output being used to train the generator to generate more realistic data and generator\u2019s output being used to train discriminator to detect fakes better. In this possibly endless competition, generator has the opportunity to improve infinitely. Ultimately, we get a neural network that can generate very trustworthy data samples. It\u2019s no surprise that GANs took the world by storm and continues to impress researchers.<\/span><\/p>\n If you followed me and created a new Python environment, you should already have it configured by now. The next step is to install Lightning:<\/span><\/p>\n We are going to start with the Lightning App quick start template. To install it, execute:<\/span><\/p>\n It will create a new project directory called We also need to make sure to install the dependencies we are going to use. It\u2019s easy to do if you replace the contents of the project\u2019s And run <\/span> Next, find the file named First, add the line saying <\/span> Next, find the <\/span>GAN<\/span> class definition and make sure that the <\/span> And last, but not least, replace:<\/span><\/p>\n With:<\/span><\/p>\n Now, we will make some changes to the Lightning App itself. Open the file What we will need to modify are the train and serve works. Open the file We will not use pre-trained weights for our GAN and we don\u2019t want to limit the number of batches.<\/span><\/p>\n Also, replace:<\/span><\/p>\n With:<\/span><\/p>\n This makes Lightning always save our model at the end of training.<\/span><\/p>\n Next, find the line that says <\/span> That\u2019s it. We are done with our training work. When we run the app, the training work will train the model, save the weights, and expose the path to the weights as the <\/span> And don\u2019t forget to add <\/span> This <\/span> Now, we can run our new GAN Lightning App:<\/span><\/p>\n The app will launch and open a new page in the browser, where you can see two tabs. The first Model Training<\/strong> tab contains TensorBoard, where you can explore any logs and metrics that your Once the app is trained, the Interactive Demo tab will load, and you\u2019ll be able to generate some MNIST numbers here.<\/span><\/p>\n If you followed this article, you now have a running GAN Lightning App. It is impressive how easily you can build a functional and interactive Lightning App with a custom UI, but it\u2019s only the beginning. I am sure you can create so much more with Lightning, and I can\u2019t wait to see what you build! <\/span>Join our Lightning community<\/u><\/a> to see what others are building and share what you\u2019ve built.<\/span><\/p>\n If you have more questions, check out <\/span>the GAN app’s source code<\/u><\/a>, or read <\/span>the Lightning documentation<\/u><\/a>.<\/span><\/p>\n PS: you can also <\/span>run your app in the cloud<\/u><\/a> without changing a line of your code! Just add <\/span>$ conda create -n lit-gan python=3.8\r\n\r\n$ conda activate lit-gan\r\n<\/pre>\n
$ pip install lightning\r\n<\/pre>\n
$ lightning install app lightning\/quick-start\r\n<\/pre>\n
lightning-quick-start<\/code> for you. Open it in the editor of your choice. I will use VS Code for this walkthrough, but if you have questions about why or which editor to use, check out <\/span>this Lightning Bits episode on IDEs<\/u><\/a>.<\/span><\/p>\n![\"Lightning](\"https:\/\/lightningaidev.wpengine.com\/wp-content\/uploads\/2022\/07\/lightning-quick-start.png\")
Before we continue, let\u2019s open <\/span>the PyTorch Lightning tutorial on GANs<\/u><\/a>. We will borrow a lot of source code from here to get our app working even faster. In this example, we will train our GAN to generate digits for the MNIST dataset.<\/span><\/p>\nrequirements.txt<\/code> file with the following:<\/span><\/p>\ntorchvision\r\npytorch-lightning==1.6.3\r\njsonargparse[signatures]==4.7.3\r\nwandb==0.12.16\r\ngradio==2.9.4\r\npyyaml==5.4.0\r\nprotobuf<4.21.0 # 4.21 breaks with wandb, tensorboard, or pytorch-lightning: https:\/\/github.com\/protocolbuffers\/protobuf\/issues\/10048 torchmetrics>=0.3\r\ntorch>=1.6, <1.9\r\n<\/pre>\n
pip install -r requirements.txt<\/code> from the terminal.<\/span><\/p>\ntrain_script.py<\/code> in your IDE and replace its content with the code from the PyTorch Lightning tutorial linked above, but don\u2019t copy the first line containing the <\/span>pip install<\/code> command. We don\u2019t need it because we already installed our dependencies. However, we are going to make some modifications to it.<\/span><\/p>\nfrom pytorch_lightning.utilities.cli import LightningCLI<\/code> to the top of the file to import the Lightning CLI. This isn\u2019t required, but it\u2019s nice to have our training script modernized to use <\/span>Lightning CLI because it significantly improves the quality of the development process<\/u><\/a>.<\/span><\/p>\n__init__<\/code> method declaration looks like this:<\/span><\/p>\nclass GAN(LightningModule):\r\n def __init__(\r\n self,\r\n channels: int = 1,\r\n width: int = 28,\r\n height: int = 28,\r\n latent_dim: int = 100,\r\n lr: float = 0.0002,\r\n b1: float = 0.5,\r\n b2: float = 0.999,\r\n batch_size: int = BATCH_SIZE,\r\n **kwargs\r\n ):\r\n<\/pre>\n
dm = MNISTDataModule()\r\nmodel = GAN(*dm.size())\r\ntrainer = Trainer(\r\n accelerator=\"auto\",\r\n devices=1 if torch.cuda.is_available() else None, # limiting got iPython runs\r\n max_epochs=5,\r\n callbacks=[TQDMProgressBar(refresh_rate=20)],\r\n)\r\ntrainer.fit(model, dm)\r\n\r\n# Start tensorboard.\r\n%load_ext tensorboard\r\n%tensorboard --logdir lightning_logs\/\r\n<\/pre>\n
if __name__ == \"__main__\":\r\n cli = LightningCLI(\r\n GAN, MNISTDataModule, seed_everything_default=42, save_config_overwrite=True, run=False\r\n )\r\n cli.trainer.fit(cli.model, datamodule=cli.datamodule)\r\n<\/pre>\n
app.py<\/code>. Our quick start Lightning App consists of two works \u2014 train work and serve work \u2014 and a Lightning Flow to orchestrate them.<\/span><\/p>\n![\"Lightning](\"https:\/\/lightningaidev.wpengine.com\/wp-content\/uploads\/2022\/07\/terminal-in-lightning-apps.png\")
This is exactly what we need for our app, so we will leave the app.py<\/code> file as-is, but you may want to increase the number of training epochs to get nicer results.<\/span><\/p>\nquick_start\/components.py<\/code> and find the class <\/span>PyTorchLightningScript<\/code>. The method we are going to focus on is <\/span>run<\/code>. Remove the following lines from it:<\/span><\/p>\ndownload_data(\"https:\/\/pl-flash-data.s3.amazonaws.com\/assets_lightning\/demo_weights.pt\", \".\/\")\r\n\"--trainer.limit_train_batches=4\",\r\n\"--trainer.limit_val_batches=4\",\r\n<\/pre>\n
\"--trainer.callbacks.monitor=val_acc\",\r\n<\/pre>\n
\"--trainer.callbacks.save_last=true<\/span>\",\r\n<\/pre>\ncheckpoint = torch.load(res[\"cli\"].trainer.checkpoint_callback.best_model_path)<\/code> inside of the <\/span>on_after_run<\/code> method and replace <\/span>best_model_path<\/code> in it with <\/span>last_model_path<\/code>. This will make sure that we are exposing our model to the serve work after training.<\/span><\/p>\nbest_model_path<\/code> property. Now, we will replace the <\/span>ImageServeGradio<\/code> class with the following code:<\/span><\/p>\nclass ImageServeGradio(ServeGradio):\r\n\r\n inputs = [\r\n gr.inputs.Slider(0, 1000, label=\"Seed\", default=42),\r\n gr.inputs.Slider(4, 64, label=\"Number of Digits\", step=1, default=10),\r\n ]\r\n outputs = \"image\"\r\n examples = [[27, 5], [18, 4], [256, 8], [1337, 35]]\r\n\r\n def __init__(self, cloud_compute, *args, **kwargs):\r\n super().__init__(*args, cloud_compute=cloud_compute, **kwargs)\r\n self.model_path = None\r\n\r\n def run(self, model_path):\r\n self.model_path = model_path\r\n super().run()\r\n\r\n def predict(self, seed, num_digits):\r\n torch.manual_seed(seed)\r\n z = torch.randn(num_digits, 100)\r\n digits = self.model(z)\r\n save_image(digits, \"digits.png\", normalize=True)\r\n return \"digits.png\"\r\n\r\n def build_model(self):\r\n model = torch.load(self.model_path)\r\n for p in model.parameters():\r\n p.requires_grad = False\r\n model.eval()\r\n return model\r\n<\/pre>\n
from torchvision.utils import save_image<\/code> at the top of the file.<\/span><\/p>\nImageServeGradio<\/code> class extends on Lightning ServeGradio, creating two inputs (sliders to control the random seed and the number of digits that our GAN will generate),one output (an image), and a list of examples. This class also defines the <\/span>run<\/code> and <\/span>build_model<\/code> methods used to load the model and the <\/span>predict<\/code> method, which runs inference on our model.<\/span><\/p>\nlightning run app app.py\r\n<\/pre>\n
train_script.py<\/code> is logging (feel free to add more TensorBoard logs there \ud83d\ude09 ), and the second Interactive Demo<\/strong> tab will show you a loading spinner until the training is complete. That\u2019s ok; you can monitor the training progress either in TensorBoard or in the terminal logs.<\/span><\/p>\n![\"TensorBoard](\"https:\/\/lightningaidev.wpengine.com\/wp-content\/uploads\/2022\/07\/lightning-tensorboard.png\")
TensorBoard tab<\/em><\/span><\/p>\n![\"Interactive](\"https:\/\/lightningaidev.wpengine.com\/wp-content\/uploads\/2022\/07\/interactive-demo-loading-in-lightning.png\")
\u201cLoading\u201d Interactive Demo tab<\/span><\/i><\/p>\n![\"Lightning](\"https:\/\/lightningaidev.wpengine.com\/wp-content\/uploads\/2022\/07\/lightning-terminal-logs-1.png\")
Logs in the terminal<\/span><\/i><\/p>\n![\"Interactive](\"https:\/\/lightningaidev.wpengine.com\/wp-content\/uploads\/2022\/07\/interactive-demo-in-lightning.png\")
Play around with the seed and the number of digits to see different results. You can also tune hyperparameters and the models in train_script.py<\/code> to improve the results.<\/span><\/p>\n--cloud<\/code> to your <\/span>run<\/code> command:<\/span><\/p>\nlightning run app app.py \u2013cloud\r\n<\/pre>\n