From 908db3bdeae4a04da2e042dcb93be784ce4e5ed5 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Tue, 30 Nov 2021 00:44:59 -0500 Subject: [PATCH] support bias in conv like linear --- models/efficientnet.py | 4 ++-- tinygrad/tensor.py | 4 ++++ 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/models/efficientnet.py b/models/efficientnet.py index 9f8f976ab8..891add39b4 100644 --- a/models/efficientnet.py +++ b/models/efficientnet.py @@ -44,8 +44,8 @@ class MBConvBlock: # has_se if self.has_se: x_squeezed = x.avg_pool2d(kernel_size=x.shape[2:4]) - x_squeezed = x_squeezed.conv2d(self._se_reduce).add(self._se_reduce_bias.reshape(shape=[1, -1, 1, 1])).swish() - x_squeezed = x_squeezed.conv2d(self._se_expand).add(self._se_expand_bias.reshape(shape=[1, -1, 1, 1])) + x_squeezed = x_squeezed.conv2d(self._se_reduce, self._se_reduce_bias).swish() + x_squeezed = x_squeezed.conv2d(self._se_expand, self._se_expand_bias) x = x.mul(x_squeezed.sigmoid()) x = self._bn2(x.conv2d(self._project_conv)) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 062f13de6e..eda8b5b4d4 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -308,6 +308,10 @@ class Tensor: def max_pool2d(self, kernel_size=(2,2)): return self._pool2d(*kernel_size).max(axis=(3,5)) + def conv2d(self, weight, bias=None, stride=1, groups=1): + ret = self._conv2d(weight, stride=stride, groups=groups) + return ret if bias is None else ret.add(bias.reshape(shape=[1, -1, 1, 1])) + # ***** functional nn ops ***** def linear(self, weight, bias):