This template is used  for GAN structure.

Github Link:

Guide List:

Jdit is a research processing oriented framework based on pytorch. Only care about your ideas. You don't need to build a long boring code to run a deep learning project to verify your ideas.

You only need to implement you ideas and don't do anything with training framework, multiply-gpus, checkpoint, process visualization, performance evaluation and so on.

If you have any problems, or you find bugs you can contact the author.


Template of GAN

For GANs, generally, two things are needed which are generative model and discriminator model. This feature is different from other tasks. So, our template need to these two models training.

Class SupGanTrainer. The key method application like this:

def train_epoch(self, subbar_disable=False):
    for iteration, batch in tqdm(enumerate(self.datasets.loader_train, 1), 
        unit="step", disable=subbar_disable):
        self.step += 1
        self.input, self.ground_truth = self.get_data_from_batch(batch, 
        self.fake = self.netG(self.input)
        self._train_iteration(self.optD, self.compute_d_loss, 
        if (self.step % self.d_turn) == 0:
            self._train_iteration(self.optG, self.compute_g_loss, 
        if iteration == 1:
  • This loop is used for one epoch training. Inside this loop, it controls the training process of self.netG(generator) and self.netD(discriminator).

  • self.d_turn controls the turns of training, such as when you train a WGAN-GP, d_trun will be set 5 or 10, to train netD 5 turns and the train netG one turn.

  • self._watch_images("Train") this method is used to show images in tensorboard.
    The default behavior is to show input image and fake image (the output of netG).

  • For self.compute_g_loss() and self.compute_d_loss() in this template, we don't need to apply for this method. It will be come true by inherit templates.

Generally, you don't need to change this template. Because it is a structure for GAN. But, the thing is not absolute. If you need something new structure then you need to  inherit and rewrite this template. I encourage you to do that to creator your own trainer.

If your find this template is good enough that you don't need to change anything. It will be good by using the next level templates, such as Pix2pixGanTrainer, GenerateGanTrainer, and using your datasets, models, optimizers. If you change anyone of them, you still can reuse this template. This is the aim of this framework.

For the next level templates: