clean up resnet
MODIFY
extra/utils.py
extra/utils.py
@@ -108,3 +108,14 @@ def fake_torch_load(b0):
np_array.strides = real_strides
return ret
+
+def get_child(parent, key):
+ obj = parent
+ for k in key.split('.'):
+ if k.isnumeric():
+ obj = obj[int(k)]
+ elif isinstance(obj, dict):
+ obj = obj[k]
+ else:
+ obj = getattr(obj, k)
+ return obj

MODIFY
models/efficientnet.py
models/efficientnet.py
@@ -2,7 +2,7 @@ import math
import numpy as np
from tinygrad.tensor import Tensor
from tinygrad.nn import BatchNorm2D
-from extra.utils import fetch, fake_torch_load
+from extra.utils import fetch, fake_torch_load, get_child
class MBConvBlock:
def __init__(self, kernel_size, strides, expand_ratio, input_filters, output_filters, se_ratio, has_se):
@@ -133,14 +133,6 @@ class EfficientNet:
7: "https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b7-dcc49843.pth"
}
- def _get_child(parent, key):
- obj = parent
- for k in key.split('.'):
- if k.isnumeric():
- obj = obj[int(k)]
- else:
- obj = getattr(obj, k)
- return obj
b0 = fake_torch_load(fetch(model_urls[self.number]))
for k,v in b0.items():
@@ -150,7 +142,7 @@ class EfficientNet:
k = k.replace('.weight', '')
#print(k, v.shape)
- mv = _get_child(self, k)
+ mv = get_child(self, k)
vnp = v.astype(np.float32)
vnp = vnp if k != '_fc' else vnp.T
vnp = vnp if vnp.shape != () else np.array([vnp])

MODIFY
models/resnet.py
models/resnet.py
from tinygrad.tensor import Tensor
import tinygrad.nn as nn
-from extra.utils import fetch, fake_torch_load
-from torch.hub import load_state_dict_from_url
+from extra.utils import fetch, fake_torch_load, get_child
import numpy as np
-model_urls = {
- 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
- 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
- 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
- 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
- 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
-}
-
-def load_from_pretrained(model, url):
- state_dict = load_state_dict_from_url(url, progress=True)
- #state_dict = fake_torch_load(fetch(url))
- layers_not_loaded = []
- for k, v in state_dict.items():
- par_name = ['model']
- for kk in k.split('.'):
- if kk.isdigit():
- par_name += [f'layers[{int(kk)}]']
- else:
- par_name += [kk]
- par_name = '.'.join(par_name)
- code = f"""
-if np.prod({par_name}.shape) == np.prod(v.shape):\n
- if "fc.weight" in par_name:\n
- {par_name}.assign(Tensor(v.detach().numpy().T))\n
- else:\n
- {par_name}.assign(Tensor(v.detach().numpy()))\n
-else:\n
- layers_not_loaded += [k]"""
- exec(code)
- print(f'Loaded from "{url}".')
- if len(layers_not_loaded) > 0:
- for l in layers_not_loaded:
- print(f'- Layer {l} not loaded.')
- return model
-
class BasicBlock:
expansion = 1
@@ -88,7 +52,8 @@ class Bottleneck:
return out
class ResNet:
- def __init__(self, block, num_blocks, num_classes=10):
+ def __init__(self, block, num_blocks, num_classes=10, url=None):
+ self.url = url
self.in_planes = 64
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, bias=False, padding=3)
@@ -97,7 +62,7 @@ class ResNet:
self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
- self.fc = (Tensor.uniform(512 * block.expansion, num_classes), Tensor.zeros(num_classes))
+ self.fc = {"weight": Tensor.uniform(512 * block.expansion, num_classes), "bias": Tensor.zeros(num_classes)}
def _make_layer(self, block, planes, num_blocks, stride):
strides = [stride] + [1] * (num_blocks-1)
@@ -114,38 +79,25 @@ class ResNet:
out = out.sequential(self.layer3)
out = out.sequential(self.layer4)
out = out.mean(3).mean(2)
- out = out.linear(*self.fc).logsoftmax()
+ out = out.linear(**self.fc).logsoftmax()
return out
def __call__(self, x):
return self.forward(x)
-def ResNet18(num_classes, pretrained=False):
- model = ResNet(BasicBlock, [2, 2, 2, 2], num_classes)
- #if pretrained:
- # model = load_from_pretrained(model, model_urls['resnet18'])
- return model
-
-def ResNet34(num_classes, pretrained=False):
- model = ResNet(BasicBlock, [3, 4, 6, 3], num_classes)
- #if pretrained:
- # model = load_from_pretrained(model, model_urls['resnet34'])
- return model
-
-def ResNet50(num_classes, pretrained=False):
- model = ResNet(Bottleneck, [3, 4, 6, 3], num_classes)
- #if pretrained:
- # model = load_from_pretrained(model, model_urls['resnet50'])
- return model
+ def load_from_pretrained(self):
+ # TODO: replace with fake torch load
+ from torch.hub import load_state_dict_from_url
+ state_dict = load_state_dict_from_url(self.url, progress=True)
+ for k, v in state_dict.items():
+ obj = get_child(self, k)
+ dat = v.detach().numpy().T if "fc.weight" in k else v.detach().numpy()
+ assert obj.shape == dat.shape
+ obj.assign(dat)
-def ResNet101(num_classes, pretrained=False):
- model = ResNet(Bottleneck, [3, 4, 23, 3], num_classes)
- #if pretrained:
- # model = load_from_pretrained(model, model_urls['resnet101'])
- return model
+ResNet18 = lambda: ResNet(BasicBlock, [2,2,2,2], 1000, 'https://download.pytorch.org/models/resnet18-5c106cde.pth')
+ResNet34 = lambda: ResNet(BasicBlock, [3,4,6,3], 1000, 'https://download.pytorch.org/models/resnet34-333f7ec4.pth')
+ResNet50 = lambda: ResNet(Bottleneck, [3,4,6,3], 1000, 'https://download.pytorch.org/models/resnet50-19c8e357.pth')
+ResNet101 = lambda: ResNet(Bottleneck, [3,4,23,3], 1000, 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth')
+ResNet101 = lambda: ResNet(Bottleneck, [3,8,36,3], 1000, 'https://download.pytorch.org/models/resnet152-b121ed2d.pth')
-def ResNet152(num_classes, pretrained=False):
- model = ResNet(Bottleneck, [3, 8, 36, 3], num_classes)
- #if pretrained:
- # model = load_from_pretrained(model, model_urls['resnet152'])
- return model

MODIFY
test/test_train.py
test/test_train.py
@@ -7,7 +7,8 @@ from tinygrad.tensor import Tensor
from extra.training import train
from extra.utils import get_parameters
from models.efficientnet import EfficientNet
-from models.transformer import Transformer, ViT
+from models.transformer import Transformer
+from models.vit import ViT
from models.resnet import ResNet18
BS = int(os.getenv("BS", "4"))
@@ -49,7 +50,8 @@ class TestTrain(unittest.TestCase):
X = np.zeros((BS, 3, 224, 224), dtype=np.float32)
Y = np.zeros((BS), dtype=np.int32)
for resnet_v in [ResNet18]:
- model = resnet_v(num_classes=1000, pretrained=True)
+ model = resnet_v()
+ model.load_from_pretrained()
train_one_step(model, X, Y)
def test_bert(self):

clean up resnet