offer training class to easily train model off an image directory

README.md

@@ -38,6 +38,39 @@ sampled_images = diffusion.p_sample_loop((1, 3, 128, 128))
 sampled_images.shape # (1, 3, 128, 128)
 ```
 
+Or, if you simply want to pass in a folder name and the desired image dimensions, you can use the `Trainer` class to easily train a model.
+
+```python
+from denoising_diffusion_pytorch import Unet, GaussianDiffusion, Trainer
+
+model = Unet(
+    dim = 64,
+    dim_mults = (1, 2, 4, 8)
+).cuda()
+
+diffusion = GaussianDiffusion(
+    model,
+    beta_start = 0.0001,
+    beta_end = 0.02,
+    num_diffusion_timesteps = 1000,   # number of steps
+    loss_type = 'l1'                  # L1 or L2
+).cuda()
+
+trainer = Trainer(
+    diffusion,
+    'path/to/your/images',
+    image_size = 128,
+    train_batch_size = 32,
+    train_lr = 3e-4,
+    train_num_steps = 100000,
+    gradient_accumulate_every = 1
+)
+
+trainer.train()
+```
+
+Todo: Command line tool for one-line training
+
 ## Citations
 
 ```bibtex