Variational Autoencoder

Generative modeling with VAEs on MNIST using PyTorch.

neural-networks

optimization

variational-inference

Implements a variational autoencoder (VAE) in PyTorch for the MNIST dataset, covering the encoder-decoder architecture, the reparameterization trick, and the ELBO loss combining reconstruction and KL divergence terms.

Published

June 18, 2025

import torch
import torch.utils.data
from torch import nn, optim
from torch.nn import functional as F
from torchvision import datasets, transforms
from torchvision.utils import save_image
import matplotlib.pyplot as plt
import os
os.makedirs("results", exist_ok=True)

This example is taken and adapted from the torch example repository

BATCH_SIZE = 128

train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('./data', train=True, download=True,
                   transform=transforms.ToTensor()),
    batch_size=BATCH_SIZE, shuffle=True)
test_loader = torch.utils.data.DataLoader(
    datasets.MNIST('./data', train=False, transform=transforms.ToTensor()),
    batch_size=BATCH_SIZE, shuffle=True)

0.3%

0.7%

1.0%

1.3%

1.7%

2.0%

2.3%

2.6%

3.0%

3.3%

3.6%

4.0%

4.3%

4.6%

5.0%

5.3%

5.6%

6.0%

6.3%

6.6%

6.9%

7.3%

7.6%

7.9%

8.3%

8.6%

8.9%

9.3%

9.6%

9.9%

10.2%

10.6%

10.9%

11.2%

11.6%

11.9%

12.2%

12.6%

12.9%

13.2%

13.6%

13.9%

14.2%

14.5%

14.9%

15.2%

15.5%

15.9%

16.2%

16.5%

16.9%

17.2%

17.5%

17.9%

18.2%

18.5%

18.8%

19.2%

19.5%

19.8%

20.2%

20.5%

20.8%

21.2%

21.5%

21.8%

22.1%

22.5%

22.8%

23.1%

23.5%

23.8%

24.1%

24.5%

24.8%

25.1%

25.5%

25.8%

26.1%

26.4%

26.8%

27.1%

27.4%

27.8%

28.1%

28.4%

28.8%

29.1%

29.4%

29.8%

30.1%

30.4%

30.7%

31.1%

31.4%

31.7%

32.1%

32.4%

32.7%

33.1%

33.4%

33.7%

34.0%

34.4%

34.7%

35.0%

35.4%

35.7%

36.0%

36.4%

36.7%

37.0%

37.4%

37.7%

38.0%

38.3%

38.7%

39.0%

39.3%

39.7%

40.0%

40.3%

40.7%

41.0%

41.3%

41.7%

42.0%

42.3%

42.6%

43.0%

43.3%

43.6%

44.0%

44.3%

44.6%

45.0%

45.3%

45.6%

45.9%

46.3%

46.6%

46.9%

47.3%

47.6%

47.9%

48.3%

48.6%

48.9%

49.3%

49.6%

49.9%

50.2%

50.6%

50.9%

51.2%

51.6%

51.9%

52.2%

52.6%

52.9%

53.2%

53.6%

53.9%

54.2%

54.5%

54.9%

55.2%

55.5%

55.9%

56.2%

56.5%

56.9%

57.2%

57.5%

57.9%

58.2%

58.5%

58.8%

59.2%

59.5%

59.8%

60.2%

60.5%

60.8%

61.2%

61.5%

61.8%

62.1%

62.5%

62.8%

63.1%

63.5%

63.8%

64.1%

64.5%

64.8%

65.1%

65.5%

65.8%

66.1%

66.4%

66.8%

67.1%

67.4%

67.8%

68.1%

68.4%

68.8%

69.1%

69.4%

69.8%

70.1%

70.4%

70.7%

71.1%

71.4%

71.7%

72.1%

72.4%

72.7%

73.1%

73.4%

73.7%

74.0%

74.4%

74.7%

75.0%

75.4%

75.7%

76.0%

76.4%

76.7%

77.0%

77.4%

77.7%

78.0%

78.3%

78.7%

79.0%

79.3%

79.7%

80.0%

80.3%

80.7%

81.0%

81.3%

81.7%

82.0%

82.3%

82.6%

83.0%

83.3%

83.6%

84.0%

84.3%

84.6%

85.0%

85.3%

85.6%

85.9%

86.3%

86.6%

86.9%

87.3%

87.6%

87.9%

88.3%

88.6%

88.9%

89.3%

89.6%

89.9%

90.2%

90.6%

90.9%

91.2%

91.6%

91.9%

92.2%

92.6%

92.9%

93.2%

93.6%

93.9%

94.2%

94.5%

94.9%

95.2%

95.5%

95.9%

96.2%

96.5%

96.9%

97.2%

97.5%

97.9%

98.2%

98.5%

98.8%

99.2%

99.5%

99.8%

100.0%

100.0%

2.0%

4.0%

6.0%

7.9%

9.9%

11.9%

13.9%

15.9%

17.9%

19.9%

21.9%

23.8%

25.8%

27.8%

29.8%

31.8%

33.8%

35.8%

37.8%

39.7%

41.7%

43.7%

45.7%

47.7%

49.7%

51.7%

53.7%

55.6%

57.6%

59.6%

61.6%

63.6%

65.6%

67.6%

69.6%

71.5%

73.5%

75.5%

77.5%

79.5%

81.5%

83.5%

85.5%

87.4%

89.4%

91.4%

93.4%

95.4%

97.4%

99.4%

100.0%

100.0%

class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()

        self.fc1 = nn.Linear(784, 400)
        self.fc21 = nn.Linear(400, 20)
        self.fc22 = nn.Linear(400, 20)
        self.fc3 = nn.Linear(20, 400)
        self.fc4 = nn.Linear(400, 784)

        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()

    def encode(self, x):
        h1 = self.relu(self.fc1(x))
        return self.fc21(h1), self.fc22(h1)

    def reparameterize(self, mu, logvar):
        if self.training:
            std = logvar.mul(0.5).exp_()
            eps = torch.randn_like(std)
            return eps.mul(std).add_(mu)
        else:
            return mu

    def decode(self, z):
        h3 = self.relu(self.fc3(z))
        return self.sigmoid(self.fc4(h3))

    def forward(self, x):
        mu, logvar = self.encode(x.view(-1, 784))
        z = self.reparameterize(mu, logvar)
        return self.decode(z), mu, logvar

model = VAE()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

def loss_function(recon_x, x, mu, logvar):
    BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduction='sum')

    # see Appendix B from VAE paper:
    # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
    # https://arxiv.org/abs/1312.6114
    # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

    return BCE + KLD

def train(epoch):
    model.train()
    train_loss = 0
    for batch_idx, (data, _) in enumerate(train_loader):
        optimizer.zero_grad()
        recon_batch, mu, logvar = model(data)
        loss = loss_function(recon_batch, data, mu, logvar)
        loss.backward()
        train_loss += loss.item()
        optimizer.step()
        if batch_idx % 20 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader),
                loss.item() / len(data)))

    print('====> Epoch: {} Average loss: {:.4f}'.format(
          epoch, train_loss / len(train_loader.dataset)))

def test(epoch):
    model.eval()
    test_loss = 0
    with torch.no_grad():
        for i, (data, _) in enumerate(test_loader):
            recon_batch, mu, logvar = model(data)
            test_loss += loss_function(recon_batch, data, mu, logvar).item()
            if i == 0:
                n = min(data.size(0), 8)
                comparison = torch.cat([data[:n],
                                      recon_batch.view(BATCH_SIZE, 1, 28, 28)[:n]])
                save_image(comparison.data,
                         'results/reconstruction_' + str(epoch) + '.png', nrow=n)

    test_loss /= len(test_loader.dataset)
    print('====> Test set loss: {:.4f}'.format(test_loss))

for epoch in range(1, 10 + 1):
    train(epoch)
    test(epoch)
    with torch.no_grad():
        sample = torch.randn(64, 20)
        sample = model.decode(sample)
        save_image(sample.view(64, 1, 28, 28),
                   'results/sample_' + str(epoch) + '.png')

Train Epoch: 1 [0/60000 (0%)]   Loss: 548.634521
Train Epoch: 1 [2560/60000 (4%)]    Loss: 236.767303

Train Epoch: 1 [5120/60000 (9%)]    Loss: 209.762161
Train Epoch: 1 [7680/60000 (13%)]   Loss: 202.013107

Train Epoch: 1 [10240/60000 (17%)]  Loss: 181.194870
Train Epoch: 1 [12800/60000 (21%)]  Loss: 173.787308

Train Epoch: 1 [15360/60000 (26%)]  Loss: 165.583420
Train Epoch: 1 [17920/60000 (30%)]  Loss: 162.148636

Train Epoch: 1 [20480/60000 (34%)]  Loss: 156.180466
Train Epoch: 1 [23040/60000 (38%)]  Loss: 158.943665

Train Epoch: 1 [25600/60000 (43%)]  Loss: 147.105698
Train Epoch: 1 [28160/60000 (47%)]  Loss: 152.154602

Train Epoch: 1 [30720/60000 (51%)]  Loss: 145.153564
Train Epoch: 1 [33280/60000 (55%)]  Loss: 147.480225

Train Epoch: 1 [35840/60000 (60%)]  Loss: 142.398788
Train Epoch: 1 [38400/60000 (64%)]  Loss: 146.010529

Train Epoch: 1 [40960/60000 (68%)]  Loss: 133.627197
Train Epoch: 1 [43520/60000 (72%)]  Loss: 136.690384

Train Epoch: 1 [46080/60000 (77%)]  Loss: 134.174911
Train Epoch: 1 [48640/60000 (81%)]  Loss: 135.501343

Train Epoch: 1 [51200/60000 (85%)]  Loss: 136.076080
Train Epoch: 1 [53760/60000 (90%)]  Loss: 134.992554

Train Epoch: 1 [56320/60000 (94%)]  Loss: 138.561661
Train Epoch: 1 [58880/60000 (98%)]  Loss: 125.460587

====> Epoch: 1 Average loss: 164.1723

====> Test set loss: 120.3354
Train Epoch: 2 [0/60000 (0%)]   Loss: 131.731873
Train Epoch: 2 [2560/60000 (4%)]    Loss: 127.556808

Train Epoch: 2 [5120/60000 (9%)]    Loss: 129.977646
Train Epoch: 2 [7680/60000 (13%)]   Loss: 130.733414

Train Epoch: 2 [10240/60000 (17%)]  Loss: 128.357269
Train Epoch: 2 [12800/60000 (21%)]  Loss: 122.836334

Train Epoch: 2 [15360/60000 (26%)]  Loss: 124.673889
Train Epoch: 2 [17920/60000 (30%)]  Loss: 121.683350

Train Epoch: 2 [20480/60000 (34%)]  Loss: 124.010063
Train Epoch: 2 [23040/60000 (38%)]  Loss: 117.151344

Train Epoch: 2 [25600/60000 (43%)]  Loss: 118.172012
Train Epoch: 2 [28160/60000 (47%)]  Loss: 114.916489

Train Epoch: 2 [30720/60000 (51%)]  Loss: 119.968933
Train Epoch: 2 [33280/60000 (55%)]  Loss: 122.616447

Train Epoch: 2 [35840/60000 (60%)]  Loss: 119.052917
Train Epoch: 2 [38400/60000 (64%)]  Loss: 124.251404

Train Epoch: 2 [40960/60000 (68%)]  Loss: 120.944016
Train Epoch: 2 [43520/60000 (72%)]  Loss: 119.940079

Train Epoch: 2 [46080/60000 (77%)]  Loss: 123.439484
Train Epoch: 2 [48640/60000 (81%)]  Loss: 121.641663

Train Epoch: 2 [51200/60000 (85%)]  Loss: 117.665115
Train Epoch: 2 [53760/60000 (90%)]  Loss: 121.112869

Train Epoch: 2 [56320/60000 (94%)]  Loss: 114.654282
Train Epoch: 2 [58880/60000 (98%)]  Loss: 118.567261

====> Epoch: 2 Average loss: 122.0352

====> Test set loss: 107.2045
Train Epoch: 3 [0/60000 (0%)]   Loss: 117.265152
Train Epoch: 3 [2560/60000 (4%)]    Loss: 120.659363

Train Epoch: 3 [5120/60000 (9%)]    Loss: 119.945419
Train Epoch: 3 [7680/60000 (13%)]   Loss: 113.292702

Train Epoch: 3 [10240/60000 (17%)]  Loss: 116.421890
Train Epoch: 3 [12800/60000 (21%)]  Loss: 117.810295

Train Epoch: 3 [15360/60000 (26%)]  Loss: 118.534966
Train Epoch: 3 [17920/60000 (30%)]  Loss: 115.239372

Train Epoch: 3 [20480/60000 (34%)]  Loss: 116.788185
Train Epoch: 3 [23040/60000 (38%)]  Loss: 119.369278

Train Epoch: 3 [25600/60000 (43%)]  Loss: 111.231339
Train Epoch: 3 [28160/60000 (47%)]  Loss: 112.620918

Train Epoch: 3 [30720/60000 (51%)]  Loss: 119.008453
Train Epoch: 3 [33280/60000 (55%)]  Loss: 111.599335

Train Epoch: 3 [35840/60000 (60%)]  Loss: 114.980438
Train Epoch: 3 [38400/60000 (64%)]  Loss: 118.566757

Train Epoch: 3 [40960/60000 (68%)]  Loss: 115.240967
Train Epoch: 3 [43520/60000 (72%)]  Loss: 113.397987

Train Epoch: 3 [46080/60000 (77%)]  Loss: 111.692711
Train Epoch: 3 [48640/60000 (81%)]  Loss: 111.034027

Train Epoch: 3 [51200/60000 (85%)]  Loss: 114.306198
Train Epoch: 3 [53760/60000 (90%)]  Loss: 114.897110

Train Epoch: 3 [56320/60000 (94%)]  Loss: 110.192017
Train Epoch: 3 [58880/60000 (98%)]  Loss: 114.898056
====> Epoch: 3 Average loss: 114.7251

====> Test set loss: 103.5963
Train Epoch: 4 [0/60000 (0%)]   Loss: 110.943779
Train Epoch: 4 [2560/60000 (4%)]    Loss: 113.126038

Train Epoch: 4 [5120/60000 (9%)]    Loss: 112.282608
Train Epoch: 4 [7680/60000 (13%)]   Loss: 109.765732

Train Epoch: 4 [10240/60000 (17%)]  Loss: 110.486633
Train Epoch: 4 [12800/60000 (21%)]  Loss: 111.794464

Train Epoch: 4 [15360/60000 (26%)]  Loss: 112.159172
Train Epoch: 4 [17920/60000 (30%)]  Loss: 112.933701

Train Epoch: 4 [20480/60000 (34%)]  Loss: 116.019157
Train Epoch: 4 [23040/60000 (38%)]  Loss: 112.606613

Train Epoch: 4 [25600/60000 (43%)]  Loss: 109.638031
Train Epoch: 4 [28160/60000 (47%)]  Loss: 113.518456

Train Epoch: 4 [30720/60000 (51%)]  Loss: 115.787491
Train Epoch: 4 [33280/60000 (55%)]  Loss: 109.828369

Train Epoch: 4 [35840/60000 (60%)]  Loss: 111.565414
Train Epoch: 4 [38400/60000 (64%)]  Loss: 112.422310

Train Epoch: 4 [40960/60000 (68%)]  Loss: 109.158768
Train Epoch: 4 [43520/60000 (72%)]  Loss: 111.735451

Train Epoch: 4 [46080/60000 (77%)]  Loss: 110.691467
Train Epoch: 4 [48640/60000 (81%)]  Loss: 112.220108

Train Epoch: 4 [51200/60000 (85%)]  Loss: 115.208496
Train Epoch: 4 [53760/60000 (90%)]  Loss: 112.054604

Train Epoch: 4 [56320/60000 (94%)]  Loss: 113.814384
Train Epoch: 4 [58880/60000 (98%)]  Loss: 108.074463
====> Epoch: 4 Average loss: 111.7308

====> Test set loss: 101.2154
Train Epoch: 5 [0/60000 (0%)]   Loss: 108.546211
Train Epoch: 5 [2560/60000 (4%)]    Loss: 112.200302

Train Epoch: 5 [5120/60000 (9%)]    Loss: 110.150421
Train Epoch: 5 [7680/60000 (13%)]   Loss: 108.479080

Train Epoch: 5 [10240/60000 (17%)]  Loss: 109.602615
Train Epoch: 5 [12800/60000 (21%)]  Loss: 109.784431

Train Epoch: 5 [15360/60000 (26%)]  Loss: 109.162491
Train Epoch: 5 [17920/60000 (30%)]  Loss: 109.937584

Train Epoch: 5 [20480/60000 (34%)]  Loss: 110.969543
Train Epoch: 5 [23040/60000 (38%)]  Loss: 111.078796

Train Epoch: 5 [25600/60000 (43%)]  Loss: 106.948372
Train Epoch: 5 [28160/60000 (47%)]  Loss: 108.658745

Train Epoch: 5 [30720/60000 (51%)]  Loss: 106.874847
Train Epoch: 5 [33280/60000 (55%)]  Loss: 109.023651

Train Epoch: 5 [35840/60000 (60%)]  Loss: 111.881744
Train Epoch: 5 [38400/60000 (64%)]  Loss: 108.394989

Train Epoch: 5 [40960/60000 (68%)]  Loss: 109.016220
Train Epoch: 5 [43520/60000 (72%)]  Loss: 112.585129

Train Epoch: 5 [46080/60000 (77%)]  Loss: 107.197525
Train Epoch: 5 [48640/60000 (81%)]  Loss: 107.831329

Train Epoch: 5 [51200/60000 (85%)]  Loss: 109.893654
Train Epoch: 5 [53760/60000 (90%)]  Loss: 106.641510

Train Epoch: 5 [56320/60000 (94%)]  Loss: 113.776169
Train Epoch: 5 [58880/60000 (98%)]  Loss: 108.198402
====> Epoch: 5 Average loss: 109.8871

====> Test set loss: 99.8494
Train Epoch: 6 [0/60000 (0%)]   Loss: 110.150063
Train Epoch: 6 [2560/60000 (4%)]    Loss: 107.798172

Train Epoch: 6 [5120/60000 (9%)]    Loss: 112.263824
Train Epoch: 6 [7680/60000 (13%)]   Loss: 108.610275

Train Epoch: 6 [10240/60000 (17%)]  Loss: 103.126335
Train Epoch: 6 [12800/60000 (21%)]  Loss: 107.636826

Train Epoch: 6 [15360/60000 (26%)]  Loss: 112.287262
Train Epoch: 6 [17920/60000 (30%)]  Loss: 109.707733

Train Epoch: 6 [20480/60000 (34%)]  Loss: 107.782089
Train Epoch: 6 [23040/60000 (38%)]  Loss: 107.041412

Train Epoch: 6 [25600/60000 (43%)]  Loss: 104.542320
Train Epoch: 6 [28160/60000 (47%)]  Loss: 108.867142

Train Epoch: 6 [30720/60000 (51%)]  Loss: 111.374702
Train Epoch: 6 [33280/60000 (55%)]  Loss: 103.376457

Train Epoch: 6 [35840/60000 (60%)]  Loss: 107.413147
Train Epoch: 6 [38400/60000 (64%)]  Loss: 108.295319

Train Epoch: 6 [40960/60000 (68%)]  Loss: 109.501320
Train Epoch: 6 [43520/60000 (72%)]  Loss: 107.235397

Train Epoch: 6 [46080/60000 (77%)]  Loss: 104.432060
Train Epoch: 6 [48640/60000 (81%)]  Loss: 114.063919

Train Epoch: 6 [51200/60000 (85%)]  Loss: 109.082703
Train Epoch: 6 [53760/60000 (90%)]  Loss: 108.684021

Train Epoch: 6 [56320/60000 (94%)]  Loss: 111.484451
Train Epoch: 6 [58880/60000 (98%)]  Loss: 109.607918
====> Epoch: 6 Average loss: 108.6958

====> Test set loss: 99.3851
Train Epoch: 7 [0/60000 (0%)]   Loss: 107.670486
Train Epoch: 7 [2560/60000 (4%)]    Loss: 111.157127

Train Epoch: 7 [5120/60000 (9%)]    Loss: 110.604736
Train Epoch: 7 [7680/60000 (13%)]   Loss: 106.069023

Train Epoch: 7 [10240/60000 (17%)]  Loss: 108.498962
Train Epoch: 7 [12800/60000 (21%)]  Loss: 106.300484

Train Epoch: 7 [15360/60000 (26%)]  Loss: 108.293655
Train Epoch: 7 [17920/60000 (30%)]  Loss: 110.514900

Train Epoch: 7 [20480/60000 (34%)]  Loss: 109.119965
Train Epoch: 7 [23040/60000 (38%)]  Loss: 103.302643

Train Epoch: 7 [25600/60000 (43%)]  Loss: 105.771667
Train Epoch: 7 [28160/60000 (47%)]  Loss: 108.977036

Train Epoch: 7 [30720/60000 (51%)]  Loss: 108.591492
Train Epoch: 7 [33280/60000 (55%)]  Loss: 100.977043

Train Epoch: 7 [35840/60000 (60%)]  Loss: 109.002068
Train Epoch: 7 [38400/60000 (64%)]  Loss: 108.723053

Train Epoch: 7 [40960/60000 (68%)]  Loss: 108.024910
Train Epoch: 7 [43520/60000 (72%)]  Loss: 101.559898

Train Epoch: 7 [46080/60000 (77%)]  Loss: 105.333687
Train Epoch: 7 [48640/60000 (81%)]  Loss: 105.032669

Train Epoch: 7 [51200/60000 (85%)]  Loss: 106.020271
Train Epoch: 7 [53760/60000 (90%)]  Loss: 106.608353

Train Epoch: 7 [56320/60000 (94%)]  Loss: 105.024284
Train Epoch: 7 [58880/60000 (98%)]  Loss: 109.724609
====> Epoch: 7 Average loss: 107.7882

====> Test set loss: 99.0749
Train Epoch: 8 [0/60000 (0%)]   Loss: 111.200348
Train Epoch: 8 [2560/60000 (4%)]    Loss: 110.275085

Train Epoch: 8 [5120/60000 (9%)]    Loss: 104.035355
Train Epoch: 8 [7680/60000 (13%)]   Loss: 105.540955

Train Epoch: 8 [10240/60000 (17%)]  Loss: 106.949150
Train Epoch: 8 [12800/60000 (21%)]  Loss: 108.091881

Train Epoch: 8 [15360/60000 (26%)]  Loss: 109.387390
Train Epoch: 8 [17920/60000 (30%)]  Loss: 108.231750

Train Epoch: 8 [20480/60000 (34%)]  Loss: 108.981506
Train Epoch: 8 [23040/60000 (38%)]  Loss: 108.095581

Train Epoch: 8 [25600/60000 (43%)]  Loss: 109.591125
Train Epoch: 8 [28160/60000 (47%)]  Loss: 107.178627

Train Epoch: 8 [30720/60000 (51%)]  Loss: 109.066597
Train Epoch: 8 [33280/60000 (55%)]  Loss: 110.745300

Train Epoch: 8 [35840/60000 (60%)]  Loss: 108.644302
Train Epoch: 8 [38400/60000 (64%)]  Loss: 104.321968

Train Epoch: 8 [40960/60000 (68%)]  Loss: 106.065765
Train Epoch: 8 [43520/60000 (72%)]  Loss: 105.985229

Train Epoch: 8 [46080/60000 (77%)]  Loss: 105.464775
Train Epoch: 8 [48640/60000 (81%)]  Loss: 106.755745

Train Epoch: 8 [51200/60000 (85%)]  Loss: 110.724197
Train Epoch: 8 [53760/60000 (90%)]  Loss: 107.127319

Train Epoch: 8 [56320/60000 (94%)]  Loss: 103.233719
Train Epoch: 8 [58880/60000 (98%)]  Loss: 106.832184
====> Epoch: 8 Average loss: 107.1373

====> Test set loss: 97.9776
Train Epoch: 9 [0/60000 (0%)]   Loss: 107.381592
Train Epoch: 9 [2560/60000 (4%)]    Loss: 108.062187

Train Epoch: 9 [5120/60000 (9%)]    Loss: 106.987198
Train Epoch: 9 [7680/60000 (13%)]   Loss: 105.263321

Train Epoch: 9 [10240/60000 (17%)]  Loss: 104.146004
Train Epoch: 9 [12800/60000 (21%)]  Loss: 105.281403

Train Epoch: 9 [15360/60000 (26%)]  Loss: 108.410995
Train Epoch: 9 [17920/60000 (30%)]  Loss: 107.729935

Train Epoch: 9 [20480/60000 (34%)]  Loss: 106.034317
Train Epoch: 9 [23040/60000 (38%)]  Loss: 106.851746

Train Epoch: 9 [25600/60000 (43%)]  Loss: 107.023514
Train Epoch: 9 [28160/60000 (47%)]  Loss: 108.563576

Train Epoch: 9 [30720/60000 (51%)]  Loss: 111.120827
Train Epoch: 9 [33280/60000 (55%)]  Loss: 109.766617

Train Epoch: 9 [35840/60000 (60%)]  Loss: 105.246010
Train Epoch: 9 [38400/60000 (64%)]  Loss: 106.029976

Train Epoch: 9 [40960/60000 (68%)]  Loss: 107.460793
Train Epoch: 9 [43520/60000 (72%)]  Loss: 105.061920

Train Epoch: 9 [46080/60000 (77%)]  Loss: 112.483910
Train Epoch: 9 [48640/60000 (81%)]  Loss: 107.765755

Train Epoch: 9 [51200/60000 (85%)]  Loss: 105.779938
Train Epoch: 9 [53760/60000 (90%)]  Loss: 106.320320

Train Epoch: 9 [56320/60000 (94%)]  Loss: 109.062149
Train Epoch: 9 [58880/60000 (98%)]  Loss: 107.750999
====> Epoch: 9 Average loss: 106.5836

====> Test set loss: 97.2389
Train Epoch: 10 [0/60000 (0%)]  Loss: 105.873734
Train Epoch: 10 [2560/60000 (4%)]   Loss: 104.698433

Train Epoch: 10 [5120/60000 (9%)]   Loss: 103.561890
Train Epoch: 10 [7680/60000 (13%)]  Loss: 109.678497

Train Epoch: 10 [10240/60000 (17%)] Loss: 107.231934
Train Epoch: 10 [12800/60000 (21%)] Loss: 106.615662

Train Epoch: 10 [15360/60000 (26%)] Loss: 102.571754
Train Epoch: 10 [17920/60000 (30%)] Loss: 109.108490

Train Epoch: 10 [20480/60000 (34%)] Loss: 107.755447
Train Epoch: 10 [23040/60000 (38%)] Loss: 103.485107

Train Epoch: 10 [25600/60000 (43%)] Loss: 107.544464
Train Epoch: 10 [28160/60000 (47%)] Loss: 106.768890

Train Epoch: 10 [30720/60000 (51%)] Loss: 106.400993
Train Epoch: 10 [33280/60000 (55%)] Loss: 108.585594

Train Epoch: 10 [35840/60000 (60%)] Loss: 101.729080
Train Epoch: 10 [38400/60000 (64%)] Loss: 104.781586

Train Epoch: 10 [40960/60000 (68%)] Loss: 104.166504
Train Epoch: 10 [43520/60000 (72%)] Loss: 109.962120

Train Epoch: 10 [46080/60000 (77%)] Loss: 102.225922
Train Epoch: 10 [48640/60000 (81%)] Loss: 105.121124

Train Epoch: 10 [51200/60000 (85%)] Loss: 105.502281
Train Epoch: 10 [53760/60000 (90%)] Loss: 100.612816

Train Epoch: 10 [56320/60000 (94%)] Loss: 104.523346
Train Epoch: 10 [58880/60000 (98%)] Loss: 108.477737
====> Epoch: 10 Average loss: 106.1819

====> Test set loss: 96.4372

from PIL import Image as PILImage
import numpy as np

fig, axes = plt.subplots(1, 2, figsize=(12, 4))
img3 = PILImage.open("results/reconstruction_3.png")
axes[0].imshow(np.array(img3), cmap="gray")
axes[0].set_title("Reconstruction (Epoch 3)")
axes[0].axis("off")
img9 = PILImage.open("results/reconstruction_9.png")
axes[1].imshow(np.array(img9), cmap="gray")
axes[1].set_title("Reconstruction (Epoch 9)")
axes[1].axis("off")
plt.tight_layout()

img_sample = PILImage.open("results/sample_10.png")
plt.figure(figsize=(8, 8))
plt.imshow(np.array(img_sample), cmap="gray")
plt.title("Generated samples (Epoch 10)")
plt.axis("off");