Hello, I’m working on integrating pix2pix into a framework of various generative models but running into issues replicating results from the original implementation.
I’ve essentially copied over the implementation from here 12 into the same object oriented structure as the rest of my models. Then I’ve added in functionality to create the multiscale generator and discriminator architecture from the pix2pixHD paper.
However, when I train the model with those features disabled (which I believe is identical to the original pix2pix implementation) the results are much worse than the repo linked above.
Both my discriminator and generator have the exact same number of parameters and structure as the original repo (although there are a couple modules wrapped around them and named slightly differently)


Generator
DataParallel(
  (module): MultiscaleGenerator(
    (mainG): UnetGenerator(
      (model): UnetSkipConnectionBlock(
        (model): Sequential(
          (0): Conv2d(3, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
          (1): UnetSkipConnectionBlock(
            (model): Sequential(
              (0): LeakyReLU(negative_slope=0.2, inplace)
              (1): Conv2d(64, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
              (2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
              (3): UnetSkipConnectionBlock(
                (model): Sequential(
                  (0): LeakyReLU(negative_slope=0.2, inplace)
                  (1): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
                  (2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
                  (3): UnetSkipConnectionBlock(
                    (model): Sequential(
                      (0): LeakyReLU(negative_slope=0.2, inplace)
                      (1): Conv2d(256, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
                      (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
                      (3): UnetSkipConnectionBlock(
                        (model): Sequential(
                          (0): LeakyReLU(negative_slope=0.2, inplace)
                          (1): Conv2d(512, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
                          (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
                          (3): UnetSkipConnectionBlock(
                            (model): Sequential(
                              (0): LeakyReLU(negative_slope=0.2, inplace)
                              (1): Conv2d(512, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
                              (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
                              (3): UnetSkipConnectionBlock(
                                (model): Sequential(
                                  (0): LeakyReLU(negative_slope=0.2, inplace)
                                  (1): Conv2d(512, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
                                  (2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
                                  (3): UnetSkipConnectionBlock(
                                    (model): Sequential(
                                      (0): LeakyReLU(negative_slope=0.2, inplace)
                                      (1): Conv2d(512, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
                                      (2): ReLU(inplace)
                                      (3): ConvTranspose2d(512, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
                                      (4): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
                                    )
                                  )
                                  (4): ReLU(inplace)
                                  (5): ConvTranspose2d(1024, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
                                  (6): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
                                  (7): Dropout(p=0.5)
                                )
                              )
                              (4): ReLU(inplace)
                              (5): ConvTranspose2d(1024, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
                              (6): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
                              (7): Dropout(p=0.5)
                            )
                          )
                          (4): ReLU(inplace)
                          (5): ConvTranspose2d(1024, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
                          (6): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
                          (7): Dropout(p=0.5)
                        )
                      )
                      (4): ReLU(inplace)
                      (5): ConvTranspose2d(1024, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
                      (6): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
                    )
                  )
                  (4): ReLU(inplace)
                  (5): ConvTranspose2d(512, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
                  (6): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
                )
              )
              (4): ReLU(inplace)
              (5): ConvTranspose2d(256, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
              (6): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
            )
          )
          (2): ReLU(inplace)
          (3): ConvTranspose2d(128, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
          (4): Tanh()
        )
      )
    )
  )
)
[Network G] Total number of parameters : 54.414 M




Discriminator
DataParallel(
  (module): MultiscaleDiscriminator(
    (discriminator_0): NLayerDiscriminator(
      (model): Sequential(
        (layer_0): Sequential(
          (0): Conv2d(6, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
          (1): LeakyReLU(negative_slope=0.2, inplace)
        )
        (layer_1): Sequential(
          (0): Conv2d(64, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): LeakyReLU(negative_slope=0.2, inplace)
        )
        (layer_2): Sequential(
          (0): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): LeakyReLU(negative_slope=0.2, inplace)
        )
        (layer_3): Sequential(
          (0): Conv2d(256, 512, kernel_size=(4, 4), stride=(1, 1), padding=(1, 1), bias=False)
          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): LeakyReLU(negative_slope=0.2, inplace)
        )
        (final_conv): Conv2d(512, 1, kernel_size=(4, 4), stride=(1, 1), padding=(1, 1))
        (sigmoid): Sigmoid()
      )
    )
  )
)
[Network D] Total number of parameters : 2.769 M


However my forward() code for the discriminator is different (to allow for extracting intermediate features for the discriminator feature loss from pix2pixHD). I think this is most likely where the issue is as the generator loss values are very similar to the original implementation while the discriminator’s losses are about an order of magnitude smaller.


My training code
for epoch in range(start_epoch, num_epochs + epochs_decay + 1):
    for i, data in enumerate(dataloader):
        # get data
        real_A = data['A'].to(self.device)
        real_B = data['B'].to(self.device)
        fake_B = self.G(real_A)

        # optimize discriminator
        self.set_requires_grad(self.D, True)
        self.optimizer_D.zero_grad()

        fake_AB = th.cat((real_A, fake_B), 1)
        pred_fake = self.D(fake_AB.detach())

        loss_D_fake = 0
        for preds in pred_fake:
            loss_D_fake += self.loss_GAN(preds, self.fake_label.expand_as(preds))

        real_AB = th.cat((real_A, real_B), 1)
        pred_real = self.D.get_features(real_AB)

        loss_D_real = 0
        for preds in self.pred_real:
            loss_D_real += self.loss_GAN(preds[-1], self.real_label.expand_as(preds[-1]))

        loss_D = (loss_D_fake + loss_D_real) * 0.5
        loss_D.backward()
        self.optimizer_D.step()

        # optimize generator
        self.set_requires_grad(self.D, False)
        self.optimizer_G.zero_grad()
        
        fake_AB = th.cat((real_A, fake_B), 1)
        pred_fake = self.D.get_features(fake_AB)

        loss_G_GAN = 0
        for preds in pred_fake:
            loss_G_GAN += self.loss_GAN(preds[-1], self.fake_label.expand_as(preds[-1]))

        loss_G_L1 = self.loss_L1(fake_B, real_B) * self.lambda_L1

        loss_G = loss_G_GAN + loss_G_L1

       if self.feature_loss: # false in this case
            # do stuff with pred_real and pred_fake here

        loss_G.backward()
        self.optimizer_G.step()

    def get_features(self, input):
        # forward with input downsampled for each scale
        # return list of lists of activations at each layer for each discriminator scale
        results = []
        for scale in range(self.n_scales): # in this case n_scales = 1
            # scale input for given discrim scale
            prev_output = interpolate(input, scale_factor=2**(scale - self.n_scales + 1))

            # get separate layers of discriminator
            discriminator = getattr(self, 'discriminator_%s'%(scale)).model
            layers = {k:v for (k,v) in discriminator.named_children() if 'layer' in k}

            # forward through the network storing discriminator features by layer
            per_scale_results = []
            for idx in range(len(layers)):
                layer = layers['layer_%s'%idx]
                prev_output = layer(prev_output)
                per_scale_results.append(prev_output)

            final_preds = getattr(discriminator, 'final_conv')(per_scale_results[-1])
            per_scale_results.append(final_preds)

            if self.use_sigmoid: # true in this case
                final_preds = getattr(discriminator, 'sigmoid')(final_preds)
                per_scale_results.append(final_preds)

            results.append(per_scale_results)

        # results look like this:
        # [[torch.Size([1, 64, 128, 128]), torch.Size([1, 128, 64, 64]), torch.Size([1, 256, 32, 32]), torch.Size([1, 512, 31, 31]), torch.Size([1, 1, 30, 30]), torch.Size([1, 1, 30, 30])]]
        # the last element in the list are the final predictions in the same format as the original pix2pix repo
        return results

    def forward(self, input):
        # forward with input downsampled for each scale
        # return list of predictions for each discriminator scale
        results = []
        for scale in range(self.n_scales):
            discriminator = getattr(self, 'discriminator_%s'%(scale))
            downsampled_input = interpolate(input, scale_factor=2**(scale - self.n_scales + 1))
            predictions = discriminator(downsampled_input)
            results.append(predictions)

        # results look like this: [[torch.Size([1, 1, 30, 30])]]
        return results




Original training code
for epoch in range(opt.epoch_count, opt.niter + opt.niter_decay + 1):
    for i, data in enumerate(dataset):
        self.real_A = data['A'].to(self.device)
        self.real_B = data['B'].to(self.device)
        self.fake_B = self.netG(self.real_A)

        # update D
        self.set_requires_grad(self.netD, True)
        self.optimizer_D.zero_grad()

        fake_AB = torch.cat((self.real_A, self.fake_B), 1)
        pred_fake = self.netD(fake_AB.detach())
        self.loss_D_fake = self.criterionGAN(pred_fake, False)

        real_AB = torch.cat((self.real_A, self.real_B), 1)
        pred_real = self.netD(real_AB)
        self.loss_D_real = self.criterionGAN(pred_real, True)

        self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5

        self.loss_D.backward()
        self.optimizer_D.step()

        # update G
        self.set_requires_grad(self.netD, False)
        self.optimizer_G.zero_grad()

        fake_AB = torch.cat((self.real_A, self.fake_B), 1)
        pred_fake = self.netD(fake_AB)
        self.loss_G_GAN = self.criterionGAN(pred_fake, True)

        self.loss_G_L1 = self.criterionL1(self.fake_B, self.real_B) * self.opt.lambda_L1

        self.loss_G = self.loss_G_GAN + self.loss_G_L1

        self.loss_G.backward()
        self.optimizer_G.step()


loss_GAN and criterionGAN from the respective implementations are nn.BCELoss()
I’ve verified that the data is being normalized in the exact same way and are identical before being given to forward()
Just to give an idea of the difference in results:
After 45 epochs mine looks like this:

sample_45.jpg776×260 96.6 KB

while the original implementation looks like this:

echte_pix2pix_45.png802×336 262 KB

Also I tried loading the checkpoint of the original implementation from epoch 35 and continuing training with my own implementation. After 1 epoch it already looks very faded:

sample_1.png776×260 224 KB

Compared with the checkpoint it’s loaded from:

After 10 epochs of continued training with my implementation it essentially looks the same as having trained it from scratch

sample_10.png776×260 247 KB

Clearly my implementation is doing something wrong, but I can’t for the life of me figure out what’s going wrong. Anyone have an idea what the issue could be?
The full codebase can be found here 9