✨ README update

joey00072 · Sep 25, 2023 · d771294 · d771294
1 parent 62c7ff3
commit d771294
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 43 deletions.
diff --git a/README.md b/README.md
@@ -6,15 +6,15 @@ this is really autograd engine backed by numpy<br><br>
 #### `tinytorch.py` shall always remain under <b>1000 lines</b>. if not we will <i>revert commit</i>
 
 
-[![Python package](https://github.com/joey00072/tinytorch/actions/workflows/unit_test.yaml/badge.svg)](https://github.com/joey00072/tinytorch/actions/workflows/unit_test.yaml)
+[![Python package](https://github.com/joey00072/tinytorch/actions/workflows/unit_test.yaml/badge.svg)](https://github.com/joey00072/tinytorch/actions/workflows/unit_test.yaml) 
 
 
 $$
 f(x) =x^3+x
 $$
 
-```python
-import tinytorch as tt
+```python 
+import tinytorch as tt      #👀
 
 def f(x):
   return x**3 + x
@@ -39,7 +39,7 @@ print(x.grad)
 ### What can you do with it?
 #### Automatic diffecrtion, yep
 ```python 
-import tinytorch as tt #👀
+import tinytorch as tt 
 
 def f(x,y):
   return x**2 + x*y + (y**3+y) **0.5
@@ -58,12 +58,19 @@ print(y.grad)
 python mnist.py
 ```
 
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/11rmsyRW65tfKLtAF_D938xsvtgGHRu-z?usp=sharing)
+
 #### GPT?? you bet (yes LLM fr fr)
 
 ```bash
 GPU=1 python mnist.py
 ```
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/14MjO1sX3oOZZDZR3J2GwzvwYYk85B9Sz?usp=sharing)
+
+
 note: numpy is too slow to train llm you need to install jax (just using it as faster numpy)
+
+
 #### Visulization 
 If you want to see your computation graph run visulize.py
 
@@ -98,5 +105,8 @@ Bcs I was bored
 - doc sources if any
 - keep tinytorch.py under 1000 lines
 
+#### Buy me Chai/Coffee
+[![ko-fi](https://ko-fi.com/img/githubbutton_sm.svg)](https://ko-fi.com/R6R8KQTZ5)
+
 ### License
 [MIT](./LICENSE)
diff --git a/mnist.py b/mnist.py
@@ -10,39 +10,43 @@
 import requests
 from tqdm import tqdm
 
-# import torch
-# import torch.nn as nn
-# import torch.nn.functional as F
-# import torch.optim as optim
-
-import tinytorch as torch
-import tinytorch as nn
-import tinytorch as optim
-import tinytorch as F
+import tinytorch as tt
 
 # Constants
 EPOCHS = 1
 BATCH_SIZE = 32
 LR = 4e-3
 MNIST_DIR = "mnist"
-BASE_URL = "http://yann.lecun.com/exdb/mnist/"
-FILES = [
-    "train-images-idx3-ubyte.gz",
-    "train-labels-idx1-ubyte.gz",
-    "t10k-images-idx3-ubyte.gz",
-    "t10k-labels-idx1-ubyte.gz",
-]
 
 
 def download_mnist():
-    if not os.path.exists(MNIST_DIR):
-        os.makedirs(MNIST_DIR)
-        for file in FILES:
-            url = f"{BASE_URL}{file}"
-            response = requests.get(url)
-            with open(f"{MNIST_DIR}/{file}", "wb") as f:
-                f.write(response.content)
-
+    # Define the URLs of the files to download
+    base_url = "https://github.com/golbin/TensorFlow-MNIST/raw/master/mnist/data/"
+    files = [
+        "train-images-idx3-ubyte.gz",
+        "train-labels-idx1-ubyte.gz",
+        "t10k-images-idx3-ubyte.gz",
+        "t10k-labels-idx1-ubyte.gz",
+    ]
+    # Define the directory to save the files
+    save_dir = MNIST_DIR
+    os.makedirs(save_dir, exist_ok=True)
+
+    for file in files:
+        file_path = os.path.join(save_dir, file)
+
+        if os.path.exists(file_path):
+            continue
+
+        url = base_url + file
+        response = requests.get(url, stream=True)
+        if response.status_code == 200:
+            with open(file_path, 'wb') as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+            print(f"Downloaded {file}")
+        else:
+            print(f"Failed to download {file}. HTTP Response Code: {response.status_code}")
 
 def load_mnist() -> tuple:
     def read_labels(filename: str) -> np.array:
@@ -68,27 +72,27 @@ def one_hot(labels: np.array) -> np.array:
     return np.eye(10)[labels]
 
 
-def get_batch(images: torch.Tensor, labels: torch.Tensor):
+def get_batch(images: tt.Tensor, labels: tt.Tensor):
     indices = list(range(0, len(images), BATCH_SIZE))
     random.shuffle(indices)
     for i in indices:
         yield images[i : i + BATCH_SIZE], labels[i : i + BATCH_SIZE]
 
 
-class Network(nn.Module):
+class Network(tt.Module):
     def __init__(self) -> None:
         super().__init__()
-        self.l1 = nn.Linear(28 * 28, 128)
-        self.l2 = nn.Linear(128, 10)
+        self.l1 = tt.Linear(28 * 28, 128)
+        self.l2 = tt.Linear(128, 10)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = F.tanh(self.l1(x))
+    def forward(self, x: tt.Tensor) -> tt.Tensor:
+        x = tt.tanh(self.l1(x))
         return self.l2(x)
 
-@torch.no_grad()
-def test(model: Network, test_images: torch.Tensor, test_labels: torch.Tensor):
+@tt.no_grad()
+def test(model: Network, test_images: tt.Tensor, test_labels: tt.Tensor):
     preds = model.forward(test_images)
-    pred_indices = torch.argmax(preds, axis=-1).numpy()
+    pred_indices = tt.argmax(preds, axis=-1).numpy()
     test_labels = test_labels.numpy()
 
     correct = 0    
@@ -100,7 +104,7 @@ def test(model: Network, test_images: torch.Tensor, test_labels: torch.Tensor):
 
 
 def train(
-    model: Network, optimizer: optim.Adam, train_images: torch.Tensor, train_labels: torch.Tensor
+    model: Network, optimizer: tt.Adam, train_images: tt.Tensor, train_labels: tt.Tensor
 ):
     model.train()
     for epoch in range(EPOCHS):
@@ -111,7 +115,7 @@ def train(
             for batch_images, batch_labels in batch_generator:
                 optimizer.zero_grad()
                 pred = model.forward(batch_images)
-                loss = F.cross_entropy(pred, batch_labels)
+                loss = tt.cross_entropy(pred, batch_labels)
                 loss.backward()
                 optimizer.step()
 
@@ -128,14 +132,14 @@ def train(
     (train_images, train_labels), (test_images, test_labels) = load_mnist()
 
     train_labels, test_labels = map(
-        torch.tensor,  [train_labels, test_labels]
+        tt.tensor,  [train_labels, test_labels]
     )
 
-    train_images = torch.tensor(train_images.reshape(-1, 28 * 28) / 255).float()
-    test_images = torch.tensor(test_images.reshape(-1, 28 * 28) / 255).float()
+    train_images = tt.tensor(train_images.reshape(-1, 28 * 28) / 255).float()
+    test_images = tt.tensor(test_images.reshape(-1, 28 * 28) / 255).float()
 
     model = Network()
-    optimizer = optim.Adam(model.parameters(), lr=LR)
+    optimizer = tt.Adam(model.parameters(), lr=LR)
 
     start_time = time.perf_counter()
     train(model, optimizer, train_images, train_labels)

diff --git a/requirements.txt b/requirements.txt
@@ -1,5 +1,7 @@
 --find-links https://download.pytorch.org/whl/torch_stable.html
 numpy
 torch==2.0.1+cpu
+jax
+jaxlib
 pytest
 graphviz