Fizz Buzz in Pytorch

fizz_buzz_pytorch.py

@@ -0,0 +1,62 @@
+import numpy as np
+
+import torch
+from torch.autograd import Variable
+
+NUM_DIGITS = 10
+NUM_HIDDEN = 100
+BATCH_SIZE = 128
+
+# Represent each input by an array of its binary digits.
+def binary_encode(i, num_digits):
+    return np.array([i >> d & 1 for d in range(num_digits)])
+
+# One-hot encode the desired outputs: [number, "fizz", "buzz", "fizzbuzz"]
+def fizz_buzz_encode(i):
+    if   i % 15 == 0: return 3
+    elif i % 5  == 0: return 2
+    elif i % 3  == 0: return 1
+    else:             return 0
+
+def fizz_buzz_decode(i, prediction):
+    return [str(i), "fizz", "buzz", "fizzbuzz"][prediction]
+
+trX = Variable(torch.Tensor([binary_encode(i, NUM_DIGITS) for i in range(101, 2 ** NUM_DIGITS)]))
+trY = Variable(torch.LongTensor([fizz_buzz_encode(i) for i in range(101, 2 ** NUM_DIGITS)]))
+
+
+# Define the model
+model = torch.nn.Sequential(
+    torch.nn.Linear(NUM_DIGITS, NUM_HIDDEN),
+    torch.nn.ReLU(),
+    torch.nn.Linear(NUM_HIDDEN, 4)
+)
+loss_fn = torch.nn.CrossEntropyLoss()
+optimizer = torch.optim.SGD(model.parameters(), lr = 0.05)
+
+
+# Start training it
+for epoch in xrange(10000):
+    for start in range(0, len(trX), BATCH_SIZE):
+        end = start + BATCH_SIZE
+        batchX = trX[start:end]
+        batchY = trY[start:end]
+
+        y_pred = model(batchX)
+        loss = loss_fn(y_pred, batchY)
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+    # Find loss on training data
+    loss = loss_fn(model(trX), trY).data[0]
+    print 'Epoch:', epoch, 'Loss:', loss
+
+
+# Output now
+testX = Variable(torch.Tensor([binary_encode(i, NUM_DIGITS) for i in range(1, 101)]))
+testY = model(testX)
+predictions = zip(range(1, 101), list(testY.max(1)[1].data.tolist()))
+
+print [fizz_buzz_decode(i, x) for (i, x) in predictions]