TEST_ENET for openpilot compiler
MODIFY
accel/opencl/ops_opencl.py
accel/opencl/ops_opencl.py
@@ -152,7 +152,7 @@ class OpenCLBuffer(GPUBuffer):
@property
def image(self):
if self._image is None:
- assert self.shape[2] == 4 and len(self.shape) == 3
+ assert len(self.shape) == 3 and self.shape[2] == 4, f"bad shape for image {self.shape}"
self._image = ECL.image(shape=(self.shape[1], self.shape[0]))
if self._buf is not None:
assert prod(self.shape) == prod(self._image.shape)*4

MODIFY
models/efficientnet.py
models/efficientnet.py
@@ -52,7 +52,7 @@ class MBConvBlock:
return x
class EfficientNet:
- def __init__(self, number=0, classes=1000, has_se=True, track_running_stats=True):
+ def __init__(self, number=0, classes=1000, has_se=True, track_running_stats=True, input_channels=3, has_fc_output=True):
self.number = number
global_params = [
# width, depth
@@ -81,7 +81,7 @@ class EfficientNet:
return int(math.ceil(global_params[1] * repeats))
out_channels = round_filters(32)
- self._conv_stem = Tensor.uniform(out_channels, 3, 3, 3)
+ self._conv_stem = Tensor.uniform(out_channels, input_channels, 3, 3)
self._bn0 = BatchNorm2D(out_channels, track_running_stats=track_running_stats)
blocks_args = [
[1, 3, (1,1), 1, 32, 16, 0.25],
@@ -117,8 +117,11 @@ class EfficientNet:
out_channels = round_filters(1280)
self._conv_head = Tensor.uniform(out_channels, in_channels, 1, 1)
self._bn1 = BatchNorm2D(out_channels, track_running_stats=track_running_stats)
+ if has_fc_output:
self._fc = Tensor.uniform(out_channels, classes)
self._fc_bias = Tensor.zeros(classes)
+ else:
+ self._fc = None
def forward(self, x):
x = self._bn0(x.conv2d(self._conv_stem, padding=(0,1,0,1), stride=2)).swish()
@@ -126,7 +129,7 @@ class EfficientNet:
x = self._bn1(x.conv2d(self._conv_head)).swish()
x = x.avg_pool2d(kernel_size=x.shape[2:4])
x = x.reshape(shape=(-1, x.shape[1]))
- return x.linear(self._fc, self._fc_bias)
+ return x.linear(self._fc, self._fc_bias) if self._fc is not None else x
def load_from_pretrained(self):
model_urls = {

MODIFY
openpilot/compile.py
openpilot/compile.py
@@ -57,10 +57,17 @@ def compile(input, output_fn):
Tensor.no_grad = True
using_graph = ops.GRAPH
ops.GRAPH = False
+ inputs, _ = get_random_input_tensors()
+ if os.getenv("TEST_ENET", None) is not None:
+ from models.efficientnet import EfficientNet
+ Tensor.training = False
+ enet = EfficientNet(number=int(os.getenv("TEST_ENET", None)), has_se=False, input_channels=12, has_fc_output=False)
+ def run_onnx(x):
+ return {"outputs": enet.forward(x['input_imgs'])}
+ else:
onnx_model = onnx.load(io.BytesIO(dat))
run_onnx = get_run_onnx(onnx_model)
- inputs, _ = get_random_input_tensors()
# initial run(s) to load weights
for _ in range(2):

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -141,9 +141,7 @@ class Tensor:
def cat(self, *args, dim=0):
dim = (dim + len(self.shape)) if dim < 0 else dim
- for y in args:
- assert len(y.shape) == len(self.shape)
- assert all(y.shape[i] == s for i,s in enumerate(self.shape) if i != dim)
+ for y in args: assert len(y.shape) == len(self.shape) and all(y.shape[i] == s for i,s in enumerate(self.shape) if i != dim)
args = [self] + list(args)
shape_cumsum = [0, *itertools.accumulate(y.shape[dim] for y in args)]
slc = [[(0, s) for s in self.shape] for _ in args]

TEST_ENET for openpilot compiler