issue-193 - Move torch loader out of efficientnet code
MODIFY
extra/efficientnet.py
extra/efficientnet.py
@@ -2,73 +2,10 @@ import math
import numpy as np
from tinygrad.tensor import Tensor
from tinygrad.nn import BatchNorm2D
-from extra.utils import fetch
+from extra.utils import fetch, fake_torch_load
USE_TORCH = False
-def fake_torch_load(b0):
- import io
- import pickle
- import struct
-
- # convert it to a file
- fb0 = io.BytesIO(b0)
-
- # skip three junk pickles
- pickle.load(fb0)
- pickle.load(fb0)
- pickle.load(fb0)
-
- key_prelookup = {}
-
- class HackTensor:
- def __new__(cls, *args):
- #print(args)
- ident, storage_type, obj_key, location, obj_size, view_metadata = args[0]
- assert ident == 'storage'
-
- ret = np.zeros(obj_size, dtype=storage_type)
- key_prelookup[obj_key] = (storage_type, obj_size, ret, args[2], args[3])
- return ret
-
- class MyPickle(pickle.Unpickler):
- def find_class(self, module, name):
- #print(module, name)
- if name == 'FloatStorage':
- return np.float32
- if name == 'LongStorage':
- return np.int64
- if module == "torch._utils" or module == "torch":
- return HackTensor
- else:
- return pickle.Unpickler.find_class(self, module, name)
-
- def persistent_load(self, pid):
- return pid
-
- ret = MyPickle(fb0).load()
-
- # create key_lookup
- key_lookup = pickle.load(fb0)
- key_real = [None] * len(key_lookup)
- for k,v in key_prelookup.items():
- key_real[key_lookup.index(k)] = v
-
- # read in the actual data
- for storage_type, obj_size, np_array, np_shape, np_strides in key_real:
- ll = struct.unpack("Q", fb0.read(8))[0]
- assert ll == obj_size
- bytes_size = {np.float32: 4, np.int64: 8}[storage_type]
- mydat = fb0.read(ll * bytes_size)
- np_array[:] = np.frombuffer(mydat, storage_type)
- np_array.shape = np_shape
-
- # numpy stores its strides in bytes
- real_strides = tuple([x*bytes_size for x in np_strides])
- np_array.strides = real_strides
-
- return ret
-
class MBConvBlock:
def __init__(self, kernel_size, strides, expand_ratio, input_filters, output_filters, se_ratio, has_se):
oup = expand_ratio * input_filters

MODIFY
extra/utils.py
extra/utils.py
@@ -25,3 +25,67 @@ def get_parameters(obj):
for k,v in obj.__dict__.items():
parameters.extend(get_parameters(v))
return parameters
+
+def fake_torch_load(b0):
+ import io
+ import pickle
+ import struct
+ import numpy as np
+
+ # convert it to a file
+ fb0 = io.BytesIO(b0)
+
+ # skip three junk pickles
+ pickle.load(fb0)
+ pickle.load(fb0)
+ pickle.load(fb0)
+
+ key_prelookup = {}
+
+ class HackTensor:
+ def __new__(cls, *args):
+ #print(args)
+ ident, storage_type, obj_key, location, obj_size, view_metadata = args[0]
+ assert ident == 'storage'
+
+ ret = np.zeros(obj_size, dtype=storage_type)
+ key_prelookup[obj_key] = (storage_type, obj_size, ret, args[2], args[3])
+ return ret
+
+ class MyPickle(pickle.Unpickler):
+ def find_class(self, module, name):
+ #print(module, name)
+ if name == 'FloatStorage':
+ return np.float32
+ if name == 'LongStorage':
+ return np.int64
+ if module == "torch._utils" or module == "torch":
+ return HackTensor
+ else:
+ return pickle.Unpickler.find_class(self, module, name)
+
+ def persistent_load(self, pid):
+ return pid
+
+ ret = MyPickle(fb0).load()
+
+ # create key_lookup
+ key_lookup = pickle.load(fb0)
+ key_real = [None] * len(key_lookup)
+ for k,v in key_prelookup.items():
+ key_real[key_lookup.index(k)] = v
+
+ # read in the actual data
+ for storage_type, obj_size, np_array, np_shape, np_strides in key_real:
+ ll = struct.unpack("Q", fb0.read(8))[0]
+ assert ll == obj_size
+ bytes_size = {np.float32: 4, np.int64: 8}[storage_type]
+ mydat = fb0.read(ll * bytes_size)
+ np_array[:] = np.frombuffer(mydat, storage_type)
+ np_array.shape = np_shape
+
+ # numpy stores its strides in bytes
+ real_strides = tuple([x*bytes_size for x in np_strides])
+ np_array.strides = real_strides
+
+ return ret

issue-193 - Move torch loader out of efficientnet code (#213)