mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-01-10 07:28:15 -05:00
ad756f6112
* use isinstance, some optimizations & whitespace removal
* revert whitespace changes
* revert more whitespace
* some more cleanup
* revert fstring (not a fan of the {{}})
* fix typo
* fix typo
76 lines
2.8 KiB
Python
76 lines
2.8 KiB
Python
from tinygrad.tensor import Tensor
|
|
import numpy
|
|
import os
|
|
|
|
# Format Details:
|
|
# A KINNE parameter set is stored as a set of files named "snoop_bin_*.bin",
|
|
# where the * is a number starting at 0.
|
|
# Each file is simply raw little-endian floats,
|
|
# as readable by: numpy.fromfile(path, "<f4")
|
|
# and as writable by: t.data.astype("<f4", "C").tofile(path)
|
|
# This format is intended to be extremely simple to get into literally anything.
|
|
# It is not intended to be structural or efficient - reloading a network when
|
|
# unnecessary is inefficient anyway.
|
|
# Ultimately, the idea behind this is as a format that, while it will always
|
|
# require code to implement, requires as little code as possible, and therefore
|
|
# works as a suitable interchange for any situation.
|
|
# To add to the usability of the format, some informal metadata is provided,
|
|
# in "meta.txt", which provides human-readable shape information.
|
|
# This is intended to help with debugging other implementations of the network,
|
|
# by providing concrete human-readable information on tensor shapes.
|
|
# It is NOT meant to be read by machines.
|
|
|
|
class KinneDir:
|
|
"""
|
|
A KinneDir is an intermediate object used to save or load a model.
|
|
"""
|
|
|
|
def __init__(self, base: str, save: bool):
|
|
"""
|
|
Opens a new KINNE directory with the given base path.
|
|
If save is true, the directory is created if possible.
|
|
(This does not create parents.)
|
|
Save being true or false determines if tensors are loaded or saved.
|
|
The base path is of the form "models/abc" - no trailing slash.
|
|
It is important that if you wish to save in the current directory,
|
|
you use ".", not the empty string.
|
|
"""
|
|
if save:
|
|
try:
|
|
os.mkdir(base)
|
|
except:
|
|
# Silence the exception - the directory may (and if reading, does) already exist.
|
|
pass
|
|
self.base = base + "/snoop_bin_"
|
|
self.next_part_index = 0
|
|
self.save = save
|
|
if save:
|
|
self.metadata = open(base + "/meta.txt", "w")
|
|
|
|
def parameter(self, t: Tensor):
|
|
"""
|
|
parameter loads or saves a parameter, given as a tensor.
|
|
"""
|
|
path = f"{self.base}{self.next_part_index}.bin"
|
|
if self.save:
|
|
t.data.astype("<f4", "C").tofile(path)
|
|
self.metadata.write(f"{self.next_part_index}: {t.shape}\n")
|
|
else:
|
|
t.assign(Tensor(numpy.fromfile(path, "<f4")).reshape(shape=t.shape))
|
|
self.next_part_index += 1
|
|
|
|
def parameters(self, params):
|
|
"""
|
|
parameters loads or saves a sequence of parameters.
|
|
It's intended for easily attaching to an existing model,
|
|
assuming that your parameters list orders are consistent.
|
|
(In other words, usage with tinygrad.utils.get_parameters isn't advised -
|
|
it's too 'implicit'.)
|
|
"""
|
|
for t in params:
|
|
self.parameter(t)
|
|
|
|
def close(self):
|
|
if self.save:
|
|
self.metadata.close()
|