Accelerate (#114)

* Refactor optimize

* Add accelerate tests

* Remove line

* Make style

* Print out accelerate

* Fixup

* Fixup

* Fixup

* Add mixed precision with accelerate

* Fixup mixed precision

* Make style

* Print logs only for rank zero

* Allow tf32

* Fixup

* Fixup

* Fix setting up logger

* Update README

* Fix linter

* Fixup

* Update test

* Handle different types in prepare

* Compute grad norm

* Fix eval_epoch single run

* Make style

* Meh linter

Author: erogol

.pylintrc

@@ -404,10 +404,9 @@ logging-modules=logging
 [MESSAGES CONTROL]
 
 # Only show warnings with the listed confidence levels. Leave empty to show
-# all. Valid levels: HIGH, CONTROL_FLOW, INFERENCE, INFERENCE_FAILURE,
+# all. Valid levels: HIGH, INFERENCE, INFERENCE_FAILURE,
 # UNDEFINED.
 confidence=HIGH,
-           CONTROL_FLOW,
            INFERENCE,
            INFERENCE_FAILURE,
            UNDEFINED

README.md

@@ -25,11 +25,65 @@ Prefer installing from Github as it is more stable.
 Subclass and overload the functions in the [```TrainerModel()```](trainer/model.py)
 
 
-## Training a model with auto optimization
+## Training a model with auto-optimization
 See the [MNIST example](examples/train_mnist.py).
 
 
 ## Training a model with advanced optimization
+With 👟 you can define the whole optimization cycle as you want as the in GAN example below. It enables more
+under-the-hood control and flexibility for more advanced training loops.
+
+You just have to use the ```scaled_backward()``` function to handle mixed precision training.
+
+```python
+...
+
+def optimize(self, batch, trainer):
+    imgs, _ = batch
+
+    # sample noise
+    z = torch.randn(imgs.shape[0], 100)
+    z = z.type_as(imgs)
+
+    # train discriminator
+    imgs_gen = self.generator(z)
+    logits = self.discriminator(imgs_gen.detach())
+    fake = torch.zeros(imgs.size(0), 1)
+    fake = fake.type_as(imgs)
+    loss_fake = trainer.criterion(logits, fake)
+
+    valid = torch.ones(imgs.size(0), 1)
+    valid = valid.type_as(imgs)
+    logits = self.discriminator(imgs)
+    loss_real = trainer.criterion(logits, valid)
+    loss_disc = (loss_real + loss_fake) / 2
+
+    # step dicriminator
+    _, _ = self.scaled_backward(loss_disc, None, trainer, trainer.optimizer[0])
+
+    if trainer.total_steps_done % trainer.grad_accum_steps == 0:
+        trainer.optimizer[0].step()
+        trainer.optimizer[0].zero_grad()
+
+    # train generator
+    imgs_gen = self.generator(z)
+
+    valid = torch.ones(imgs.size(0), 1)
+    valid = valid.type_as(imgs)
+
+    logits = self.discriminator(imgs_gen)
+    loss_gen = trainer.criterion(logits, valid)
+
+    # step generator
+    _, _ = self.scaled_backward(loss_gen, None, trainer, trainer.optimizer[1])
+    if trainer.total_steps_done % trainer.grad_accum_steps == 0:
+        trainer.optimizer[1].step()
+        trainer.optimizer[1].zero_grad()
+    return {"model_outputs": logits}, {"loss_gen": loss_gen, "loss_disc": loss_disc}
+
+...
+```
+
 See the [GAN training example](examples/train_simple_gan.py) with Gradient Accumulation
 
 
@@ -51,6 +105,18 @@ We don't use ```.spawn()``` to initiate multi-gpu training since it causes certa
 - ```.spawn()``` trains the model in subprocesses and the model in the main process is not updated.
 - DataLoader with N processes gets really slow when the N is large.
 
+## Training with [Accelerate](https://huggingface.co/docs/accelerate/index)
+
+Setting `use_accelerate` in `TrainingArgs` to `True` will enable training with Accelerate.
+
+You can also use it for multi-gpu or distributed training.
+
+```console
+CUDA_VISIBLE_DEVICES="0,1,2" accelerate launch --multi_gpu --num_processes 3 train_recipe_autoregressive_prompt.py
+```
+
+See the [Accelerate docs](https://huggingface.co/docs/accelerate/basic_tutorials/launch).
+
 ## Adding a callback
 👟 Supports callbacks to customize your runs. You can either set callbacks in your model implementations or give them
 explicitly to the Trainer.

requirements.dev.txt

@@ -3,3 +3,4 @@ coverage
 isort
 pytest
 pylint
+accelerate  # for testing

tests/test_train_gan.py

@@ -1,5 +1,6 @@
 import os
 from dataclasses import dataclass
+from typing import Any, Dict, Tuple
 
 import numpy as np
 import torch
@@ -159,6 +160,103 @@ def get_data_loader(
     assert loss_g1 > loss_g2, f"Generator loss should decrease. {loss_g1} > {loss_g2}"
 
 
+def test_overfit_accelerate_mnist_simple_gan():
+    @dataclass
+    class GANModelConfig(TrainerConfig):
+        epochs: int = 1
+        print_step: int = 2
+        training_seed: int = 666
+
+    class GANModel(TrainerModel):
+        def __init__(self):
+            super().__init__()
+            data_shape = (1, 28, 28)
+            self.generator = Generator(latent_dim=100, img_shape=data_shape)
+            self.discriminator = Discriminator(img_shape=data_shape)
+
+        def forward(self, x):
+            ...
+
+        def train_step(self, batch, criterion, optimizer_idx):
+            imgs, _ = batch
+
+            # sample noise
+            z = torch.randn(imgs.shape[0], 100)
+            z = z.type_as(imgs)
+
+            # train discriminator
+            if optimizer_idx == 0:
+                imgs_gen = self.generator(z)
+                logits = self.discriminator(imgs_gen.detach())
+                fake = torch.zeros(imgs.size(0), 1)
+                fake = fake.type_as(imgs)
+                loss_fake = criterion(logits, fake)
+
+                valid = torch.ones(imgs.size(0), 1)
+                valid = valid.type_as(imgs)
+                logits = self.discriminator(imgs)
+                loss_real = loss = criterion(logits, valid)
+                loss = (loss_real + loss_fake) / 2
+                return {"model_outputs": logits}, {"loss": loss}
+
+            # train generator
+            if optimizer_idx == 1:
+                imgs_gen = self.generator(z)
+
+                valid = torch.ones(imgs.size(0), 1)
+                valid = valid.type_as(imgs)
+
+                logits = self.discriminator(imgs_gen)
+                loss_real = criterion(logits, valid)
+                return {"model_outputs": logits}, {"loss": loss_real}
+
+        @torch.no_grad()
+        def eval_step(self, batch, criterion, optimizer_idx):
+            return self.train_step(batch, criterion, optimizer_idx)
+
+        def get_optimizer(self):
+            discriminator_optimizer = torch.optim.Adam(self.discriminator.parameters(), lr=0.0001, betas=(0.5, 0.999))
+            generator_optimizer = torch.optim.Adam(self.generator.parameters(), lr=0.001, betas=(0.5, 0.999))
+            return [discriminator_optimizer, generator_optimizer]
+
+        def get_criterion(self):
+            return nn.BCELoss()
+
+        def get_data_loader(
+            self, config, assets, is_eval, samples, verbose, num_gpus, rank=0
+        ):  # pylint: disable=unused-argument
+            transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
+            dataset = MNIST(os.getcwd(), train=not is_eval, download=True, transform=transform)
+            dataset.data = dataset.data[:64]
+            dataset.targets = dataset.targets[:64]
+            dataloader = DataLoader(dataset, batch_size=config.batch_size, drop_last=True, shuffle=False)
+            return dataloader
+
+    config = GANModelConfig()
+    config.batch_size = 64
+    config.grad_clip = None
+    config.training_seed = 333
+
+    model = GANModel()
+    trainer = Trainer(
+        TrainerArgs(use_accelerate=True), config, model=model, output_path=os.getcwd(), gpu=0 if is_cuda else None
+    )
+
+    trainer.eval_epoch()
+    loss_d1 = trainer.keep_avg_eval["avg_loss_0"]
+    loss_g1 = trainer.keep_avg_eval["avg_loss_1"]
+
+    trainer.config.epochs = 5
+    trainer.fit()
+    loss_d2 = trainer.keep_avg_train["avg_loss_0"]
+    loss_g2 = trainer.keep_avg_train["avg_loss_1"]
+
+    print(f"loss_d1: {loss_d1}, loss_d2: {loss_d2}")
+    print(f"loss_g1: {loss_g1}, loss_g2: {loss_g2}")
+    assert loss_d1 > loss_d2, f"Discriminator loss should decrease. {loss_d1} > {loss_d2}"
+    assert loss_g1 > loss_g2, f"Generator loss should decrease. {loss_g1} > {loss_g2}"
+
+
 def test_overfit_manual_optimize_mnist_simple_gan():
     @dataclass
     class GANModelConfig(TrainerConfig):
@@ -390,7 +488,131 @@ def get_data_loader(
     assert loss_g1 > loss_g2, f"Generator loss should decrease. {loss_g1} > {loss_g2}"
 
 
+def test_overfit_manual_accelerate_optimize_grad_accum_mnist_simple_gan():
+    @dataclass
+    class GANModelConfig(TrainerConfig):
+        epochs: int = 1
+        print_step: int = 2
+        training_seed: int = 666
+
+    class GANModel(TrainerModel):
+        def __init__(self):
+            super().__init__()
+            data_shape = (1, 28, 28)
+            self.generator = Generator(latent_dim=100, img_shape=data_shape)
+            self.discriminator = Discriminator(img_shape=data_shape)
+
+        def train_step():
+            ...
+
+        def forward(self, x):
+            ...
+
+        def optimize(self, batch, trainer):
+            imgs, _ = batch
+
+            # sample noise
+            z = torch.randn(imgs.shape[0], 100)
+            z = z.type_as(imgs)
+
+            # train discriminator
+            imgs_gen = self.generator(z)
+            logits = self.discriminator(imgs_gen.detach())
+            fake = torch.zeros(imgs.size(0), 1)
+            fake = fake.type_as(imgs)
+            loss_fake = trainer.criterion(logits, fake)
+
+            valid = torch.ones(imgs.size(0), 1)
+            valid = valid.type_as(imgs)
+            logits = self.discriminator(imgs)
+            loss_real = trainer.criterion(logits, valid)
+            loss_disc = (loss_real + loss_fake) / 2
+
+            # step dicriminator
+            self.scaled_backward(loss_disc, trainer, trainer.optimizer[0])
+
+            if trainer.total_steps_done % trainer.grad_accum_steps == 0:
+                trainer.optimizer[0].step()
+                trainer.optimizer[0].zero_grad()
+
+            # train generator
+            imgs_gen = self.generator(z)
+
+            valid = torch.ones(imgs.size(0), 1)
+            valid = valid.type_as(imgs)
+
+            logits = self.discriminator(imgs_gen)
+            loss_gen = trainer.criterion(logits, valid)
+
+            # step generator
+            self.scaled_backward(loss_gen, trainer, trainer.optimizer[1])
+            if trainer.total_steps_done % trainer.grad_accum_steps == 0:
+                trainer.optimizer[1].step()
+                trainer.optimizer[1].zero_grad()
+            return {"model_outputs": logits}, {"loss_gen": loss_gen, "loss_disc": loss_disc}
+
+        @torch.no_grad()
+        def eval_step(self, batch, criterion):
+            imgs, _ = batch
+
+            # sample noise
+            z = torch.randn(imgs.shape[0], 100)
+            z = z.type_as(imgs)
+
+            imgs_gen = self.generator(z)
+            valid = torch.ones(imgs.size(0), 1)
+            valid = valid.type_as(imgs)
+
+            logits = self.discriminator(imgs_gen)
+            loss_gen = trainer.criterion(logits, valid)
+            return {"model_outputs": logits}, {"loss_gen": loss_gen}
+
+        def get_optimizer(self):
+            discriminator_optimizer = torch.optim.Adam(self.discriminator.parameters(), lr=0.0001, betas=(0.5, 0.999))
+            generator_optimizer = torch.optim.Adam(self.generator.parameters(), lr=0.001, betas=(0.5, 0.999))
+            return [discriminator_optimizer, generator_optimizer]
+
+        def get_criterion(self):
+            return nn.BCELoss()
+
+        def get_data_loader(
+            self, config, assets, is_eval, samples, verbose, num_gpus, rank=0
+        ):  # pylint: disable=unused-argument
+            transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
+            dataset = MNIST(os.getcwd(), train=not is_eval, download=True, transform=transform)
+            dataset.data = dataset.data[:64]
+            dataset.targets = dataset.targets[:64]
+            dataloader = DataLoader(dataset, batch_size=config.batch_size, drop_last=True, shuffle=True)
+            return dataloader
+
+    config = GANModelConfig()
+    config.batch_size = 64
+    config.grad_clip = None
+
+    model = GANModel()
+    trainer = Trainer(
+        TrainerArgs(use_accelerate=True), config, model=model, output_path=os.getcwd(), gpu=0 if is_cuda else None
+    )
+
+    trainer.config.epochs = 1
+    trainer.fit()
+    loss_d1 = trainer.keep_avg_train["avg_loss_disc"]
+    loss_g1 = trainer.keep_avg_train["avg_loss_gen"]
+
+    trainer.config.epochs = 5
+    trainer.fit()
+    loss_d2 = trainer.keep_avg_train["avg_loss_disc"]
+    loss_g2 = trainer.keep_avg_train["avg_loss_gen"]
+
+    print(f"loss_d1: {loss_d1}, loss_d2: {loss_d2}")
+    print(f"loss_g1: {loss_g1}, loss_g2: {loss_g2}")
+    assert loss_d1 > loss_d2, f"Discriminator loss should decrease. {loss_d1} > {loss_d2}"
+    assert loss_g1 > loss_g2, f"Generator loss should decrease. {loss_g1} > {loss_g2}"
+
+
 if __name__ == "__main__":
     test_overfit_mnist_simple_gan()
+    test_overfit_accelerate_mnist_simple_gan()
     test_overfit_manual_optimize_mnist_simple_gan()
     test_overfit_manual_optimize_grad_accum_mnist_simple_gan()
+    test_overfit_manual_accelerate_optimize_grad_accum_mnist_simple_gan()

trainer/generic_utils.py

@@ -13,7 +13,7 @@ def isimplemented(obj, method_name):
     """Check if a method is implemented in a class."""
     if method_name in dir(obj) and callable(getattr(obj, method_name)):
         try:
-            obj.__getattribute__(method_name)()  # pylint: disable=unnecessary-dunder-call
+            obj.__getattribute__(method_name)()  # pylint: disable=bad-option-value, unnecessary-dunder-call
         except NotImplementedError:
             return False
         except:  # pylint: disable=bare-except