openpilot v0.9.6 release

date: 2024-01-12T10:13:37
master commit: ba792d576a49a0899b88a753fa1c52956bedf9e6

tinygrad_repo/tinygrad/nn/__init__.py

@@ -0,0 +1,128 @@
+import math
+from typing import Optional, Union, Tuple
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import prod, all_int
+
+class BatchNorm2d:
+  def __init__(self, sz, eps=1e-5, affine=True, track_running_stats=True, momentum=0.1):
+    self.eps, self.track_running_stats, self.momentum = eps, track_running_stats, momentum
+
+    if affine: self.weight, self.bias = Tensor.ones(sz), Tensor.zeros(sz)
+    else: self.weight, self.bias = None, None
+
+    self.running_mean, self.running_var = Tensor.zeros(sz, requires_grad=False), Tensor.ones(sz, requires_grad=False)
+    self.num_batches_tracked = Tensor.zeros(1, requires_grad=False)
+
+  def __call__(self, x:Tensor):
+    if Tensor.training:
+      # This requires two full memory accesses to x
+      # https://github.com/pytorch/pytorch/blob/c618dc13d2aa23625cb0d7ada694137532a4fa33/aten/src/ATen/native/cuda/Normalization.cuh
+      # There's "online" algorithms that fix this, like https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_Online_algorithm
+      batch_mean = x.mean(axis=(0,2,3))
+      y = (x - batch_mean.reshape(shape=[1, -1, 1, 1]))
+      batch_var = (y*y).mean(axis=(0,2,3))
+      batch_invstd = batch_var.add(self.eps).pow(-0.5)
+
+      # NOTE: wow, this is done all throughout training in most PyTorch models
+      if self.track_running_stats:
+        self.running_mean.assign((1 - self.momentum) * self.running_mean + self.momentum * batch_mean.detach())
+        self.running_var.assign((1 - self.momentum) * self.running_var + self.momentum * prod(y.shape)/(prod(y.shape) - y.shape[1]) * batch_var.detach() )
+        self.num_batches_tracked += 1
+    else:
+      batch_mean = self.running_mean
+      # NOTE: this can be precomputed for static inference. we expand it here so it fuses
+      batch_invstd = self.running_var.reshape(1, -1, 1, 1).expand(x.shape).add(self.eps).rsqrt()
+
+    return x.batchnorm(self.weight, self.bias, batch_mean, batch_invstd)
+
+# TODO: these Conv lines are terrible
+def Conv1d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True):
+  return Conv2d(in_channels, out_channels, (kernel_size,), stride, padding, dilation, groups, bias)
+
+class Conv2d:
+  def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True):
+    self.kernel_size = (kernel_size, kernel_size) if isinstance(kernel_size, int) else tuple(kernel_size)
+    self.stride, self.padding, self.dilation, self.groups = stride, padding, dilation, groups
+    self.weight = self.initialize_weight(out_channels, in_channels, groups)
+    assert all_int(self.weight.shape), "does not support symbolic shape"
+    bound = 1 / math.sqrt(prod(self.weight.shape[1:]))
+    self.bias = Tensor.uniform(out_channels, low=-bound, high=bound) if bias else None
+
+  def __call__(self, x:Tensor):
+    return x.conv2d(self.weight, self.bias, padding=self.padding, stride=self.stride, dilation=self.dilation, groups=self.groups)
+
+  def initialize_weight(self, out_channels, in_channels, groups): return Tensor.kaiming_uniform(out_channels, in_channels//groups, *self.kernel_size, a=math.sqrt(5))
+
+def ConvTranspose1d(in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, dilation=1, groups=1, bias=True):
+  return ConvTranspose2d(in_channels, out_channels, (kernel_size,), stride, padding, output_padding, dilation, groups, bias)
+
+class ConvTranspose2d(Conv2d):
+  def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, dilation=1, groups=1, bias=True):
+    super().__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias)
+    self.output_padding = output_padding
+
+  def __call__(self, x:Tensor):
+    return x.conv_transpose2d(self.weight, self.bias, padding=self.padding, output_padding=self.output_padding, stride=self.stride, dilation=self.dilation, groups=self.groups)
+
+  def initialize_weight(self, out_channels, in_channels, groups): return Tensor.kaiming_uniform(in_channels, out_channels//groups, *self.kernel_size, a=math.sqrt(5))
+
+class Linear:
+  def __init__(self, in_features, out_features, bias=True):
+    self.weight = Tensor.kaiming_uniform(out_features, in_features, a=math.sqrt(5))
+    # TODO: remove this once we can represent Tensor with int shape in typing
+    assert isinstance(self.weight.shape[1], int), "does not support symbolic shape"
+    bound = 1 / math.sqrt(self.weight.shape[1])
+    self.bias = Tensor.uniform(out_features, low=-bound, high=bound) if bias else None
+
+  def __call__(self, x:Tensor):
+    return x.linear(self.weight.transpose(), self.bias)
+
+class GroupNorm:
+  def __init__(self, num_groups:int, num_channels:int, eps:float=1e-5, affine:bool=True):
+    self.num_groups, self.num_channels, self.eps = num_groups, num_channels, eps
+    self.weight: Optional[Tensor] = Tensor.ones(num_channels) if affine else None
+    self.bias: Optional[Tensor] = Tensor.zeros(num_channels) if affine else None
+
+  def __call__(self, x:Tensor):
+    # reshape for layernorm to work as group norm
+    # subtract mean and divide stddev
+    x = x.reshape(x.shape[0], self.num_groups, -1).layernorm(eps=self.eps).reshape(x.shape)
+
+    if self.weight is None or self.bias is None: return x
+    # elementwise_affine on channels
+    return x * self.weight.reshape(1, -1, *[1] * (len(x.shape)-2)) + self.bias.reshape(1, -1, *[1] * (len(x.shape)-2))
+
+class InstanceNorm:
+  def __init__(self, num_features:int, eps:float=1e-5, affine:bool=True):
+    self.num_features, self.eps = num_features, eps
+    self.weight: Optional[Tensor] = Tensor.ones(num_features) if affine else None
+    self.bias: Optional[Tensor] = Tensor.zeros(num_features) if affine else None
+
+  def __call__(self, x:Tensor):
+    x = x.reshape(x.shape[0], self.num_features, -1).layernorm(eps=self.eps).reshape(x.shape)
+    if self.weight is None or self.bias is None: return x
+    return x * self.weight.reshape(1, -1, *[1] * (len(x.shape)-2)) + self.bias.reshape(1, -1, *[1] * (len(x.shape)-2))
+
+class LayerNorm:
+  def __init__(self, normalized_shape:Union[int, Tuple[int, ...]], eps:float=1e-5, elementwise_affine:bool=True):
+    self.normalized_shape = (normalized_shape,) if isinstance(normalized_shape, int) else tuple(normalized_shape)
+    self.axis, self.eps, self.elementwise_affine = tuple(-1-i for i in range(len(self.normalized_shape))), eps, elementwise_affine
+    self.weight, self.bias = (Tensor.ones(*self.normalized_shape), Tensor.zeros(*self.normalized_shape)) if elementwise_affine else (None, None)
+
+  def __call__(self, x:Tensor):
+    assert self.normalized_shape == x.shape[-len(self.normalized_shape):], f"last dimensions of {x.shape} must match {self.normalized_shape}"
+    x = x.layernorm(eps=self.eps, axis=self.axis)
+    if not self.elementwise_affine: return x
+    return x * self.weight + self.bias
+
+class LayerNorm2d(LayerNorm):
+  def __call__(self, x): return super().__call__(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+
+class Embedding:
+  def __init__(self, vocab_size:int, embed_size:int):
+    self.vocab_size = vocab_size
+    self.weight = Tensor.glorot_uniform(vocab_size, embed_size)
+
+  def __call__(self, idx:Tensor) -> Tensor:
+    if not hasattr(self, 'vocab_counter'): self.vocab_counter = Tensor.arange(self.vocab_size, requires_grad=False).reshape(1, 1, self.vocab_size)
+    return (self.vocab_counter == idx.unsqueeze(2)).expand(*idx.shape, self.vocab_size) @ self.weight

tinygrad_repo/tinygrad/nn/optim.py

@@ -0,0 +1,68 @@
+# sorted in order of increasing complexity
+from typing import List
+from tinygrad.helpers import dedup
+from tinygrad.tensor import Tensor
+
+class Optimizer:
+  def __init__(self, params: List[Tensor], lr: float):
+    # if it's None, but being put into an optimizer, set it to True
+    for x in params:
+      if x.requires_grad is None: x.requires_grad = True
+
+    self.params: List[Tensor] = dedup([x for x in params if x.requires_grad])
+    self.buffers: List[Tensor] = dedup([x for x in params if not x.requires_grad])   # buffers are still realized
+    self.lr = Tensor([lr], requires_grad=False).contiguous()
+
+  def zero_grad(self):
+    for param in self.params: param.grad = None
+
+  def realize(self, extra=None):
+    # NOTE: in extra is too late for most of the params due to issues with assign
+    Tensor.corealize(extra + self.params + self.buffers if extra is not None else self.params + self.buffers)
+
+class SGD(Optimizer):
+  def __init__(self, params: List[Tensor], lr=0.001, momentum=0, weight_decay=0.0, nesterov=False):
+    super().__init__(params, lr)
+    self.momentum, self.wd, self.nesterov = momentum, weight_decay, nesterov
+    self.b = [Tensor.zeros(*t.shape, device=t.device, requires_grad=False) for t in self.params] if self.momentum else []
+
+  # https://pytorch.org/docs/stable/generated/torch.optim.SGD.html
+  def step(self) -> None:
+    for i, t in enumerate(self.params):
+      assert t.grad is not None
+      g = t.grad.realize() + self.wd * t.detach()
+      if self.momentum:
+        self.b[i].assign(self.momentum * self.b[i] + g).realize()  # NOTE: self.b[i] is zero on the first run, no if required
+        g = (g + self.momentum * self.b[i]) if self.nesterov else self.b[i]
+      t.assign(t.detach() - g * self.lr)
+    self.realize(self.b)
+
+# LAMB is essentially just the trust ratio part of LARS applied to Adam/W so if we just set the trust ratio to 1.0 its just Adam/W.
+def AdamW(params: List[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-8, wd=0.01): return LAMB(params, lr, b1, b2, eps, wd, adam=True)
+def Adam(params: List[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-8): return LAMB(params, lr, b1, b2, eps, 0.0, adam=True)
+
+class LAMB(Optimizer):
+  def __init__(self, params: List[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, wd=0.0, adam=False):
+    super().__init__(params, lr)
+    self.b1, self.b2, self.eps, self.wd, self.adam, self.t = b1, b2, eps, wd, adam, Tensor([0], requires_grad=False).realize()
+    self.m = [Tensor.zeros(*t.shape, device=t.device, requires_grad=False) for t in self.params]
+    self.v = [Tensor.zeros(*t.shape, device=t.device, requires_grad=False) for t in self.params]
+
+  def step(self) -> None:
+    self.t.assign(self.t + 1).realize()
+    for i, t in enumerate(self.params):
+      assert t.grad is not None
+      g = t.grad.realize()
+      self.m[i].assign(self.b1 * self.m[i] + (1.0 - self.b1) * g).realize()
+      self.v[i].assign(self.b2 * self.v[i] + (1.0 - self.b2) * (g * g)).realize()
+      m_hat = self.m[i] / (1.0 - self.b1**self.t)
+      v_hat = self.v[i] / (1.0 - self.b2**self.t)
+      up = (m_hat / (v_hat.sqrt() + self.eps)) + self.wd * t.detach()
+      if not self.adam:
+        r1 = t.detach().square().sum().sqrt()
+        r2 = up.square().sum().sqrt()
+        r = Tensor.where(r1 > 0, Tensor.where(r2 > 0, r1 / r2, 1.0), 1.0)
+      else:
+        r = 1.0
+      t.assign(t.detach() - self.lr * r * up)
+    self.realize([self.t] + self.m + self.v)

tinygrad_repo/tinygrad/nn/state.py

@@ -0,0 +1,124 @@
+import os, json, pathlib, zipfile, pickle
+from tqdm import tqdm
+from typing import Dict, Union, List, Optional, Any, Tuple
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import dtypes, prod, argsort, DEBUG, Timing, GlobalCounters, CI
+from tinygrad.shape.view import strides_for_shape
+from tinygrad.ops import Device
+
+safe_dtypes = {"F16": dtypes.float16, "F32": dtypes.float32, "U8": dtypes.uint8, "I8": dtypes.int8, "I32": dtypes.int32, "I64": dtypes.int64}
+inverse_safe_dtypes = {v:k for k,v in safe_dtypes.items()}
+
+def safe_load_metadata(fn:Union[Tensor,str]) -> Tuple[Tensor, int, Any]:
+  t = fn if isinstance(fn, Tensor) else Tensor.empty(os.stat(fn).st_size, dtype=dtypes.uint8, device=f"disk:{fn}")
+  json_len = t[0:1].cast(dtypes.int64).numpy()[0]
+  return (t, json_len, json.loads(t[8:8+json_len].numpy().tobytes()))
+
+def safe_load(fn:Union[Tensor,str]) -> Dict[str, Tensor]:
+  t, json_len, metadata = safe_load_metadata(fn)
+  return {k:t[8+json_len+v['data_offsets'][0]:].cast(safe_dtypes[v['dtype']])[:prod(v['shape'])].reshape(v['shape']) for k,v in metadata.items() if k != "__metadata__"}
+
+def safe_save(tensors:Dict[str, Tensor], fn:str, metadata:Optional[Dict[str, Any]]=None):
+  headers, offset = {}, 0
+  if metadata: headers['__metadata__'] = metadata
+  for k,v in tensors.items():
+    headers[k] = {'dtype': inverse_safe_dtypes[v.dtype], 'shape': list(v.shape), 'data_offsets':[offset, offset+v.nbytes()]}
+    offset += v.nbytes()
+  j = json.dumps(headers, separators=(',', ':'))
+  j += "\x20"*((8-len(j)%8)%8)
+  pathlib.Path(fn).unlink(missing_ok=True)
+  t = Tensor.empty(8+len(j)+offset, dtype=dtypes.uint8, device=f"disk:{fn}")
+  t[0:1].cast(dtypes.int64).assign([len(j)])
+  t[8:8+len(j)].assign(Tensor(list(j.encode('utf-8')), dtype=dtypes.uint8, device="cpu"))
+  for k,v in safe_load(t).items(): v.assign(tensors[k])
+
+# state dict
+
+from collections import OrderedDict
+def get_state_dict(obj, prefix:str='', tensor_type=Tensor) -> Dict[str, Tensor]:
+  if isinstance(obj, tensor_type): return {prefix.strip('.'):obj}
+  if hasattr(obj, '_asdict'): return get_state_dict(obj._asdict(), prefix, tensor_type)  # namedtuple
+  if isinstance(obj, OrderedDict): return get_state_dict(dict(obj), prefix, tensor_type)
+  if hasattr(obj, '__dict__'): return get_state_dict(obj.__dict__, prefix, tensor_type)
+  state_dict = {}
+  if isinstance(obj, (list, tuple)):
+    for i,x in enumerate(obj): state_dict.update(get_state_dict(x, f"{prefix}{str(i)}.", tensor_type))
+  elif isinstance(obj, dict):
+    for k,v in obj.items(): state_dict.update(get_state_dict(v, f"{prefix}{str(k)}.", tensor_type))
+  return state_dict
+def get_parameters(obj) -> List[Tensor]: return list(get_state_dict(obj).values())
+
+def load_state_dict(model, state_dict, strict=True, verbose=True):
+  with Timing("loaded weights in ", lambda et_ns: f", {GlobalCounters.mem_used/1e9:.2f} GB loaded at {GlobalCounters.mem_used/et_ns:.2f} GB/s"):
+    model_state_dict = get_state_dict(model)
+    if DEBUG >= 1 and len(state_dict) > len(model_state_dict): print("WARNING: unused weights in state_dict", sorted(list(state_dict.keys() - model_state_dict.keys())))
+    for k,v in (t := tqdm(model_state_dict.items(), disable=CI or not verbose)):
+      t.set_description(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB, {k:50s}")
+      if k not in state_dict and not strict:
+        if DEBUG >= 1: print(f"WARNING: not loading {k}")
+        continue
+      v.assign(state_dict[k].to(v.device)).realize()
+
+# torch support!
+
+def torch_load(fn:str):
+  t = Tensor.empty(os.stat(fn).st_size, dtype=dtypes.uint8, device=f"disk:{fn}")
+
+  offsets: Dict[str, int] = {}
+  lens: Dict[str, int] = {}
+  def _rebuild_tensor_v2(storage, storage_offset, size, stride, requires_grad, backward_hooks, metadata=None):
+    #print(storage, storage_offset, size, stride, requires_grad, backward_hooks, metadata)
+    lens[storage[2]] = storage[4] * storage[1].itemsize
+    if storage[2] not in offsets: return None
+    byte_offset = offsets[storage[2]]+storage_offset*storage[1].itemsize
+    ret = t[byte_offset:byte_offset+prod(size)].cast(storage[1])
+    # convert bfloat16 -> float16 using LLVM for Llama 2
+    # upstream LLaMA also does this conversion:
+    # https://github.com/facebookresearch/llama/blob/6c7fe276574e78057f917549435a2554000a876d/llama/generation.py#L95
+    # TODO: should this be done in the example instead? or maybe we don't need this anymore with better bfloat16 support
+    if storage[1] == dtypes.bfloat16:
+      ret = ret.bitcast(dtypes.uint16).to("CPU").cast(dtypes.uint32).mul(1<<16).bitcast(dtypes.float32).to(Device.DEFAULT).half()
+      #ret = ret.to("LLVM").half().to(Device.DEFAULT)
+
+    # 7 lines to deal with permuted tensors. NOTE: this currently requires reading off the disk
+    shape_strides = [(s, st) for s,st in zip(size, stride) if s != 1]
+    permute_indexes = [len(shape_strides)-1-y for y in argsort([x[1] for x in shape_strides])]
+    if tuple(permute_indexes) != tuple(range(len(permute_indexes))):
+      intermediate_shape = tuple([shape_strides[x][0] for x in argsort(permute_indexes)])
+      assert tuple([shape_strides[i][1] for i in argsort(permute_indexes)]) == strides_for_shape(intermediate_shape), "nonpermutable strides"
+      if DEBUG >= 2: print(f"WARNING: this torch load is slow. CPU to permute {intermediate_shape} with {permute_indexes}")
+      # TODO: find a nice way to support all shapetracker on disktensors
+      ret = ret.cpu().reshape(intermediate_shape).permute(permute_indexes)
+
+    return ret.reshape(size)
+
+  intercept = {"HalfStorage": dtypes.float16, "FloatStorage": dtypes.float32, "BFloat16Storage": dtypes.bfloat16, "IntStorage": dtypes.int32, "LongStorage": dtypes.int64, "_rebuild_tensor_v2": _rebuild_tensor_v2}
+  whitelist = {"torch", "collections", "numpy", "_codecs"}  # NOTE: this is not for security, only speed
+  class Dummy: pass
+  class TorchPickle(pickle.Unpickler):
+    def find_class(self, module, name):
+      module_root = module.split(".")[0]
+      if module_root not in whitelist:
+        if DEBUG >= 2: print(f"WARNING: returning Dummy for {module} {name}")
+        return Dummy
+      return intercept[name] if module_root == "torch" else super().find_class(module, name)
+    def persistent_load(self, pid): return pid
+
+  if tuple(t[0:2].numpy()) == (0x50, 0x4b):
+    myzip = zipfile.ZipFile(fn, 'r')
+    base_name = myzip.namelist()[0].split('/', 1)[0]
+    for n in myzip.namelist():
+      if n.startswith(f'{base_name}/data/'):
+        with myzip.open(n) as myfile:
+          offsets[n.split("/")[-1]] = myfile._orig_compress_start # type: ignore
+    with myzip.open(f'{base_name}/data.pkl') as myfile:
+      return TorchPickle(myfile).load()
+  else:
+    with open(fn, "rb") as f:
+      pkl = TorchPickle(f)
+      _, _, _, rwd, _, ids, base_offset = pkl.load(), pkl.load(), pkl.load(), f.tell(), pkl.load(), pkl.load(), f.tell()
+      for i in ids:
+        offsets[i] = base_offset + 8
+        base_offset += 8 + lens[i]
+      f.seek(rwd)
+      return TorchPickle(f).load()