1import torch
2import torch.nn.functional as F
3
4# G: noise -> sample D: sample -> P(real)
5G = make_mlp(in_dim=1, hidden=12, out_dim=1)
6D = make_mlp(in_dim=1, hidden=12, out_dim=1) # + sigmoid on output
7opt_G = torch.optim.SGD(G.parameters(), lr=0.05, momentum=0.5)
8opt_D = torch.optim.SGD(D.parameters(), lr=0.08, momentum=0.5)
9
10
11def train_round(real_batch):
12 B = real_batch.shape[0]
13 ones, zeros = torch.ones(B, 1), torch.zeros(B, 1)
14
15 # ---- 1. Detective step: real -> 1, fake -> 0 ----
16 z = torch.rand(B, 1) * 2 - 1
17 fakes = G(z).detach() # detach: don't teach G here
18 d_loss = (
19 F.binary_cross_entropy(torch.sigmoid(D(real_batch)), ones)
20 + F.binary_cross_entropy(torch.sigmoid(D(fakes)), zeros)
21 )
22 opt_D.zero_grad(); d_loss.backward(); opt_D.step()
23
24 # ---- 2. Forger step: make D say "real" ----
25 z = torch.rand(B, 1) * 2 - 1
26 verdict = torch.sigmoid(D(G(z))) # NO detach: gradients flow
27 # through the detective into G!
28 # Non-saturating loss: maximize log D(G(z))
29 g_loss = F.binary_cross_entropy(verdict, ones)
30 opt_G.zero_grad(); g_loss.backward(); opt_G.step()
31
32 return d_loss.item(), g_loss.item()
33
34# At equilibrium both losses hover near log 2 = 0.693 --
35# the detective is coin-flipping. Watch it happen in the lab.