Pytorch GAN model doesn't train: matrix multiplication error
I'm trying to build a basic GAN to familiarise myself with Pytorch. I have some (limited) experience with Keras, but since I'm bound to do a larger project in Pytorch, I wanted to explore first using 'basic' networks.
I'm using Pytorch Lightning. I think I've added all necessary components. I tried passing some noise through the generator and the discriminator separately, and I think the output has the expected shape. Nonetheless, I get a runtime error when I try to train the GAN (full traceback below):
RuntimeError: mat1 and mat2 shapes cannot be multiplied (7x9 and 25x1)
I noticed that 7 is the size of the batch (by printing out the batch dimensions), even though I specified batch_size to be 64. Other than that, quite honestly, I don't know where to begin: the error traceback doesn't help me.
Chances are, I made multiple mistakes. However, I'm hoping some of you will be able to spot the current error from the code, since the multiplication error seems to point towards a dimensionality problem somewhere. Here's the code.
import os

import pytorch_lightning as pl
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
from skimage import io
from torch.utils.data import Dataset, DataLoader, random_split
from torchvision.utils import make_grid
from torchvision.transforms import Resize, ToTensor, ToPILImage, Normalize  

class DoppelDataset(Dataset):
    &quot;&quot;&quot;
    Dataset class for face data
    &quot;&quot;&quot;

    def __init__(self, face_dir: str, transform=None):

        self.face_dir = face_dir
        self.face_paths = os.listdir(face_dir)
        self.transform = transform

    def __len__(self):

        return len(self.face_paths)

    def __getitem__(self, idx):

        if torch.is_tensor(idx):
            idx = idx.tolist()

        face_path = os.path.join(self.face_dir, self.face_paths[idx])
        face = io.imread(face_path)

        sample = {'image': face}

        if self.transform:
            sample = self.transform(sample['image'])

        return sample


class DoppelDataModule(pl.LightningDataModule):

    def __init__(self, data_dir='../data/faces', batch_size: int = 64, num_workers: int = 0):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.num_workers = num_workers

        self.transforms = transforms.Compose([
            ToTensor(),
            Resize(100),
            Normalize(mean=(123.26290927634774, 95.90498110733365, 86.03763122875182),
                      std=(63.20679012922922, 54.86211954409834, 52.31266645797249))
        ])

    def setup(self, stage=None):
        # Initialize dataset
        doppel_data = DoppelDataset(face_dir=self.data_dir, transform=self.transforms)

        # Train/val/test split
        n = len(doppel_data)
        train_size = int(.8 * n)
        val_size = int(.1 * n)
        test_size = n - (train_size + val_size)

        self.train_data, self.val_data, self.test_data = random_split(dataset=doppel_data,
                                                                      lengths=[train_size, val_size, test_size])

    def train_dataloader(self) -&gt; DataLoader:
        return DataLoader(dataset=self.test_data, batch_size=self.batch_size, num_workers=self.num_workers)

    def val_dataloader(self) -&gt; DataLoader:
        return DataLoader(dataset=self.val_data, batch_size=self.batch_size, num_workers=self.num_workers)

    def test_dataloader(self) -&gt; DataLoader:
        return DataLoader(dataset=self.test_data, batch_size=self.batch_size, num_workers=self.num_workers)


class DoppelGenerator(nn.Sequential):
    &quot;&quot;&quot;
    Generator network that produces images based on latent vector
    &quot;&quot;&quot;

    def __init__(self, latent_dim: int):
        super().__init__()

        def block(in_channels: int, out_channels: int, padding: int = 1, stride: int = 2, bias=False):
            return nn.Sequential(
                nn.ConvTranspose2d(in_channels=in_channels, out_channels=out_channels, kernel_size=4, stride=stride,
                                   padding=padding, bias=bias),
                nn.BatchNorm2d(num_features=out_channels),
                nn.ReLU(True)
            )

        self.model = nn.Sequential(
            block(latent_dim, 512, padding=0, stride=1),
            block(512, 256),
            block(256, 128),
            block(128, 64),
            block(64, 32),
            nn.ConvTranspose2d(32, 3, kernel_size=4, stride=2, padding=1, bias=False),
            nn.Tanh()
        )

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


class DoppelDiscriminator(nn.Sequential):
    &quot;&quot;&quot;
    Discriminator network that classifies images in two categories
    &quot;&quot;&quot;

    def __init__(self):
        super().__init__()

        def block(in_channels: int, out_channels: int):
            return nn.Sequential(
                nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=4, stride=2, padding=1,
                          bias=False),
                nn.BatchNorm2d(num_features=out_channels),
                nn.LeakyReLU(0.2, inplace=True),
            )

        self.model = nn.Sequential(
            block(3, 64),
            block(64, 128),
            block(128, 256),
            block(256, 512),
            nn.Conv2d(512, 1, kernel_size=4, stride=1, padding=0, bias=False),
            nn.Flatten(),
            nn.Linear(25, 1),
            nn.Sigmoid()
        )

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


class DoppelGAN(pl.LightningModule):

    def __init__(self,
                 channels: int,
                 width: int,
                 height: int,
                 lr: float = 0.0002,
                 b1: float = 0.5,
                 b2: float = 0.999,
                 batch_size: int = 64,
                 **kwargs):

        super().__init__()

        # Save all keyword arguments as hyperparameters, accessible through self.hparams.X)
        self.save_hyperparameters()

        # Initialize networks
        # data_shape = (channels, width, height)
        self.generator = DoppelGenerator(latent_dim=self.hparams.latent_dim, )
        self.discriminator = DoppelDiscriminator()

        self.validation_z = torch.randn(8, self.hparams.latent_dim,1,1)

    def forward(self, input):
        return self.generator(input)

    def adversarial_loss(self, y_hat, y):
        return F.binary_cross_entropy(y_hat, y)

    def training_step(self, batch, batch_idx, optimizer_idx):
        images = batch

        # Sample noise (batch_size, latent_dim,1,1)
        z = torch.randn(images.size(0), self.hparams.latent_dim,1,1)

        # Train generator
        if optimizer_idx == 0:

            # Generate images (call generator -- see forward -- on latent vector)
            self.generated_images = self(z)

            # Log sampled images (visualize what the generator comes up with)
            sample_images = self.generated_images[:6]
            grid = make_grid(sample_images)
            self.logger.experiment.add_image('generated_images', grid, 0)

            # Ground truth result (ie: all fake)
            valid = torch.ones(images.size(0), 1)

            # Adversarial loss is binary cross-entropy
            generator_loss = self.adversarial_loss(self.discriminator(self(z)), valid)
            tqdm_dict = {'gen_loss': generator_loss}

            output = {
                'loss': generator_loss,
                'progress_bar': tqdm_dict,
                'log': tqdm_dict
            }
            return output

        # Train discriminator: classify real from generated samples
        if optimizer_idx == 1:

            # How well can it label as real?
            valid = torch.ones(images.size(0), 1)
            real_loss = self.adversarial_loss(self.discriminator(images), valid)

            # How well can it label as fake?
            fake = torch.zeros(images.size(0), 1)
            fake_loss = self.adversarial_loss(
                self.discriminator(self(z).detach()), fake)

            # Discriminator loss is the average of these
            discriminator_loss = (real_loss + fake_loss) / 2
            tqdm_dict = {'d_loss': discriminator_loss}
            output = {
                'loss': discriminator_loss,
                'progress_bar': tqdm_dict,
                'log': tqdm_dict
            }
            return output

    def configure_optimizers(self):
        lr = self.hparams.lr
        b1 = self.hparams.b1
        b2 = self.hparams.b2

        # Optimizers
        opt_g = torch.optim.Adam(self.generator.parameters(), lr=lr, betas=(b1, b2))
        opt_d = torch.optim.Adam(self.discriminator.parameters(), lr=lr, betas=(b1, b2))

        # Return optimizers/schedulers (currently no scheduler)
        return [opt_g, opt_d], []

    def on_epoch_end(self):

        # Log sampled images
        sample_images = self(self.validation_z)
        grid = make_grid(sample_images)
        self.logger.experiment.add_image('generated_images', grid, self.current_epoch)


if __name__ == '__main__':

    # Global parameter
    image_dim = 128
    latent_dim = 100
    batch_size = 64

    # Initialize dataset
    tfs = transforms.Compose([
        ToPILImage(),
        Resize(image_dim),
        ToTensor()
    ])
    doppel_dataset = DoppelDataset(face_dir='../data/faces', transform=tfs)

    # Initialize data module
    doppel_data_module = DoppelDataModule(batch_size=batch_size)

    # Build models
    generator = DoppelGenerator(latent_dim=latent_dim)
    discriminator = DoppelDiscriminator()

    # Test generator
    x = torch.rand(batch_size, latent_dim, 1, 1)
    y = generator(x)
    print(f'Generator: x {x.size()} --&gt; y {y.size()}')

    # Test discriminator
    x = torch.rand(batch_size, 3, 128, 128)
    y = discriminator(x)
    print(f'Discriminator: x {x.size()} --&gt; y {y.size()}')

    # Build GAN
    doppelgan = DoppelGAN(batch_size=batch_size, channels=3, width=image_dim, height=image_dim, latent_dim=latent_dim)

    # Fit GAN
    trainer = pl.Trainer(gpus=0, max_epochs=5, progress_bar_refresh_rate=1)
    trainer.fit(model=doppelgan, datamodule=doppel_data_module)

Full traceback:
Traceback (most recent call last):
  File &quot;/usr/local/lib/python3.9/site-packages/IPython/core/interactiveshell.py&quot;, line 3437, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File &quot;&lt;ipython-input-2-28805d67d74b&gt;&quot;, line 1, in &lt;module&gt;
    runfile('/Users/wouter/Documents/OneDrive/Hardnose/Projects/Coding/0002_DoppelGANger/doppelganger/gan.py', wdir='/Users/wouter/Documents/OneDrive/Hardnose/Projects/Coding/0002_DoppelGANger/doppelganger')
  File &quot;/Applications/PyCharm.app/Contents/plugins/python/helpers/pydev/_pydev_bundle/pydev_umd.py&quot;, line 197, in runfile
    pydev_imports.execfile(filename, global_vars, local_vars)  # execute the script
  File &quot;/Applications/PyCharm.app/Contents/plugins/python/helpers/pydev/_pydev_imps/_pydev_execfile.py&quot;, line 18, in execfile
    exec(compile(contents+&quot;\n&quot;, file, 'exec'), glob, loc)
  File &quot;/Users/wouter/Documents/OneDrive/Hardnose/Projects/Coding/0002_DoppelGANger/doppelganger/gan.py&quot;, line 298, in &lt;module&gt;
    trainer.fit(model=doppelgan, datamodule=doppel_data_module)
  File &quot;/usr/local/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py&quot;, line 510, in fit
    results = self.accelerator_backend.train()
  File &quot;/usr/local/lib/python3.9/site-packages/pytorch_lightning/accelerators/accelerator.py&quot;, line 57, in train
    return self.train_or_test()
  File &quot;/usr/local/lib/python3.9/site-packages/pytorch_lightning/accelerators/accelerator.py&quot;, line 74, in train_or_test
    results = self.trainer.train()
  File &quot;/usr/local/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py&quot;, line 561, in train
    self.train_loop.run_training_epoch()
  File &quot;/usr/local/lib/python3.9/site-packages/pytorch_lightning/trainer/training_loop.py&quot;, line 550, in run_training_epoch
    batch_output = self.run_training_batch(batch, batch_idx, dataloader_idx)
  File &quot;/usr/local/lib/python3.9/site-packages/pytorch_lightning/trainer/training_loop.py&quot;, line 718, in run_training_batch
    self.optimizer_step(optimizer, opt_idx, batch_idx, train_step_and_backward_closure)
  File &quot;/usr/local/lib/python3.9/site-packages/pytorch_lightning/trainer/training_loop.py&quot;, line 485, in optimizer_step
    model_ref.optimizer_step(
  File &quot;/usr/local/lib/python3.9/site-packages/pytorch_lightning/core/lightning.py&quot;, line 1298, in optimizer_step
    optimizer.step(closure=optimizer_closure)
  File &quot;/usr/local/lib/python3.9/site-packages/pytorch_lightning/core/optimizer.py&quot;, line 286, in step
    self.__optimizer_step(*args, closure=closure, profiler_name=profiler_name, **kwargs)
  File &quot;/usr/local/lib/python3.9/site-packages/pytorch_lightning/core/optimizer.py&quot;, line 144, in __optimizer_step
    optimizer.step(closure=closure, *args, **kwargs)
  File &quot;/usr/local/lib/python3.9/site-packages/torch/autograd/grad_mode.py&quot;, line 26, in decorate_context
    return func(*args, **kwargs)
  File &quot;/usr/local/lib/python3.9/site-packages/torch/optim/adam.py&quot;, line 66, in step
    loss = closure()
  File &quot;/usr/local/lib/python3.9/site-packages/pytorch_lightning/trainer/training_loop.py&quot;, line 708, in train_step_and_backward_closure
    result = self.training_step_and_backward(
  File &quot;/usr/local/lib/python3.9/site-packages/pytorch_lightning/trainer/training_loop.py&quot;, line 806, in training_step_and_backward
    result = self.training_step(split_batch, batch_idx, opt_idx, hiddens)
  File &quot;/usr/local/lib/python3.9/site-packages/pytorch_lightning/trainer/training_loop.py&quot;, line 319, in training_step
    training_step_output = self.trainer.accelerator_backend.training_step(args)
  File &quot;/usr/local/lib/python3.9/site-packages/pytorch_lightning/accelerators/cpu_accelerator.py&quot;, line 62, in training_step
    return self._step(self.trainer.model.training_step, args)
  File &quot;/usr/local/lib/python3.9/site-packages/pytorch_lightning/accelerators/cpu_accelerator.py&quot;, line 58, in _step
    output = model_step(*args)
  File &quot;/Users/wouter/Documents/OneDrive/Hardnose/Projects/Coding/0002_DoppelGANger/doppelganger/gan.py&quot;, line 223, in training_step
    real_loss = self.adversarial_loss(self.discriminator(images), valid)
  File &quot;/usr/local/lib/python3.9/site-packages/torch/nn/modules/module.py&quot;, line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File &quot;/Users/wouter/Documents/OneDrive/Hardnose/Projects/Coding/0002_DoppelGANger/doppelganger/gan.py&quot;, line 154, in forward
    return self.model(input)
  File &quot;/usr/local/lib/python3.9/site-packages/torch/nn/modules/module.py&quot;, line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File &quot;/usr/local/lib/python3.9/site-packages/torch/nn/modules/container.py&quot;, line 117, in forward
    input = module(input)
  File &quot;/usr/local/lib/python3.9/site-packages/torch/nn/modules/module.py&quot;, line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File &quot;/usr/local/lib/python3.9/site-packages/torch/nn/modules/linear.py&quot;, line 93, in forward
    return F.linear(input, self.weight, self.bias)
  File &quot;/usr/local/lib/python3.9/site-packages/torch/nn/functional.py&quot;, line 1690, in linear
    ret = torch.addmm(bias, input, weight.t())
RuntimeError: mat1 and mat2 shapes cannot be multiplied (7x9 and 25x1)


This multiplication problem comes from the DoppelDiscriminator. There is a linear layer
    nn.Linear(25, 1),

that should be
    nn.Linear(9, 1),

based on the error message.
