def __init__(self, syms, maxlen, cnt, embed_dim, num_heads):
self.maxlen, self.syms = maxlen, syms
self.embed = Tensor.uniform(maxlen + syms, embed_dim)
self.tbs = []
for i in range(cnt):
self.tbs.append(TransformerBlock(embed_dim, num_heads))
self.final = Tensor.uniform(embed_dim, syms)
def __init__(self, kernel_size, strides, expand_ratio, input_filters, output_filters, se_ratio, has_se):
oup = expand_ratio * input_filters
if expand_ratio != 1:
self._expand_conv = Tensor.uniform(oup, input_filters, 1, 1)
self._bn0 = BatchNorm2D(oup)
else:
self._expand_conv = None
self.strides = strides
if strides == (2,2):
self.pad = [(kernel_size-1)//2-1, (kernel_size-1)//2]*2
else:
self.pad = [(kernel_size-1)//2]*4
self._depthwise_conv = Tensor.uniform(oup, 1, kernel_size, kernel_size)
self._bn1 = BatchNorm2D(oup)
self.has_se = has_se
if self.has_se:
num_squeezed_channels = max(1, int(input_filters * se_ratio))
self._se_reduce = Tensor.uniform(num_squeezed_channels, oup, 1, 1)
self._se_reduce_bias = Tensor.zeros(num_squeezed_channels)
self._se_expand = Tensor.uniform(oup, num_squeezed_channels, 1, 1)
self._se_expand_bias = Tensor.zeros(oup)
self._project_conv = Tensor.uniform(output_filters, oup, 1, 1)
self._bn2 = BatchNorm2D(output_filters)
def __init__(self, classes=10): conv = 3 inter_chan, out_chan = 8, 16 # for speed self.c1 = Tensor.uniform(inter_chan,3,conv,conv) #self.bn1 = BatchNorm2D(inter_chan) self.c2 = Tensor.uniform(out_chan,inter_chan,conv,conv) #self.bn2 = BatchNorm2D(out_chan) self.l1 = Tensor.uniform(out_chan*6*6, classes)
def __init__(self): # https://keras.io/examples/vision/mnist_convnet/ conv = 3 #inter_chan, out_chan = 32, 64 inter_chan, out_chan = 8, 16 # for speed self.c1 = Tensor.uniform(inter_chan,1,conv,conv) self.c2 = Tensor.uniform(out_chan,inter_chan,conv,conv) self.l1 = Tensor.uniform(out_chan*5*5, 10)
def __init__(self):
self.conv = [
ConvBlock(28, 28, 1),
ConvBlock(28, 28, 128),
ConvBlock(14, 14, 128)
]
self.weight1 = Tensor.uniform(128, 10)
self.weight2 = Tensor.uniform(128, 10)
def __init__(self, number=0, classes=1000, has_se=True):
self.number = number
global_params = [
# width, depth
(1.0, 1.0), # b0
(1.0, 1.1), # b1
(1.1, 1.2), # b2
(1.2, 1.4), # b3
(1.4, 1.8), # b4
(1.6, 2.2), # b5
(1.8, 2.6), # b6
(2.0, 3.1), # b7
(2.2, 3.6), # b8
(4.3, 5.3), # l2
][number]
def round_filters(filters):
multiplier = global_params[0]
divisor = 8
filters *= multiplier
new_filters = max(divisor, int(filters + divisor / 2) // divisor * divisor)
if new_filters < 0.9 * filters: # prevent rounding by more than 10%
new_filters += divisor
return int(new_filters)
def round_repeats(repeats):
return int(math.ceil(global_params[1] * repeats))
out_channels = round_filters(32)
self._conv_stem = Tensor.uniform(out_channels, 3, 3, 3)
self._bn0 = BatchNorm2D(out_channels)
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, (2,2), 6, 112, 192, 0.25],
[1, 3, (1,1), 6, 192, 320, 0.25],
]
# num_repeats, kernel_size, strides, expand_ratio, input_filters, output_filters, se_ratio
self._blocks = []
for b in blocks_args:
args = b[1:]
args[3] = round_filters(args[3])
args[4] = round_filters(args[4])
for n in range(round_repeats(b[0])):
self._blocks.append(MBConvBlock(*args, has_se=has_se))
args[3] = args[4]
args[1] = (1,1)
in_channels = round_filters(320)
out_channels = round_filters(1280)
self._conv_head = Tensor.uniform(out_channels, in_channels, 1, 1)
self._bn1 = BatchNorm2D(out_channels)
self._fc = Tensor.uniform(out_channels, classes)
self._fc_bias = Tensor.zeros(classes)
def __init__(self, h, w, inp, filters=128, conv=3): self.h, self.w = h, w self.inp = inp #init weights self.cweights = [Tensor.uniform(filters, inp if i==0 else filters, conv, conv) for i in range(3)] self.cbiases = [Tensor.uniform(1, filters, 1, 1) for i in range(3)] #init layers self._bn = BatchNorm2D(128, training=True) self._seb = SqueezeExciteBlock2D(filters)
def __init__(self, syms, maxlen, layers, embed_dim, num_heads, ff_dim):
self.maxlen, self.syms = maxlen, syms
self.embed = Tensor.uniform(maxlen + syms,
embed_dim,
requires_grad=False)
self.tbs = []
for i in range(layers):
self.tbs.append(TransformerBlock(embed_dim, num_heads, ff_dim))
self.final = Tensor.uniform(embed_dim, syms)
def __init__(self, classes, saved=None):
conv = 3
inter_chan, out_chan = 8, 16
if saved:
self.c1 = Tensor(saved["arr_0"])
self.c2 = Tensor(saved["arr_1"])
self.l1 = Tensor(saved["arr_2"])
else:
self.c1 = to_gpu(Tensor.uniform(inter_chan, 4, conv, conv))
self.c2 = to_gpu(Tensor.uniform(out_chan, inter_chan, conv, conv))
self.l1 = to_gpu(Tensor.uniform(out_chan * 7 * 7, classes))
def __init__(self):
self.blocks = 3
self.block_convs = 3
self.chans = 128
self.convs = [
Tensor.uniform(self.chans, self.chans if i > 0 else 1, 3, 3)
for i in range(self.blocks * self.block_convs)
]
# TODO: Make batchnorm work at train time
#self.bn = [BatchNorm2D(self.chans) for i in range(3)]
self.fc = Tensor.uniform(self.chans, 10)
def __init__(self, inC, outC, last = False):
# Massively overstate the weights to get them to be focused on,
# since otherwise the biases overrule everything
self.weight = Tensor.uniform(outC, inC, 3, 3) * 16.0
# Layout-wise, blatant cheat, but serious_mnist does it. I'd guess channels either have to have a size of 1 or whatever the target is?
# Values-wise, entirely different blatant cheat.
# In most cases, use uniform bias, but tiny.
# For the last layer, use just 0.5, constant.
if last:
self.bias = Tensor.zeros(1, outC, 1, 1) + 0.5
else:
self.bias = Tensor.uniform(1, outC, 1, 1)
def __init__(self, layers=12, embed_dim=192, num_heads=3):
self.embedding = (Tensor.uniform(embed_dim, 3, 16,
16), Tensor.zeros(embed_dim))
self.embed_dim = embed_dim
self.cls = Tensor.ones(1, 1, embed_dim)
self.pos_embedding = Tensor.ones(1, 197, embed_dim)
self.tbs = [
TransformerBlock(embed_dim=embed_dim,
num_heads=num_heads,
ff_dim=embed_dim * 4,
prenorm=True,
act=lambda x: x.gelu()) for i in range(layers)
]
self.encoder_norm = (Tensor.uniform(embed_dim),
Tensor.zeros(embed_dim))
self.head = (Tensor.uniform(embed_dim, 1000), Tensor.zeros(1000))
def test_conv2d(self):
BS, C1, H, W = 4, 16, 224, 224
C2, K, S, P = 64, 7, 2, 1
# create in tinygrad
layer = Conv2d(C1, C2, kernel_size=K, stride=S, padding=P)
# create in torch
with torch.no_grad():
torch_layer = torch.nn.Conv2d(C1,
C2,
kernel_size=K,
stride=S,
padding=P).eval()
torch_layer.weight[:] = torch.tensor(layer.weight.data,
dtype=torch.float32)
torch_layer.bias[:] = torch.tensor(layer.bias.data,
dtype=torch.float32)
# test
x = Tensor.uniform(BS, C1, H, W)
z = layer(x)
torch_x = torch.tensor(x.cpu().data)
torch_z = torch_layer(torch_x)
np.testing.assert_allclose(z.data,
torch_z.detach().numpy(),
atol=5e-4,
rtol=1e-5)
def __init__(self):
self.blocks = 3
self.block_convs = 3
# TODO: raise back to 128 when it's fast
self.chans = 32
self.convs = [
Tensor.uniform(self.chans, self.chans if i > 0 else 1, 3, 3)
for i in range(self.blocks * self.block_convs)
]
self.cbias = [
Tensor.uniform(1, self.chans, 1, 1)
for i in range(self.blocks * self.block_convs)
]
self.bn = [BatchNorm2D(self.chans, training=True) for i in range(3)]
self.fc1 = Tensor.uniform(self.chans, 10)
self.fc2 = Tensor.uniform(self.chans, 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)
self.bn1 = nn.BatchNorm2D(64)
self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=2)
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 = {"weight": Tensor.uniform(512 * block.expansion, num_classes), "bias": Tensor.zeros(num_classes)}
def __init__(self,
in_channels,
out_channels,
kernel_size,
stride=1,
padding=0,
groups=1,
bias=True):
self.in_channels, self.out_channels, self.stride, self.padding, self.groups, self.bias = in_channels, out_channels, stride, padding, groups, bias # Wow this is terrible
self.kernel_size = (kernel_size, kernel_size) if isinstance(
kernel_size, int) else kernel_size
assert out_channels % groups == 0 and in_channels % groups == 0
self.weight = Tensor.uniform(out_channels, in_channels // groups,
*self.kernel_size)
if self.bias:
self.bias = Tensor.uniform(1, out_channels, 1, 1)
else:
self.bias = None
def __init__(self,
embed_dim,
num_heads,
ff_dim,
prenorm=False,
act=lambda x: x.relu()):
self.num_heads = num_heads
self.head_size = embed_dim // num_heads
assert self.head_size * self.num_heads == embed_dim
self.prenorm, self.act = prenorm, act
self.query = (Tensor.uniform(embed_dim,
embed_dim), Tensor.zeros(embed_dim))
self.key = (Tensor.uniform(embed_dim,
embed_dim), Tensor.zeros(embed_dim))
self.value = (Tensor.uniform(embed_dim,
embed_dim), Tensor.zeros(embed_dim))
self.out = (Tensor.uniform(embed_dim,
embed_dim), Tensor.zeros(embed_dim))
self.ff1 = (Tensor.uniform(embed_dim, ff_dim), Tensor.zeros(ff_dim))
self.ff2 = (Tensor.uniform(ff_dim, embed_dim), Tensor.zeros(embed_dim))
self.ln1 = (Tensor.ones(embed_dim), Tensor.zeros(embed_dim))
self.ln2 = (Tensor.ones(embed_dim), Tensor.zeros(embed_dim))
def __init__(self, num, num_classes):
self.num = num
self.block = {
18: BasicBlock,
34: BasicBlock,
50: Bottleneck,
101: Bottleneck,
152: Bottleneck
}[num]
self.num_blocks = {
18: [2, 2, 2, 2],
34: [3, 4, 6, 3],
50: [3, 4, 6, 3],
101: [3, 4, 23, 3],
152: [3, 8, 36, 3]
}[num]
self.in_planes = 64
self.conv1 = nn.Conv2d(3,
64,
kernel_size=7,
stride=2,
bias=False,
padding=3)
self.bn1 = nn.BatchNorm2D(64)
self.layer1 = self._make_layer(self.block,
64,
self.num_blocks[0],
stride=2)
self.layer2 = self._make_layer(self.block,
128,
self.num_blocks[1],
stride=2)
self.layer3 = self._make_layer(self.block,
256,
self.num_blocks[2],
stride=2)
self.layer4 = self._make_layer(self.block,
512,
self.num_blocks[3],
stride=2)
self.fc = {
"weight": Tensor.uniform(512 * self.block.expansion, num_classes),
"bias": Tensor.zeros(num_classes)
}
def __init__(self, embed_dim, num_heads):
# Multi-Head Attention
self.num_heads = num_heads
self.head_size = embed_dim // num_heads
assert self.head_size * self.num_heads == embed_dim
# looks like bias is useless
self.query_dense = Tensor.uniform(embed_dim, embed_dim)
self.key_dense = Tensor.uniform(embed_dim, embed_dim)
self.value_dense = Tensor.uniform(embed_dim, embed_dim)
self.final = Tensor.uniform(embed_dim, embed_dim)
self.ff1 = Tensor.uniform(embed_dim, embed_dim)
self.ff2 = Tensor.uniform(embed_dim, embed_dim)
def _test_linear(x):
# create in tinygrad
layer = (Tensor.uniform(in_dim, out_dim), Tensor.zeros(out_dim))
z = x.linear(*layer)
# create in torch
with torch.no_grad():
torch_layer = torch.nn.Linear(in_dim, out_dim).eval()
torch_layer.weight[:] = torch.tensor(layer[0].data.T,
dtype=torch.float32)
torch_layer.bias[:] = torch.tensor(layer[1].data,
dtype=torch.float32)
torch_x = torch.tensor(x.cpu().data, dtype=torch.float32)
torch_z = torch_layer(torch_x)
# test
np.testing.assert_allclose(z.data,
torch_z.detach().numpy(),
atol=5e-4,
rtol=1e-5)
def __init__(self, filters): self.filters = filters self.weight1 = Tensor.uniform(self.filters, self.filters//32) self.bias1 = Tensor.uniform(1,self.filters//32) self.weight2 = Tensor.uniform(self.filters//32, self.filters) self.bias2 = Tensor.uniform(1, self.filters)
def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, bias=True): self.kernel_size = (kernel_size, kernel_size) if isinstance(kernel_size, int) else (kernel_size[0], kernel_size[1]) self.stride = (stride, stride) if isinstance(stride, int) else (stride[0], stride[1]) self.padding = (padding, ) * 4 if isinstance(padding, int) else (padding[0], padding[0], padding[1], padding[1]) self.weight = Tensor.uniform(out_channels, in_channels, self.kernel_size[0], self.kernel_size[1]) self.bias = Tensor.uniform(out_channels) if bias else None
def __init__(self): self.l1 = Tensor.uniform(784, 128) self.l2 = Tensor.uniform(128, 10)
def __init__(self): lv = 128 self.l1 = Tensor.uniform(128, 256) self.l2 = Tensor.uniform(256, 512) self.l3 = Tensor.uniform(512, 1024) self.l4 = Tensor.uniform(1024, 784)
def __init__(self): in_sh = 784 self.l1 = Tensor.uniform(784, 1024) self.l2 = Tensor.uniform(1024, 512) self.l3 = Tensor.uniform(512, 256) self.l4 = Tensor.uniform(256, 2)
