eriklindernoren / PyTorch-GAN

PyTorch implementations of Generative Adversarial Networks.
MIT License
16.22k stars 4.05k forks source link

Model Test in DualGAN ([B, 3, 256, 256] => Discriminator => [B, 1, 30, 30] , instead of [B, 1, 1, 1]) #111

Open Sean16SYSU opened 4 years ago

Sean16SYSU commented 4 years ago

I have tested the models in DualGAN and got something unexpected. Given input tensor whose shape is [B, 3, 256, 256], I put it into Discriminator in DualGAN and the shape of output tensor is [B, 1, 30, 30], instead of [B, 1, 1, 1].

The code I run will be list at the bottom, most of the code is copied from models.py in dualgan, except for the import and main.

the output of the code is shown as follows:

torch.Size([2, 3, 256, 256])
torch.Size([2, 3, 256, 256])
torch.Size([2, 1, 30, 30])

Code:

import torch.nn as nn
import torch.nn.functional as F
import torch
from torchvision.models import vgg19
import math
from dataloader import MyDataset
import torch.utils.data as Data

def weights_init_normal(m):
    classname = m.__class__.__name__
    if classname.find("Conv") != -1:
        torch.nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find("BatchNorm2d") != -1:
        torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
        torch.nn.init.constant_(m.bias.data, 0.0)

##############################
#           U-NET
##############################

class UNetDown(nn.Module):
    def __init__(self, in_size, out_size, normalize=True, dropout=0.0):
        super(UNetDown, self).__init__()
        layers = [nn.Conv2d(in_size, out_size, 4, stride=2, padding=1, bias=False)]
        if normalize:
            layers.append(nn.InstanceNorm2d(out_size, affine=True))
        layers.append(nn.LeakyReLU(0.2))
        if dropout:
            layers.append(nn.Dropout(dropout))
        self.model = nn.Sequential(*layers)

    def forward(self, x):
        return self.model(x)

class UNetUp(nn.Module):
    def __init__(self, in_size, out_size, dropout=0.0):
        super(UNetUp, self).__init__()
        layers = [
            nn.ConvTranspose2d(in_size, out_size, 4, stride=2, padding=1, bias=False),
            nn.InstanceNorm2d(out_size, affine=True),
            nn.ReLU(inplace=True),
        ]
        if dropout:
            layers.append(nn.Dropout(dropout))

        self.model = nn.Sequential(*layers)

    def forward(self, x, skip_input):
        x = self.model(x)
        x = torch.cat((x, skip_input), 1)

        return x

class Generator(nn.Module):
    def __init__(self, channels=3):
        super(Generator, self).__init__()

        self.down1 = UNetDown(channels, 64, normalize=False)
        self.down2 = UNetDown(64, 128)
        self.down3 = UNetDown(128, 256)
        self.down4 = UNetDown(256, 512, dropout=0.5)
        self.down5 = UNetDown(512, 512, dropout=0.5)
        self.down6 = UNetDown(512, 512, dropout=0.5)
        self.down7 = UNetDown(512, 512, dropout=0.5, normalize=False)

        self.up1 = UNetUp(512, 512, dropout=0.5)
        self.up2 = UNetUp(1024, 512, dropout=0.5)
        self.up3 = UNetUp(1024, 512, dropout=0.5)
        self.up4 = UNetUp(1024, 256)
        self.up5 = UNetUp(512, 128)
        self.up6 = UNetUp(256, 64)

        self.final = nn.Sequential(nn.ConvTranspose2d(128, channels, 4, stride=2, padding=1), nn.Tanh())

    def forward(self, x):
        # Propogate noise through fc layer and reshape to img shape
        d1 = self.down1(x)
        d2 = self.down2(d1)
        d3 = self.down3(d2)
        d4 = self.down4(d3)
        d5 = self.down5(d4)
        d6 = self.down6(d5)
        d7 = self.down7(d6)
        u1 = self.up1(d7, d6)
        u2 = self.up2(u1, d5)
        u3 = self.up3(u2, d4)
        u4 = self.up4(u3, d3)
        u5 = self.up5(u4, d2)
        u6 = self.up6(u5, d1)

        return self.final(u6)

##############################
#        Discriminator
##############################

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

        def discrimintor_block(in_features, out_features, normalize=True):
            """Discriminator block"""
            layers = [nn.Conv2d(in_features, out_features, 4, stride=2, padding=1)]
            if normalize:
                layers.append(nn.BatchNorm2d(out_features, 0.8))
            layers.append(nn.LeakyReLU(0.2, inplace=True))
            return layers

        self.model = nn.Sequential(
            *discrimintor_block(in_channels, 64, normalize=False),
            *discrimintor_block(64, 128),
            *discrimintor_block(128, 256),
            nn.ZeroPad2d((1, 0, 1, 0)),
            nn.Conv2d(256, 1, kernel_size=4)
        )

    def forward(self, img):
        return self.model(img)

if __name__ == '__main__':
    G = Generator(3)
    path = r'D:\Software\DataSet\CUHK_Face_Sketch\CUHK_training_photo\photo'
    dataset = MyDataset(path=path, resize=256, Len=10, img_type='jpg', sketch=False)
    train_loader = Data.DataLoader(dataset=dataset, batch_size=2, shuffle=True)

    D = Discriminator(3)
    for step, x in enumerate(train_loader):
        print(x.shape)
        print(G(x).shape)
        print(D(x).shape)
        break
AverageName commented 4 years ago

I don't know the details, but your discriminator might be PatchGAN, it looks at overlapping patches 70x70 and tells if this path is generated or not. So this is might be the case. Hope I've helped you :)

Sean16SYSU commented 4 years ago

I have implemented DualGAN now, and found that wgan works well only when D without BN layers.

D may not suitable for any size image. And we would modify slightly before applying it for your own dataset.

My fault.

Thanks and I will shut down this issue. And I am always welcomed for further communicated for GANs model implementation with Pytorch. Anyone is free to email me.

Sean16SYSU commented 4 years ago

@AverageName Thanks for your help. Actually, DualGAN ultilizes WGAN instead of PatchGAN which explained in CycleGAN. But now I feel this case reasonable, as the models may not suitable for any possible input of any size.