From bf116deb5a26a9b8dd9f4ab7976c577cfece9dfc Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Thu, 13 Nov 2025 13:44:55 -0800
Subject: [PATCH] imagenet labels

---
 tinygrad/apps/resnet.py    | 18 ++++++++++--------
 tinygrad/mixin/movement.py |  2 +-
 tinygrad/nn/datasets.py    |  6 ++++++
 tinygrad/nn/state.py       | 17 +++++++++++------
 tinygrad/tensor.py         |  6 ++++++
 5 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/tinygrad/apps/resnet.py b/tinygrad/apps/resnet.py
index 302d11ee64..3563201aa2 100644
--- a/tinygrad/apps/resnet.py
+++ b/tinygrad/apps/resnet.py
@@ -12,15 +12,15 @@ class Bottleneck:
     self.downsample = (stride != 1 or in_c != out_c) and [nn.Conv2d(in_c, out_c, 1, stride, bias=False), nn.BatchNorm2d(out_c)] or []
 
   def __call__(self, x:Tensor) -> Tensor:
-    id = x.sequential(self.downsample)
+    identity = x.sequential(self.downsample)
     x = self.bn1(self.conv1(x)).relu()
     x = self.bn2(self.conv2(x)).relu()
     x = self.bn3(self.conv3(x))
-    return (x + id).relu()
+    return (x + identity).relu()
 
 class ResNet50:
   def __init__(self, num_classes=1000):
-    self.conv1, self.bn1 = nn.Conv2d(3, 64, 7, 2, 3, bias=False), nn.BatchNorm2d(64)
+    self.conv1, self.bn1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False), nn.BatchNorm2d(64)
     self.layer1 = self._make_layer(64,  64,  3, 1)
     self.layer2 = self._make_layer(256, 128, 4, 2)
     self.layer3 = self._make_layer(512, 256, 6, 2)
@@ -34,14 +34,16 @@ class ResNet50:
 
   def __call__(self, x:Tensor) -> Tensor:
     x = self.bn1(self.conv1(x)).relu()
-    x = x.max_pool2d()
-    x = x.sequential([self.layer1, self.layer2, self.layer3, self.layer4])
+    # TODO: max_pool2d return type is Tensor | tuple[Tensor, Tensor], this should be type specialised
+    x = x.max_pool2d()  # type: ignore
+    x = x.sequential([*self.layer1, *self.layer2, *self.layer3, *self.layer4])
     x = x.mean((2, 3))
     return self.fc(x)
 
 if __name__ == "__main__":
-  state_dict = nn.state.safe_load(Tensor.from_url("https://huggingface.co/timm/resnet50.a1_in1k/resolve/main/model.safetensors"))
   model = ResNet50()
+  state_dict = nn.state.safe_load(Tensor.from_url("https://huggingface.co/timm/resnet50.a1_in1k/resolve/main/model.safetensors"))
   nn.state.load_state_dict(model, state_dict)
-  img = nn.state.png_load(Tensor.from_url(sys.argv[1]))
-  print(model(img).argmax())
\ No newline at end of file
+  img = nn.state.png_load(Tensor.from_url(sys.argv[1] if len(sys.argv) > 1 else "https://upload.wikimedia.org/wikipedia/commons/0/05/Cat.png"))
+  value = model(img.rearrange("h w c -> 1 c h w")).argmax().item()
+  print(nn.datasets.imagenet_labels()[value])
diff --git a/tinygrad/mixin/movement.py b/tinygrad/mixin/movement.py
index e7967885f3..5b5a931263 100644
--- a/tinygrad/mixin/movement.py
+++ b/tinygrad/mixin/movement.py
@@ -89,7 +89,7 @@ class MovementMixin:
     # resolve -1
     if (c := new_shape.count(-1)) > 1: raise RuntimeError(f"only one dimension can be inferred using -1, getting {new_shape}")
     if c: new_shape = tuple([-prod(self.shape) // prod(new_shape) if s == -1 else s for s in new_shape])
-    if prod(self.shape) != prod(new_shape): raise ValueError(f"size mismatch, can't reshape ({self.shape}) -> ({new_shape})")
+    if prod(self.shape) != prod(new_shape): raise ValueError(f"size mismatch, can't reshape {self.shape=} -> {new_shape=}")
     ret = self._mop(Ops.RESHAPE, arg=new_shape)
     return self if ret.shape == self.shape else ret
 
diff --git a/tinygrad/nn/datasets.py b/tinygrad/nn/datasets.py
index 30439fb49b..4134dbc6ef 100644
--- a/tinygrad/nn/datasets.py
+++ b/tinygrad/nn/datasets.py
@@ -1,3 +1,4 @@
+import ast
 from tinygrad.tensor import Tensor
 from tinygrad.nn.state import tar_extract
 
@@ -12,3 +13,8 @@ def cifar(device=None):
   train = Tensor.cat(*[tt[f"cifar-10-batches-bin/data_batch_{i}.bin"].reshape(-1, 3073).to(device) for i in range(1,6)])
   test = tt["cifar-10-batches-bin/test_batch.bin"].reshape(-1, 3073).to(device)
   return train[:, 1:].reshape(-1,3,32,32), train[:, 0], test[:, 1:].reshape(-1,3,32,32), test[:, 0]
+
+def imagenet_labels():
+  return ast.literal_eval(Tensor.from_url(
+    "https://gist.githubusercontent.com/yrevar/942d3a0ac09ec9e5eb3a/raw/238f720ff059c1f82f368259d1ca4ffa5dd8f9f5/imagenet1000_clsidx_to_labels.txt"
+  ).tobytes().decode())
diff --git a/tinygrad/nn/state.py b/tinygrad/nn/state.py
index 0fb1cd7993..2fb3bd9cc7 100644
--- a/tinygrad/nn/state.py
+++ b/tinygrad/nn/state.py
@@ -1,4 +1,4 @@
-import json, pathlib, zipfile, pickle, tarfile, struct, functools, io
+import json, pathlib, zipfile, pickle, tarfile, struct, functools, io, zlib
 from collections import OrderedDict
 from typing import Any, Callable, BinaryIO, Iterable, cast
 from tinygrad.tensor import Tensor
@@ -360,12 +360,17 @@ def gguf_load(tensor: Tensor) -> tuple[dict, dict[str, Tensor]]:
 def png_load(t:Tensor) -> Tensor:
   f = io.BufferedReader(TensorIO(t))
   assert f.read(8) == b'\x89PNG\r\n\x1a\n', "not a PNG"
+  idats = []
   while (slen:=f.read(4)):
-    len, typ = struct.unpack(">I", slen)[0], f.read(4)
-    dat = f.read(len)
+    ilen, typ = struct.unpack(">I", slen)[0], f.read(4)
+    dat = f.read(ilen)
     if typ == b'IHDR':
       width, height, depth, color_type, compression, filter_method, interlace = struct.unpack(">IIBBBBB", dat)
-      print(width, height, depth, color_type)
-
-    print(len, typ)
+      assert depth == 8 and color_type == 2 and compression == 0 and filter_method == 0 and interlace == 0, "only RGB PNG is supported"
+    if typ == b'IDAT':
+      idats.append(dat)
+    if DEBUG >= 3: print(ilen, typ)
     f.seek(4, 1)
+  decompressed = Tensor(zlib.decompress(b''.join(idats)))
+  # the first pixel in each scanline is a filter pixel
+  return decompressed.reshape(height, width*3+1)[:, 1:].reshape(height, width, 3)
diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py
index a6b86ae080..e5bf3984a8 100644
--- a/tinygrad/tensor.py
+++ b/tinygrad/tensor.py
@@ -326,6 +326,12 @@ class Tensor(OpMixin):
     assert all_int(self.shape), f"no data if shape is symbolic, {self.shape=}"
     return self._buffer().as_typed_buffer(self.shape)
 
+  def tobytes(self) -> bytes:
+    """
+    Returns the data of this tensor as bytes, like numpy's `.tobytes()`.
+    """
+    return bytes(self.data())
+
   def item(self) -> ConstType:
     """
     Returns the value of this tensor as a standard Python number.