commaai
/
tinygrad
mirror of https://github.com/commaai/tinygrad.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fix batchnorm shapes for resnet.load_pretrained (#5167) 

* Fix batchnorm shapes

* make it general reshape
This commit is contained in:
reddyn12 2024-06-26 18:44:10 -04:00 committed by GitHub
parent 396ce6cfc9
commit f1c7944c44
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
extra/models/resnet.py
View File

@ -146,11 +146,8 @@ class ResNet:
print("skipping fully connected layer")
continue # Skip FC if transfer learning
if dat.shape == ():
assert obj.shape == (1,), obj.shape
dat = dat.reshape(1)
assert obj.shape == dat.shape, (k, obj.shape, dat.shape)
obj.assign(dat)
if 'bn' not in k and 'downsample' not in k: assert obj.shape == dat.shape, (k, obj.shape, dat.shape)
obj.assign(dat.reshape(obj.shape))
ResNet18 = lambda num_classes=1000: ResNet(18, num_classes=num_classes)
ResNet34 = lambda num_classes=1000: ResNet(34, num_classes=num_classes)