1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
|
# Define Generator
class Generator(nn.Module):
def __init__(self, input_dim, output_dim):
super(Generator, self).__init__()
self.model = nn.Sequential(
nn.Linear(input_dim, 256),
nn.LeakyReLU(0.2),
nn.BatchNorm1d(256),
nn.Linear(256, 512),
nn.LeakyReLU(0.2),
nn.BatchNorm1d(512),
nn.Linear(512, 512),
nn.LeakyReLU(0.2),
nn.BatchNorm1d(512),
nn.Linear(512, 256),
nn.LeakyReLU(0.2),
nn.BatchNorm1d(256),
nn.Linear(256, output_dim),
)
def forward(self, z):
return self.model(z)
# Define Discriminator
class Discriminator(nn.Module):
def __init__(self, input_dim):
super(Discriminator, self).__init__()
self.model = nn.Sequential(
nn.Linear(input_dim, 512), # 增加神经元数量
nn.LeakyReLU(0.2),
nn.Linear(512, 256),
nn.LeakyReLU(0.2),
nn.Linear(256, 128),
nn.LeakyReLU(0.2),
nn.Linear(128, 64),
nn.LeakyReLU(0.2),
nn.Linear(64, 1),
nn.Sigmoid()
)
def forward(self, x):
return self.model(x)
# Generate training data
def generate_checkerboard(*, num: int, rng: np.random.Generator) -> np.ndarray:
x1 = rng.uniform(size=num) * 4 - 2
x2_ = rng.uniform(size=num) - rng.choice([0, 1], size=(num,)) * 2
x2 = x2_ + (np.floor(x1) % 2)
x = np.hstack([x1[:, None], x2[:, None]]) * 2
return x
# Training Data
X_train = generate_checkerboard(num=2000, rng=np.random.default_rng())
X_train = torch.tensor(X_train, dtype=torch.float32)
# Hyperparameters
z_dim = 20
lr_g = 0.0002
lr_d = 0.0001
batch_size = 512
epochs = 40000
# Initialize models
generator = Generator(input_dim=z_dim, output_dim=2)
discriminator = Discriminator(input_dim=2)
g_optimizer = optim.Adam(generator.parameters(), lr=lr_g, betas=(0.7, 0.999))
d_optimizer = optim.Adam(discriminator.parameters(), lr=lr_d, betas=(0.7, 0.999))
loss_fn = nn.BCELoss()
# Training loop
for epoch in range(epochs):
# Sample real data
idx = torch.randint(0, X_train.shape[0], (batch_size,))
real_data = X_train[idx]
# Sample noise
z = torch.randn(batch_size, z_dim) * 8
# Generate fake data
fake_data = generator(z)
# Train Discriminator
d_optimizer.zero_grad()
real_labels = torch.ones(batch_size, 1)
fake_labels = torch.zeros(batch_size, 1)
real_loss = loss_fn(discriminator(real_data), real_labels)
fake_loss = loss_fn(discriminator(fake_data.detach()), fake_labels)
d_loss = real_loss + fake_loss
d_loss.backward()
d_optimizer.step()
# Train Generator
g_optimizer.zero_grad()
g_loss = loss_fn(discriminator(fake_data), real_labels) # Flip labels
g_loss.backward()
g_optimizer.step()
if epoch % 500 == 0:
print(f"Epoch {epoch}, D Loss: {d_loss.item()}, G Loss: {g_loss.item()}")
# Generate new data
z = torch.randn(1000, z_dim) * 2
generated_data = generator(z).detach().numpy()
# Plot results
plt.scatter(X_train[:, 0], X_train[:, 1], label='Real Data', alpha=0.3)
plt.scatter(generated_data[:, 0], generated_data[:, 1], label='Generated Data', alpha=0.3)
plt.legend()
plt.show()
|