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

more enet

This commit is contained in:
George Hotz 2020-10-27 19:37:21 -07:00
parent 41828d768f
commit 9166eb58bb
1 changed files with 25 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
examples/efficientnet.py
View File

@ -3,6 +3,7 @@
# https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b0-355c32eb.pth
# a rough copy of
# https://github.com/lukemelas/EfficientNet-PyTorch/blob/master/efficientnet_pytorch/model.py
from tinygrad.tensor import Tensor
class BatchNorm2D:
def __init__(self, sz):
@ -15,12 +16,16 @@ class BatchNorm2D:
return x * self.weight + self.bias
class MBConvBlock:
def __init__(self, input_filters, expand_ratio, se_ratio, output_filters):
def __init__(self, kernel_size, strides, expand_ratio, input_filters, output_filters, se_ratio):
oup = expand_ratio * input_filters
if expand_ratio != 1:
self._expand_conv = Tensor.zeros(oup, input_filters, 1, 1)
self._bn0 = BatchNorm2D(oup)
self._depthwise_conv = Tensor.zeros(oup, 1, 3, 3)
self.pad = (kernel_size-1)//2
self.strides = strides
self._depthwise_conv = Tensor.zeros(oup, 1, kernel_size, kernel_size)
self._bn1 = BatchNorm2D(oup)
num_squeezed_channels = max(1, int(input_filters * se_ratio))
@ -34,7 +39,8 @@ class MBConvBlock:
def __call__(self, x):
x = self._bn0(x.conv2d(self._expand_conv)).swish()
x = self._bn1(x.conv2d(self._depthwise_conv)).swish() # TODO: repeat on axis 1
x = x.pad(self.pad, self.pad, self.pad, self.pad)
x = self._bn1(x.conv2d(self._depthwise_conv, stride=self.stride)).swish() # TODO: repeat on axis 1
# has_se
x_squeezed = x.avg_pool2d()
@ -49,9 +55,20 @@ class EfficientNet:
def __init__(self):
self._conv_stem = Tensor.zeros(32, 3, 3, 3)
self._bn0 = BatchNorm2D(32)
blocks_args = [
[1, 3, (1,1), 1, 32, 16, 0.25],
[2, 3, (2,2), 6, 16, 24, 0.25],
[2, 5, (2,2), 6, 24, 40, 0.25],
[3, 3, (2,2), 6, 40, 80, 0.25],
[3, 5, (1,1), 6, 80, 112, 0.25],
[4, 5, (1,1), 6, 112, 192, 0.25],
[1, 3, (1,1), 6, 192, 320, 0.25],
]
self._blocks = []
# TODO: create blocks
# num_repeats, kernel_size, strides, expand_ratio, input_filters, output_filters, se_ratio
for b in blocks_args:
for n in range(b[0]):
self._blocks.append(MBConvBlock(*b[1:]))
self._conv_head = Tensor.zeros(1280, 320, 1, 1)
self._bn1 = BatchNorm2D(1280)
self._fc = Tensor.zeros(1280, 1000)
@ -65,3 +82,6 @@ class EfficientNet:
x = x.dropout(0.2)
return x.dot(self_fc).swish()
if __name__ == "__main__":
model = EfficientNet()