I have trained my implementation 2 of Pix2Pix on the face2comics dataset and although the generated images are sharp and realistic, they are too bright. A tanh activation in the last layer of the generator outputs the generated images in the range [-1, 1]. The training images are normalized to have zero mean and std of one by computing the statistics of the training set.

fake776×260 69.8 KB


comic776×260 256 KB

However, when I check if the training and output images are in those ranges, only the generated images are in the desired range and individual training images are well above and/or below the [-1, 1] range. It is only when I compute the statistics of the entire training set after normalization that I get zero mean and unit variance.
Is this behaviour expected? How can I fix the “saturation” of the generated images?
Dataset module:
import os
import torch

from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from tqdm import tqdm


DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
torch.cuda.empty_cache()


def mean_std(dataset):
    """Return the mean and std of the dataset."""
    loader = DataLoader(dataset, batch_size=128, num_workers=0, shuffle=False)

    mean_inputs = 0.
    std_inputs = 0.
    mean_targets = 0.
    std_targets = 0.

    for inputs, targets in tqdm(loader):

        inputs = inputs.to(DEVICE).view(inputs.size(0), inputs.size(1), -1)
        mean_inputs += inputs.mean(2).sum(0)
        std_inputs += inputs.std(2).sum(0)

        targets = targets.to(DEVICE).view(targets.size(0), inputs.size(1), -1)
        mean_targets += targets.mean(2).sum(0)
        std_targets += targets.std(2).sum(0)

    mean_inputs /= len(loader.dataset)
    std_inputs /= len(loader.dataset)
    mean_targets /= len(loader.dataset)
    std_targets /= len(loader.dataset)

    return (mean_inputs, std_inputs), (mean_targets, std_targets)


class Face2Comic(Dataset):
    """A paired face-to-comics dataset."""

    def __init__(self, data_dir, train=True):
        super(Face2Comic, self).__init__()
        self.data_dir = data_dir
        self.faces_dir = os.path.join(data_dir, "faces")
        self.faces = os.listdir(self.faces_dir)
        self.comics_dir = os.path.join(data_dir, "comics")
        self.comics = os.listdir(self.comics_dir)
        self.len = len(self.faces)
        self.train = train

    def apply_transforms(self, face, comic):
        """Apply the same transforms to the input and the target."""
        common_transform = transforms.Compose([transforms.Resize((256, 256)),
                                               transforms.ToTensor()])

        normalize_face = transforms.Normalize(mean=[0.5129, 0.4136, 0.3671],
                                              std=[0.2372, 0.1972, 0.1883])

        normalize_comic = transforms.Normalize(mean=[0.4445, 0.3650, 0.3226],
                                               std=[0.2594, 0.2051, 0.1840])

        face = normalize_face(common_transform(face))
        comic = normalize_comic(common_transform(comic))

        if self.train:
            train_transform = transforms.Compose([
                transforms.RandomHorizontalFlip(p=0.5),
                transforms.ColorJitter()])
            face = train_transform(face)
            comic = train_transform(comic)

        return face, comic

    def __len__(self):
        """Get the number of samples in the dataset."""
        return self.len

    def __getitem__(self, index):
        """Return the transformed input (face) and target (comic)."""
        face = Image.open(os.path.join(self.faces_dir, self.faces[index]))
        comic = Image.open(os.path.join(self.comics_dir, self.comics[index]))
        return self.apply_transforms(face, comic)


if __name__ == '__main__':

    data_dir_train = os.getcwd() + '\\data\\train\\'
    dataset_train = Face2Comic(data_dir=data_dir_train, train=True)

    stats_faces, stats_comics = mean_std(dataset_train)

    print(f"Faces: mean = {stats_faces[0]}, std = {stats_faces[1]}")
    print(f"Comics: mean = {stats_comics[0]}, std = {stats_comics[1]}")

    data_dir_val = os.getcwd() + '\\data\\val\\'
    dataset_val = Face2Comic(data_dir=data_dir_val, train=False)

    stats_faces, stats_comics = mean_std(dataset_val)

    print(f"Faces: mean = {stats_faces[0]}, std = {stats_faces[1]}")
    print(f"Comics: mean = {stats_comics[0]}, std = {stats_comics[1]}")

Model module:
import torch
import torch.nn as nn


DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
torch.cuda.empty_cache()


def make_conv(in_size, out_size, encode, batch_norm, activation, drop_out):
    """Convolutional blocks of the Generator and the Discriminator.

    Let Ck denote a Convolution-BtachNorm-ReLU block with k filters.
    CDk denotes a Convolution-BtachNorm-Dropout-ReLU block with 50% dropout.
    All convolutions are 4 x 4 spatial filters with stride 2. Convolutions in
    the encoder and discriminator downsample by a factor of 2, whereas in the
    decoder they upsample by a factor of 2.
    """
    block = [nn.Conv2d(in_size, out_size,
                       kernel_size=4, stride=2, padding=1,
                       padding_mode="reflect",
                       bias=False if batch_norm else True)
             if encode else
             nn.ConvTranspose2d(in_size, out_size,
                                kernel_size=4, stride=2, padding=1,
                                bias=False if batch_norm else True)]

    if batch_norm:
        block.append(nn.BatchNorm2d(out_size))
    if activation == "leaky":
        block.append(nn.LeakyReLU(0.2))
    elif activation == "sigmoid":
        block.append(nn.Sigmoid())
    elif activation == "tanh":
        block.append(nn.Tanh())
    elif activation == "relu":
        block.append(nn.ReLU())
    if drop_out:
        block.append(nn.Dropout(0.5))

    return nn.Sequential(*block)


def init_weights(model, mean=0.0, std=0.02):
    """Initialize weights from a Gaussian distribution."""
    for module in model.modules():
        if isinstance(module, (nn.Conv2d, nn.BatchNorm2d)):
            nn.init.normal_(module.weight.data, mean=mean, std=std)


class Generator(nn.Module):
    """UNet Generator architecture.

    encoder:
        C64-C128-C256-C512-C512-C512-C512-C512
    decoder:
        CD512-CD1024-CD1024-C1024-C1024-C512-C256-C128

    After the C128 block in the decoder, a convolution is applied to map to the
    number of output channels, followed by a Tanh function. BatchNorm is not
    applied to the C64 block in the encoder. All ReLUs in the econder are
    leaky with slope 0.2, while ReLUs in the decoder are not leaky.
    """

    def __init__(self, in_channels=3, out_channels=3):
        super(Generator, self).__init__()

        encoder = [in_channels, 64, 128, 256, 512, 512, 512, 512, 512]
        encoder = zip(encoder, encoder[1:])

        self.encoder = nn.ModuleList()
        for idx, (input_size, output_size) in enumerate(encoder):
            if idx == 0:
                input_size *= 2
                batch_norm = False
            else:
                batch_norm = True
            self.encoder.append(make_conv(in_size=input_size,
                                          out_size=output_size,
                                          encode=True,
                                          batch_norm=batch_norm,
                                          activation="leaky",
                                          drop_out=False))

        decoder = [512, 1024, 1024, 1024, 1024, 512, 256, 128, out_channels]
        layers_decoder = len(decoder)
        decoder = zip(decoder, decoder[1:])

        self.decoder = nn.ModuleList()
        for idx, (input_size, output_size) in enumerate(decoder):
            if idx < layers_decoder - 2:
                batch_norm = True
                activation = "relu"
                output_size //= 2
            else:
                batch_norm = False
                activation = "tanh"
            self.decoder.append(make_conv(in_size=input_size,
                                          out_size=output_size,
                                          encode=False,
                                          batch_norm=batch_norm,
                                          activation=activation,
                                          drop_out=True if idx < 3 else False))

        init_weights(self, mean=0.0, std=0.02)

    def forward(self, x, z):
        """Generate a translation of x conditioned on the noise z."""
        x = torch.cat((x, z), dim=1)
        skip = [None]*len(self.encoder)

        for idx, block in zip(range(len(skip)-1, -1, -1), self.encoder):
            x = block(x)
            skip[idx] = x

        for idx, block in enumerate(self.decoder):
            if idx > 0:
                x = torch.cat((x, skip[idx]), dim=1)
            x = block(x)

        return x


class Discriminator(nn.Module):
    """C64-C128-C256-C512 PatchGAN Discriminator architecture.

    After the C512 block, a convolution is applied to map to a 1-d output,
    followed by a Sigmoid function. BatchNorm is not applied to the c64 block.
    All ReLUs are leaky with slope of 0.2.
    """

    def __init__(self, in_channels=3):
        super(Discriminator, self).__init__()

        channels = [in_channels, 64, 128, 256, 512, 1]
        layers = len(channels)
        channels = zip(channels, channels[1:])

        self.blocks = nn.ModuleList()
        for layer, (input_size, output_size) in enumerate(channels):
            if layer == 0:
                input_size *= 2
                batch_norm = False
                activation = "leaky"
            elif layer < layers - 2:
                batch_norm = True
                activation = "leaky"
            else:
                batch_norm = False
                activation = "sigmoid"
            self.blocks.append(make_conv(in_size=input_size,
                                         out_size=output_size,
                                         encode=True,
                                         batch_norm=batch_norm,
                                         activation=activation,
                                         drop_out=False))

        init_weights(self, mean=0.0, std=0.02)

    def forward(self, x, y):
        """Return a nxn tensor of patch probabilities."""
        x = torch.cat((x, y), dim=1)
        for block in self.blocks:
            x = block(x)
        return x


if __name__ == '__main__':

    batch_size = 8
    channels = 3
    height = 256
    width = 256

    x = torch.randn((batch_size, channels, height, width), device=DEVICE)
    y = torch.randn((batch_size, channels, height, width), device=DEVICE)
    z = torch.randn((batch_size, channels, height, width), device=DEVICE)

    generator = Generator().to(DEVICE)
    total_params = sum(p.numel() for p in generator.parameters())
    print(f"Number of parameters in Generator: {total_params:,}")

    G_z = generator(x, z)
    print(G_z.shape)

    discriminator = Discriminator().to(DEVICE)
    total_params = sum(p.numel() for p in discriminator.parameters())
    print(f"Number of parameters in Discriminator: {total_params:,}")

    D_x = discriminator(x, y)
    print(D_x.shape)