Hi, I am a newbie in PyTorch, GAN, and I don’t have much experience in Python (Although I am a C/C++ programmer).
I have a simple tutorial code for DCGAN for generating fake image, and it was ok when i run the code with “DATASETNAME = ‘MNIST’”. However, when i change the dataset to ‘CIFAR10’, the program produces error related to “running_mean”.
The code is as below

import torch.nn as nn


def weights_init(module):
    if isinstance(module, nn.Conv2d) or isinstance(module, nn.ConvTranspose2d):
        module.weight.detach().normal_(mean=0., std=0.02)

    elif isinstance(module, nn.BatchNorm2d):
        module.weight.detach().normal_(1., 0.02)
        module.bias.detach().zero_()

    else:
        pass


class View(nn.Module):
    def __init__(self, output_shape):
        super(View, self).__init__()
        self.output_shape = output_shape

    def forward(self, x):
        return x.view(x.shape[0], *self.output_shape)


class Generator(nn.Module):
    def __init__(self, dataset_name):
        super(Generator, self).__init__()
        act = nn.ReLU(inplace=True)
        norm = nn.BatchNorm2d

        if dataset_name == 'CIFAR10':  # Output shape 3x32x32
            model = [nn.Linear(100, 512 * 4 * 4), View([512, 4, 4]), norm(512), act]  # 4x4
            model += [nn.ConvTranspose2d(512, 256, 5, stride=2, padding=2, output_padding=1), norm(256), act]  # 8x8
            model += [nn.ConvTranspose2d(256, 128, 5, stride=2, padding=2, output_padding=1), norm(128), act]  # 16x16
            model += [nn.ConvTranspose2d(128, 3, 5, stride=2, padding=2, output_padding=1), nn.Tanh()]  # 32x32

        elif dataset_name == 'LSUN':  # Output shape 3x64x64
            model = [nn.Linear(100, 1024 * 4 * 4), View([1024, 4, 4]), norm(1024), act]  # 4x4
            model += [nn.ConvTranspose2d(1024, 512, 5, stride=2, padding=2, output_padding=1), norm(512), act]  # 8x8
            model += [nn.ConvTranspose2d(512, 256, 5, stride=2, padding=2, output_padding=1), norm(256), act]  # 16x16
            model += [nn.ConvTranspose2d(256, 128, 5, stride=2, padding=2, output_padding=1), norm(128), act]  # 32x32
            model += [nn.ConvTranspose2d(128, 3, 5, stride=2, padding=2, output_padding=1), nn.Tanh()]  # 64x64

        elif dataset_name == 'MNIST':  # Output shape 1x28x28
            model = [nn.Linear(100, 256 * 4 * 4), View([256, 4, 4]), norm(256), act]  # 4x4
            model += [nn.ConvTranspose2d(256, 128, 5, stride=2, padding=2), norm(128), act]  # 7x7
            model += [nn.ConvTranspose2d(128, 64, 5, stride=2, padding=2, output_padding=1), norm(64), act]  # 14x14
            model += [nn.ConvTranspose2d(64, 1, 5, stride=2, padding=2, output_padding=1), nn.Tanh()]  # 28x28

        else:
            raise NotImplementedError

        self.model = nn.Sequential(*model)

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


class Discriminator(nn.Module):
    def __init__(self, dataset_name):
        super(Discriminator, self).__init__()
        act = nn.LeakyReLU(inplace=True, negative_slope=0.2)
        norm = nn.BatchNorm2d

        if dataset_name == 'CIFAR10':  # Input shape 3x32x32
            model = [nn.Conv2d(3, 128, 5, stride=2, padding=2, bias=False), act]  # 16x16
            model += [nn.Conv2d(128, 256, 5, stride=2, padding=2, bias=False), norm(128), act]  # 8x8
            model += [nn.Conv2d(256, 512, 5, stride=2, padding=2, bias=False), norm(256), act]  # 4x4
            model += [nn.Conv2d(512, 1, 4, stride=2, padding=2, bias=False), nn.Sigmoid()]  # 1x1

        elif dataset_name == 'LSUN':  # Input shape 3x64x64
            model = [nn.Conv2d(3, 128, 5, stride=2, padding=2, bias=False), act]  # 128x32x32
            model += [nn.Conv2d(128, 256, 5, stride=2, padding=2, bias=False), norm(128), act]  # 256x16x16
            model += [nn.Conv2d(256, 512, 5, stride=2, padding=2, bias=False), norm(256), act]  # 512x8x8
            model += [nn.Conv2d(512, 1024, 5, stride=2, padding=2, bias=False), norm(512), act]  # 1024x4x4
            model += [nn.Conv2d(1024, 1, 4), nn.Sigmoid()]  # 1x1x1

        elif dataset_name == 'MNIST':  # Input shape 1x28x28
            model = [nn.Conv2d(1, 64, 5, stride=2, padding=2, bias=False), act]  # 14x14
            model += [nn.Conv2d(64, 128, 5, stride=2, padding=2, bias=False), norm(128), act]  # 7x7
            model += [nn.Conv2d(128, 256, 5, stride=2, padding=2, bias=False), norm(256), act]  # 4x4
            model += [nn.Conv2d(256, 1, 4, bias=False), nn.Sigmoid()]  # 1x1

        else:
            raise NotImplementedError

        self.model = nn.Sequential(*model)

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

if __name__ == '__main__':
    import os
    from torchvision.transforms import Compose, Normalize, Resize, ToTensor
    from torch.utils.data import DataLoader
    #from models import Discriminator, Generator, weights_init
    import torch
    import torch.nn as nn
    import matplotlib.pyplot as plt
    from time import time
    from tqdm import tqdm
    from torchvision.utils import save_image
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    

    BETA1, BETA2 = 0.5, 0.99
    BATCH_SIZE = 16
    DATASET_NAME = 'CIFAR10'
    DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu:0')
    EPOCHS = 1
    ITER_REPORT = 10
    LATENT_DIM = 100
    LR = 2e-4
    N_D_STEP = 1
    ITER_DISPLAY = 500

    IMAGE_DIR = './GAN/checkpoints/'+DATASET_NAME+'/Image'
    MODEL_DIR = './GAN/checkpoints/'+DATASET_NAME+'/Model'
    
    if DATASET_NAME == 'CIFAR10':
        IMAGE_SIZE = 32
        OUT_CHANNEL = 3
        from torchvision.datasets import CIFAR10
        transforms = Compose([ToTensor(), Normalize(mean=[0.5], std=[0.5])])
        dataset = CIFAR10(root='./datasets', train=True, transform=transforms, download=True)
    elif DATASET_NAME == 'LSUN':
        IMAGE_SIZE = 64
        OUT_CHANNEL = 3
        from torchvision.datasets import LSUN
        transforms = Compose([Resize(64), ToTensor(), Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])
        dataset = LSUN(root='./datasets/LSUN', classes=['bedroom_train'], transform=transforms)
    elif DATASET_NAME == 'MNIST':
        IMAGE_SIZE = 28
        OUT_CHANNEL = 1
        from torchvision.datasets import MNIST
        transforms = Compose([ToTensor(), Normalize(mean=[0.5], std=[0.5])])
        dataset = MNIST(root='./datasets', train=True, transform=transforms, download=True)
    else:
        raise NotImplementedError

    data_loader = DataLoader(dataset=dataset, batch_size=BATCH_SIZE, num_workers=0, shuffle=True)

    D = Discriminator(DATASET_NAME).apply(weights_init).to(DEVICE)
    G = Generator(DATASET_NAME).apply(weights_init).to(DEVICE)
    print(D, G)
    criterion = nn.BCELoss()

    optim_D = torch.optim.Adam(D.parameters(), lr=LR, betas=(BETA1, BETA2))
    optim_G = torch.optim.Adam(G.parameters(), lr=LR, betas=(BETA1, BETA2))

    list_D_loss = list()
    list_G_loss = list()
    total_step = 0

    st = time()
    for epoch in range(EPOCHS):
        for data in tqdm(data_loader):
            total_step += 1
            real, label = data[0].to(DEVICE), data[1].to(DEVICE)
            z = torch.randn(BATCH_SIZE, LATENT_DIM).to(DEVICE)

            fake = G(z)

            real_score = D(real)
            fake_score = D(fake.detach())

            D_loss = 0.5 * (criterion(fake_score, torch.zeros_like(fake_score).to(DEVICE))
                            + criterion(real_score, torch.ones_like(real_score).to(DEVICE)))
            optim_D.zero_grad()
            D_loss.backward()
            optim_D.step()
            list_D_loss.append(D_loss.detach().cpu().item())
            
            if total_step % ITER_DISPLAY == 0:
                #(BatchSize, Channel*ImageSize*ImageSize)-->(BatchSize, Channel, ImageSize, ImageSize)
                fake = fake.view(BATCH_SIZE, OUT_CHANNEL, IMAGE_SIZE, IMAGE_SIZE)
                real = real.view(BATCH_SIZE, OUT_CHANNEL, IMAGE_SIZE, IMAGE_SIZE)
                save_image(fake, IMAGE_DIR + '/{}_fake.png'.format(epoch + 1), nrow=4, normalize=True)
                save_image(real, IMAGE_DIR + '/{}_real.png'.format(epoch + 1), nrow=4, normalize=True)
            
            if total_step % N_D_STEP == 0:
                fake_score = D(fake)
                G_loss = criterion(fake_score, torch.ones_like(fake_score))
                optim_G.zero_grad()
                G_loss.backward()
                optim_G.step()
                list_G_loss.append(G_loss.detach().cpu().item())

                if total_step % ITER_REPORT == 0:
                    print("Epoch: {}, D_loss: {:.{prec}} G_loss: {:.{prec}}"
                          .format(epoch, D_loss.detach().cpu().item(), G_loss.detach().cpu().item(), prec=4))

    torch.save(D.state_dict(), '{}_D.pt'.format(DATASET_NAME))
    torch.save(G.state_dict(), '{}_G.pt'.format(DATASET_NAME))

    plt.figure()
    plt.plot(range(0, len(list_D_loss)), list_D_loss, linestyle='--', color='r', label='Discriminator loss')
    plt.plot(range(0, len(list_G_loss) * N_D_STEP, N_D_STEP), list_G_loss, linestyle='--', color='g',
             label='Generator loss')
    plt.xlabel('Iteration')
    plt.ylabel('Loss')
    plt.legend()
    plt.savefig('Loss.png')

    print(time() - st)
---------------------------------------------------------------------------------------------------------------------------------


The error seems to come from Discriminator .forward as follows:

---------------------------------------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-2-0e0e520d9b12> in <module>
     71             fake = G(z)
     72 
---> 73             real_score = D(real)
     74             fake_score = D(fake.detach())
     75 

C:\Anaconda3\lib\site-packages\torch\nn\modules\module.py in __call__(self, *input, **kwargs)
    491             result = self._slow_forward(*input, **kwargs)
    492         else:
--> 493             result = self.forward(*input, **kwargs)
    494         for hook in self._forward_hooks.values():
    495             hook_result = hook(self, input, result)

<ipython-input-1-5160dda25717> in forward(self, x)
     87 
     88     def forward(self, x):
---> 89         return self.model(x)

C:\Anaconda3\lib\site-packages\torch\nn\modules\module.py in __call__(self, *input, **kwargs)
    491             result = self._slow_forward(*input, **kwargs)
    492         else:
--> 493             result = self.forward(*input, **kwargs)
    494         for hook in self._forward_hooks.values():
    495             hook_result = hook(self, input, result)

C:\Anaconda3\lib\site-packages\torch\nn\modules\container.py in forward(self, input)
     90     def forward(self, input):
     91         for module in self._modules.values():
---> 92             input = module(input)
     93         return input
     94 

C:\Anaconda3\lib\site-packages\torch\nn\modules\module.py in __call__(self, *input, **kwargs)
    491             result = self._slow_forward(*input, **kwargs)
    492         else:
--> 493             result = self.forward(*input, **kwargs)
    494         for hook in self._forward_hooks.values():
    495             hook_result = hook(self, input, result)

C:\Anaconda3\lib\site-packages\torch\nn\modules\batchnorm.py in forward(self, input)
     81             input, self.running_mean, self.running_var, self.weight, self.bias,
     82             self.training or not self.track_running_stats,
---> 83             exponential_average_factor, self.eps)
     84 
     85     def extra_repr(self):

C:\Anaconda3\lib\site-packages\torch\nn\functional.py in batch_norm(input, running_mean, running_var, weight, bias, training, momentum, eps)
   1695     return torch.batch_norm(
   1696         input, weight, bias, running_mean, running_var,
-> 1697         training, momentum, eps, torch.backends.cudnn.enabled
   1698     )
   1699 

RuntimeError: running_mean should contain 256 elements not 128


Can anyone tell me what is this error about? It seems to come from size setting of something in the model, but that’s all I can guess.
Thank you in advance.