Update innerRace.py

innerRace.py

@@ -1,14 +1,23 @@
 # -*- coding: utf-8 -*-
 """beths butterfly training.ipynb
+
 Automatically generated by Colab.
+
 Original file is located at
     https://colab.research.google.com/drive/1SbxWXhffEnCJ6tVT6ZfTDbY2-cxb063U
+
 # Train a diffusion model
+
 Unconditional image generation is a popular application of diffusion models that generates images that look like those in the dataset used for training. Typically, the best results are obtained from finetuning a pretrained model on a specific dataset. You can find many of these checkpoints on the [Hub](https://huggingface.co/search/full-text?q=unconditional-image-generation&type=model), but if you can't find one you like, you can always train your own!
+
 This tutorial will teach you how to train a [UNet2DModel](https://huggingface.co/docs/diffusers/main/en/api/models/unet2d#diffusers.UNet2DModel) from scratch on a subset of the [Smithsonian Butterflies](https://huggingface.co/datasets/huggan/smithsonian_butterflies_subset) dataset to generate your own 🦋 butterflies 🦋.
+
 <Tip>
+
 💡 This training tutorial is based on the [Training with 🧨 Diffusers](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/training_example.ipynb) notebook. For additional details and context about diffusion models like how they work, check out the notebook!
+
 </Tip>
+
 Before you begin, make sure you have 🤗 Datasets installed to load and preprocess image datasets, and 🤗 Accelerate, to simplify training on any number of GPUs. The following command will also install [TensorBoard](https://www.tensorflow.org/tensorboard) to visualize training metrics (you can also use [Weights & Biases](https://docs.wandb.ai/) to track your training).
 """
 
@@ -25,29 +34,34 @@ notebook_login()
 #!git config --global credential.helper store
 
 """Or login in from the terminal:
+
 ```bash
 huggingface-cli login
 ```
+
 Since the model checkpoints are quite large, install [Git-LFS](https://git-lfs.com/) to version these large files:
+
 ```bash
 !sudo apt -qq install git-lfs
 !git config --global credential.helper store
 ```
+
 ## Training configuration
+
 For convenience, create a `TrainingConfig` class containing the training hyperparameters (feel free to adjust them):
 """
 
 from dataclasses import dataclass
-import datetime
+
 
 @dataclass
 class TrainingConfig:
     image_size = 256  # the generated image resolution
     train_batch_size = 10
-    eval_batch_size = 16 # how many images to sample during evaluation
-    num_epochs = 50
+    eval_batch_size = 16  # how many images to sample during evaluation
+    num_epochs = 250
     gradient_accumulation_steps = 1
-    learning_rate = 1e-4 #Was 1e-4
+    learning_rate = 1e-4
     lr_warmup_steps = 500
     save_image_epochs = 25
     save_model_epochs = 25
@@ -59,13 +73,11 @@ class TrainingConfig:
     overwrite_output_dir = False  # KEEP THIS AS FALSE
     seed = 0
 
-    time_started = datetime.datetime.now().strftime('%a %d %b %Y, %I:%M%p')
 
 config = TrainingConfig()
 
-
-
 """## Load the dataset
+
 You can easily load the [Smithsonian Butterflies](https://huggingface.co/datasets/huggan/smithsonian_butterflies_subset) dataset with the 🤗 Datasets library:
 """
 
@@ -75,8 +87,11 @@ config.dataset_name = "GaumlessGraham/7inchInnerRace1730AugData"
 dataset = load_dataset(config.dataset_name, split="train")
 
 """<Tip>
+
 💡 You can find additional datasets from the [HugGan Community Event](https://huggingface.co/huggan) or you can use your own dataset by creating a local [`ImageFolder`](https://huggingface.co/docs/datasets/image_dataset#imagefolder). Set `config.dataset_name` to the repository id of the dataset if it is from the HugGan Community Event, or `imagefolder` if you're using your own images.
+
 </Tip>
+
 🤗 Datasets uses the [Image](https://huggingface.co/docs/datasets/main/en/package_reference/main_classes#datasets.Image) feature to automatically decode the image data and load it as a [`PIL.Image`](https://pillow.readthedocs.io/en/stable/reference/Image.html) which we can visualize:
 """
 
@@ -91,7 +106,9 @@ fig.show()
 """<div class="flex justify-center">
     <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/butterflies_ds.png"/>
 </div>
+
 The images are all different sizes though, so you'll need to preprocess them first:
+
 * `Resize` changes the image size to the one defined in `config.image_size`.
 * `RandomHorizontalFlip` augments the dataset by randomly mirroring the images.
 * `Normalize` is important to rescale the pixel values into a [-1, 1] range, which is what the model expects.
@@ -126,6 +143,7 @@ train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=config.train_
 fig.show()
 
 """## Create a UNet2DModel
+
 Pretrained models in 🧨 Diffusers are easily created from their model class with the parameters you want. For example, to create a [UNet2DModel](https://huggingface.co/docs/diffusers/main/en/api/models/unet2d#diffusers.UNet2DModel):
 """
 
@@ -163,8 +181,11 @@ print("Input shape:", sample_image.shape)
 print("Output shape:", model(sample_image, timestep=0).sample.shape)
 
 """Great! Next, you'll need a scheduler to add some noise to the image.
+
 ## Create a scheduler
+
 The scheduler behaves differently depending on whether you're using the model for training or inference. During inference, the scheduler generates image from the noise. During training, the scheduler takes a model output - or a sample - from a specific point in the diffusion process and applies noise to the image according to a *noise schedule* and an *update rule*.
+
 Let's take a look at the [DDPMScheduler](https://huggingface.co/docs/diffusers/main/en/api/schedulers/ddpm#diffusers.DDPMScheduler) and use the `add_noise` method to add some random noise to the `sample_image` from before:
 """
 
@@ -182,6 +203,7 @@ noisy_image = noise_scheduler.add_noise(sample_image, noise, timesteps)
 """<div class="flex justify-center">
     <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/noisy_butterfly.png"/>
 </div>
+
 The training objective of the model is to predict the noise added to the image. The loss at this step can be calculated by:
 """
 
@@ -191,7 +213,9 @@ noise_pred = model(noisy_image, timesteps).sample
 loss = F.mse_loss(noise_pred, noise)
 
 """## Train the model
+
 By now, you have most of the pieces to start training the model and all that's left is putting everything together.
+
 First, you'll need an optimizer and a learning rate scheduler:
 """
 
@@ -209,7 +233,7 @@ lr_scheduler = get_cosine_schedule_with_warmup(
 from diffusers import DDPMPipeline
 import math
 import os
-import sys
+
 
 def make_grid(images, rows, cols):
     w, h = images[0].size
@@ -231,13 +255,16 @@ def evaluate(config, epoch, pipeline):
     image_grid = make_grid(images, rows=4, cols=4)
 
     # Save the images
-    test_dir = os.path.join(config.output_dir, "samples"+config.time_started)
+    test_dir = os.path.join(config.output_dir, "samples")
     os.makedirs(test_dir, exist_ok=True)
     image_grid.save(f"{test_dir}/{epoch:04d}.png")
 
 """Now you can wrap all these components together in a training loop with 🤗 Accelerate for easy TensorBoard logging, gradient accumulation, and mixed precision training. To upload the model to the Hub, write a function to get your repository name and information and then push it to the Hub.
+
 <Tip>
+
 💡 The training loop below may look intimidating and long, but it'll be worth it later when you launch your training in just one line of code! If you can't wait and want to start generating images, feel free to copy and run the code below. You can always come back and examine the training loop more closely later, like when you're waiting for your model to finish training. 🤗
+
 </Tip>
 """
 
@@ -272,7 +299,7 @@ def train_loop(config, model, noise_scheduler, optimizer, train_dataloader, lr_s
             repo = Repository(config.output_dir, clone_from=repo_name)
         elif config.output_dir is not None:
             os.makedirs(config.output_dir, exist_ok=True)
-        accelerator.init_trackers("log_file " + config.time_started)
+        accelerator.init_trackers("train_example")
 
     # Prepare everything
     # There is no specific order to remember, you just need to unpack the
@@ -329,6 +356,11 @@ def train_loop(config, model, noise_scheduler, optimizer, train_dataloader, lr_s
 
             if (epoch + 1) % config.save_model_epochs == 0 or epoch == config.num_epochs - 1:
                 if config.push_to_hub:
+
+                    model_dir = os.path.join(config.output_dir, str(epoch))
+                    os.makedirs(model_dir, exist_ok=True)
+                    
+                    pipeline.save_pretrained(model_dir)
                     repo.push_to_hub(commit_message=f"Epoch {epoch}", blocking=True)
                 else:
                     pipeline.save_pretrained(config.output_dir)
@@ -352,4 +384,18 @@ notebook_launcher(train_loop, args, num_processes=1)
 import glob
 
 sample_images = sorted(glob.glob(f"{config.output_dir}/samples/*.png"))
-Image.open(sample_images[-1])
+Image.open(sample_images[-1])
+
+"""<div class="flex justify-center">
+    <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/butterflies_final.png"/>
+</div>
+
+## Next steps
+
+Unconditional image generation is one example of a task that can be trained. You can explore other tasks and training techniques by visiting the [🧨 Diffusers Training Examples](https://huggingface.co/docs/diffusers/main/en/tutorials/../training/overview) page. Here are some examples of what you can learn:
+
+* [Textual Inversion](https://huggingface.co/docs/diffusers/main/en/tutorials/../training/text_inversion), an algorithm that teaches a model a specific visual concept and integrates it into the generated image.
+* [DreamBooth](https://huggingface.co/docs/diffusers/main/en/tutorials/../training/dreambooth), a technique for generating personalized images of a subject given several input images of the subject.
+* [Guide](https://huggingface.co/docs/diffusers/main/en/tutorials/../training/text2image) to finetuning a Stable Diffusion model on your own dataset.
+* [Guide](https://huggingface.co/docs/diffusers/main/en/tutorials/../training/lora) to using LoRA, a memory-efficient technique for finetuning really large models faster.
+"""