IQ.Pilot Release Commit @ f2a861c

artifacts/package_sources/tinygrad/examples/mnist_gan.py

@@ -0,0 +1,106 @@
+from pathlib import Path
+import torch
+from torchvision.utils import make_grid, save_image
+from tinygrad.nn.state import get_parameters
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import trange, Context
+from tinygrad.nn import optim
+from tinygrad.nn.datasets import mnist
+
+class LinearGen:
+  def __init__(self):
+    self.l1 = Tensor.scaled_uniform(128, 256)
+    self.l2 = Tensor.scaled_uniform(256, 512)
+    self.l3 = Tensor.scaled_uniform(512, 1024)
+    self.l4 = Tensor.scaled_uniform(1024, 784)
+
+  def forward(self, x):
+    x = x.dot(self.l1).leaky_relu(0.2)
+    x = x.dot(self.l2).leaky_relu(0.2)
+    x = x.dot(self.l3).leaky_relu(0.2)
+    x = x.dot(self.l4).tanh()
+    return x
+
+class LinearDisc:
+  def __init__(self):
+    self.l1 = Tensor.scaled_uniform(784, 1024)
+    self.l2 = Tensor.scaled_uniform(1024, 512)
+    self.l3 = Tensor.scaled_uniform(512, 256)
+    self.l4 = Tensor.scaled_uniform(256, 2)
+
+  def forward(self, x):
+    # balance the discriminator inputs with const bias (.add(1))
+    x = x.dot(self.l1).add(1).leaky_relu(0.2).dropout(0.3)
+    x = x.dot(self.l2).leaky_relu(0.2).dropout(0.3)
+    x = x.dot(self.l3).leaky_relu(0.2).dropout(0.3)
+    x = x.dot(self.l4).log_softmax()
+    return x
+
+def make_batch(images):
+  sample = Tensor.randint(batch_size, low=0, high=images.shape[0])
+  return images[sample].reshape(batch_size, 28*28).cast('float').div(127.5).sub(1.0)
+
+def make_labels(bs, col, val=-2.0):
+  y = Tensor.zeros(bs, 2)
+  if col == 0: y = y + Tensor([val, 0.0])
+  else: y = y + Tensor([0.0, val])
+  return y
+
+def train_discriminator(optimizer, data_real, data_fake):
+  real_labels = make_labels(batch_size, 1)
+  fake_labels = make_labels(batch_size, 0)
+  optimizer.zero_grad()
+  output_real = discriminator.forward(data_real)
+  output_fake = discriminator.forward(data_fake)
+  loss_real = (output_real * real_labels).mean()
+  loss_fake = (output_fake * fake_labels).mean()
+  loss_real.backward()
+  loss_fake.backward()
+  optimizer.step()
+  return (loss_real + loss_fake).numpy()
+
+def train_generator(optimizer, data_fake):
+  real_labels = make_labels(batch_size, 1)
+  optimizer.zero_grad()
+  output = discriminator.forward(data_fake)
+  loss = (output * real_labels).mean()
+  loss.backward()
+  optimizer.step()
+  return loss.numpy()
+
+if __name__ == "__main__":
+  # data for training and validation
+  X_train, _, _, _ = mnist()
+  ds_noise = Tensor.randn(64, 128)
+  # parameters
+  epochs, batch_size, k = 300, 512, 1
+  sample_interval = epochs // 10
+  n_steps = X_train.shape[0] // batch_size
+  # models and optimizer
+  generator = LinearGen()
+  discriminator = LinearDisc()
+  # path to store results
+  output_dir = Path(".").resolve() / "outputs"
+  output_dir.mkdir(exist_ok=True)
+  # optimizers
+  optim_g = optim.Adam(get_parameters(generator), lr=0.0002, b1=0.5)  # 0.0002 for equilibrium!
+  optim_d = optim.Adam(get_parameters(discriminator), lr=0.0002, b1=0.5)
+  # training loop
+  with Context(TRAINING=1):
+    for epoch in (t := trange(epochs)):
+      loss_g, loss_d = 0.0, 0.0
+      for _ in range(n_steps):
+        data_real = make_batch(X_train)
+        for step in range(k):  # Try with k = 5 or 7.
+          noise = Tensor.randn(batch_size, 128)
+          data_fake = generator.forward(noise).detach()
+          loss_d += train_discriminator(optim_d, data_real, data_fake)
+        noise = Tensor.randn(batch_size, 128)
+        data_fake = generator.forward(noise)
+        loss_g += train_generator(optim_g, data_fake)
+      if (epoch + 1) % sample_interval == 0:
+        fake_images = generator.forward(ds_noise).detach().numpy()
+        fake_images = (fake_images.reshape(-1, 1, 28, 28) + 1) / 2  # 0 - 1 range.
+        save_image(make_grid(torch.tensor(fake_images)), output_dir / f"image_{epoch+1}.jpg")
+      t.set_description(f"Generator loss: {loss_g/n_steps}, Discriminator loss: {loss_d/n_steps}")
+  print("Training Completed!")