import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
from torchvision.datasets import MNIST

from vis_util import visual_mnist


##### settings 
x_dim = 28 * 28 # size of mnist digit
z_dim = 100     # random noise
h_dim = 128     # hidden layer
batch_size = 60000 // 1000
lr = 1e-3
epochs = 120


##### load data and generate targets
# NOTE: we only need train data
trainset = MNIST('mnist', train=True, transform=ToTensor(), download=True)
dataloader = DataLoader(trainset, batch_size=batch_size, shuffle=True)
# targets for D, who only output 1 for real data, 0 for fake data
real_target = torch.ones(batch_size, 1) 
fake_target = torch.zeros(batch_size, 1)


##### network arch 
class Generator(nn.Module):
    """generator, generates fake data which is similar to the real data, to 
    fool the discriminator"""

    def __init__(self, z_dim, h_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(z_dim, h_dim),
            nn.ReLU(inplace=True),
            nn.Linear(h_dim, x_dim),
            nn.Sigmoid(),
        )

    def forward(self, x):
        x = self.net(x)
        return x


class Discriminator(nn.Module):
    """discriminator, is response for judging whether an input is real data 
    or fake data that generated by the generator"""

    def __init__(self, x_dim, h_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(x_dim, h_dim),
            nn.ReLU(inplace=True),
            nn.Linear(h_dim, 1),
            nn.Sigmoid(),
        )

    def forward(self, x):
        x = self.net(x)
        return x


##### init model
G = Generator(z_dim, h_dim)
D = Discriminator(x_dim, h_dim)


##### optimizer and loss
G_optim = optim.Adam(G.parameters(), lr=lr)
D_optim = optim.Adam(D.parameters(), lr=lr)

D_loss_real = nn.BCELoss()
D_loss_fake = nn.BCELoss()
G_loss_f2r = nn.BCELoss()


##### training loop
G.train()
D.train()
for epoch_i in range(epochs):
    for x, _ in dataloader:
        # step 1: G generate fake data fx using noise z
        z = torch.randn(batch_size, z_dim)
        fx = G(z)

        # step 2: D judge on (fx, x), then update itself
        fake = D(fx)
        x = x.view(-1, x_dim)
        real = D(x)
        D_loss = D_loss_real(real, real_target) + D_loss_fake(fake, fake_target)
        D_optim.zero_grad()
        D_loss.backward()
        D_optim.step()

        # step 3: G update
        z = torch.randn(batch_size, z_dim)
        fx = G(z)
        fake = D(fx)
        G_loss = G_loss_f2r(fake, real_target)
        G_optim.zero_grad()
        G_loss.backward()
        G_optim.step()

    # for each epoch, visualize result
    print(f"Epoch-{epoch_i}: D_loss: {D_loss:.5f}, G_loss: {G_loss:.5f}")
    visual_mnist(epoch_i, fx.detach().numpy()[:16], (4, 4))
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
from torchvision.datasets import MNIST

from vis_util import visual_mnist


##### settings 
x_dim = 28 * 28 # size of mnist digit
z_dim = 100     # random noise
h_dim = 128     # hidden layer
batch_size = 60000 // 1000
lr = 1e-3
epochs = 120


##### load data and generate targets
# NOTE: we only need train data
trainset = MNIST('mnist', train=True, transform=ToTensor(), download=True)
dataloader = DataLoader(trainset, batch_size=batch_size, shuffle=True)
# targets for D, who only output 1 for real data, 0 for fake data
real_target = torch.ones(batch_size, 1) 
fake_target = torch.zeros(batch_size, 1)


##### network arch 
class Generator(nn.Module):
    """generator, generates fake data which is similar to the real data, to 
    fool the discriminator"""

    def __init__(self, z_dim, h_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(z_dim, h_dim),
            nn.ReLU(inplace=True),
            nn.Linear(h_dim, x_dim),
            nn.Sigmoid(),
        )

    def forward(self, x):
        x = self.net(x)
        return x


class Discriminator(nn.Module):
    """discriminator, is response for judging whether an input is real data 
    or fake data that generated by the generator"""

    def __init__(self, x_dim, h_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(x_dim, h_dim),
            nn.ReLU(inplace=True),
            nn.Linear(h_dim, 1),
            nn.Sigmoid(),
        )

    def forward(self, x):
        x = self.net(x)
        return x


##### init model
G = Generator(z_dim, h_dim)
D = Discriminator(x_dim, h_dim)


##### optimizer and loss
G_optim = optim.Adam(G.parameters(), lr=lr)
D_optim = optim.Adam(D.parameters(), lr=lr)

D_loss_real = nn.BCELoss()
D_loss_fake = nn.BCELoss()
G_loss_f2r = nn.BCELoss()


##### training loop
G.train()
D.train()
for epoch_i in range(epochs):
    for x, _ in dataloader:
        # step 1: G generate fake data fx using noise z
        z = torch.randn(batch_size, z_dim)
        fx = G(z)

        # step 2: D judge on (fx, x), then update itself
        fake = D(fx)
        x = x.view(-1, x_dim)
        real = D(x)
        D_loss = D_loss_real(real, real_target) + D_loss_fake(fake, fake_target)
        D_optim.zero_grad()
        D_loss.backward()
        D_optim.step()

        # step 3: G update
        z = torch.randn(batch_size, z_dim)
        fx = G(z)
        fake = D(fx)
        G_loss = G_loss_f2r(fake, real_target)
        G_optim.zero_grad()
        G_loss.backward()
        G_optim.step()

    # for each epoch, visualize result
    print(f"Epoch-{epoch_i}: D_loss: {D_loss:.5f}, G_loss: {G_loss:.5f}")
    visual_mnist(epoch_i, fx.detach().numpy()[:16], (4, 4))
##### settings # size of mnist digit # random noise # hidden layer ##### load data and generate targets # NOTE: we only need train data # targets for D, who only output 1 for real data, 0 for fake data ##### network arch generator, generates fake data which is similar to the real data, to      fool the discriminator discriminator, is response for judging whether an input is real data      or fake data that generated by the generator ##### init model ##### optimizer and loss ##### training loop # step 1: G generate fake data fx using noise z # step 2: D judge on (fx, x), then update itself # step 3: G update # for each epoch, visualize result