<jupyter_start><jupyter_text>DCGAN loss and training loop<jupyter_code>#hide
IN_COLAB = 'google.colab' in str(get_ipython())
#hide
if IN_COLAB:
  !pip3 install -Uqq fastbook
#hide
if IN_COLAB:
  from pathlib import Path
  from nbdev.imports import Config
  project_path = Path('/content/drive/My Drive/Colab Notebooks/github/dcgan')
  get_ipython().magic(f'cd {project_path}')
  get_ipython().magic(f'cd {Config().nbs_path}')
# default_exp train
#hide
from nbdev.showdoc import *
#export
from fastai.vision.all import *
from dcgan.data import *
from dcgan.models import *<jupyter_output><empty_output><jupyter_text>Loss Function<jupyter_code>#export
class DCGANLoss(nn.Module):

  def __init__(self, dcgan: nn.Module):
    super().__init__()
    store_attr()

    self.crit = BCEWithLogitsLossFlat()
  
  def set_input(self, noise, real): "set `self.noise` and `self.real` for future loss calculation"; self.noise,self.real = noise, real

  def forward(self, output, target):
      """
      Generator Loss
      """
      fake_image     = self.dcgan.G(self.noise)
      disc_fake_pred = self.dcgan.D(fake_image)
      self.gen_loss  = self.crit(disc_fake_pred, torch.ones_like(disc_fake_pred))

      return self.gen_loss<jupyter_output><empty_output><jupyter_text>Trainer<jupyter_code>#export
class DCGANTrainer(Callback):
    """`Learner` Callback for training a DCGAN model."""
    run_before = Recorder

    def __init__(self): pass

    def _set_trainable(self, disc=False):
        """Put the generators or discriminators in training mode depending on arguments."""
        def set_requires_grad(m, rg):
            for p in m.parameters(): p.requires_grad_(rg)
        set_requires_grad(self.learn.model.G, not disc)
        set_requires_grad(self.learn.model.D, disc)
        if disc: self.opt_D.hypers = self.learn.opt.hypers

    def before_train(self, **kwargs):
        self.G = self.learn.model.G
        self.D = self.learn.model.D
        self.crit = self.learn.loss_func.crit
        
        if not getattr(self,'opt_G',None):
            self.opt_G = self.learn.opt_func(self.learn.splitter(nn.Sequential(*flatten_model(self.G))), self.learn.lr)
        else:
            self.opt_G.hypers = self.learn.opt.hypers
        if not getattr(self, 'opt_D',None):
            self.opt_D = self.learn.opt_func(self.learn.splitter(nn.Sequential(*flatten_model(self.D))), self.learn.lr)
        else:
            self.opt_D.hypers = self.learn.opt.hypers

        self.learn.opt = self.opt_G

    def before_batch(self, **kwargs):
        self._set_trainable()
        self._training = self.learn.model.training
        self.learn.xb = (self.learn.xb[0],self.learn.yb[0])

        self.learn.loss_func.set_input(*self.learn.xb)

    def after_step(self):
        self.opt_D.hypers = self.learn.opt.hypers

    def after_batch(self, **kwargs):
        "Discriminator training loop"
        if self._training:
            # Obtain images
            (noise, real) = self.learn.xb
            self._set_trainable(disc=True)
            
            noise = torch.randn(noise.shape).to('cuda')

            fake_image = self.G(noise)
            
            disc_fake_pred = self.D(fake_image)
            disc_fake_loss = self.crit(disc_fake_pred, torch.zeros_like(disc_fake_pred))
            disc_real_pred = self.D(real)
            disc_real_loss = self.crit(disc_real_pred, torch.ones_like(disc_real_pred))
            disc_loss      = (disc_fake_loss + disc_real_loss) / 2
            
            disc_loss.backward()

            self.learn.loss_func.disc_loss = disc_loss.detach().cpu()
            
            # Optimizer stepping
            self.opt_D.step()
            self.opt_D.zero_grad()
            self._set_trainable()

    def before_validate(self, **kwargs):
        self.G    = self.learn.model.G
        self.D    = self.learn.model.D
        self.crit = self.learn.loss_func.crit
#export
class ShowDCGANImgsCallback(Callback):
    "Update the progress bar with input and prediction images"
    run_after,run_valid=DCGANTrainer,False

    
    def __init__(self): pass
        
    def before_fit(self):
        self.run = not hasattr(self.learn, 'lr_finder') and not hasattr(self, "gather_preds")
        self.nb_batches = []
        self.imgs = []
        self.titles = []
        self.steps = 0
        assert hasattr(self.learn, 'progress')

    def after_batch(self):
        if (self.steps > 0) and (self.steps % 1000 == 0):
          noise = torch.randn(1, 64, 1, 1).to('cuda')
          
          fake_image = self.learn.model.G(noise)
          img = TensorImage(self.learn.dls.after_batch.decode(TensorImage(fake_image))[0])

          self.imgs.append(img)
          self.titles.append(f'Step {self.steps}')
          self.progress.mbar.show_imgs(self.imgs, self.titles, imgsize=5, figsize=(8, 8))

        self.steps += 1
#export
@delegates(Learner.__init__)
def dcgan_learner(dls:DataLoader, m:DCGAN, opt_func=Adam, show_imgs:bool=True, metrics:list=[], cbs:list=[], **kwargs):
    
    lms = LossMetrics(['gen_loss', 'disc_loss'])
    learn = Learner(dls, m, loss_func=DCGANLoss(m), opt_func=opt_func,
                    cbs=[DCGANTrainer, *cbs], metrics=[*lms, *[AvgMetric(metric) for metric in [*metrics]]])
    
    if show_imgs: learn.add_cbs(ShowDCGANImgsCallback())
    learn.recorder.train_metrics = True
    learn.recorder.valid_metrics = False
    return learn
show_doc(dcgan_learner)<jupyter_output><empty_output><jupyter_text>Test<jupyter_code>#ignore
path = untar_data(URLs.MNIST)
dls  = get_dls(path, bs=128)
#ignore
dc_gan = DCGAN(ch_in=1,
               z_dim=64
               )

learn = dcgan_learner(dls, dc_gan, opt_func=partial(Adam, betas=(0.5, 0.999)))
#ignore
learn.lr_find()
#ignore
learn.fit(n_epoch=75, lr=2e-4)
#ignore
learn.show_results(ds_idx=0)
#ignore
from nbdev.export import notebook2script
notebook2script()
#ignore
from nbdev.cli import nbdev_build_lib
nbdev_build_lib()
#ignore
from nbdev.cli import nbdev_build_docs
nbdev_build_docs()
#ignore
from nbdev.cli import nbdev_test_nbs
nbdev_test_nbs()
#ignore
!git status
#ignore
!git add *.ipynb README.md docs/ ; git config --global user.email '<EMAIL>'; git config --global user.name 'numb3r33'; git commit -m 'Split cells.'
#ignore
!git pull origin main --no-edit
#ignore
!git push origin main
#ignore
from nbdev.cli import nbdev_clean_nbs
nbdev_clean_nbs()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12206134205200063, 'fraction_numerical': 0.006694690954382687, 'mean_word_length': 3.0351758793969847, 'pattern_counts': {'":': 0, '<': 18, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '553776', 'license': 'Apache-2.0', 'n_tokens_mistral': 2327, 'n_tokens_neox': 2227, 'n_words': 415, 'path': '02_train.ipynb', 'repo_name': 'numb3r33/dcgan', 'size': '6436'}