fix bug in graph with use of id
MODIFY
test/test_mnist.py
test/test_mnist.py
@@ -51,6 +51,12 @@ class TestMNIST(unittest.TestCase):
optimizer = optim.SGD(model.parameters(), lr=0.001)
train(model, X_train, Y_train, optimizer, BS=69, steps=1)
+ def test_conv_onestep(self):
+ np.random.seed(1337)
+ model = TinyConvNet()
+ optimizer = optim.SGD(model.parameters(), lr=0.001)
+ train(model, X_train, Y_train, optimizer, BS=69, steps=1)
+
def test_conv(self):
np.random.seed(1337)
model = TinyConvNet()

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -22,6 +22,7 @@ if DEBUG:
print(f"{name:>20} : {debug_counts[name]:>6} {debug_times[name]:>10.2f} ms")
atexit.register(print_debug_exit)
+global_num_max = 0
class ProfileOp:
def __init__(self, ctx, name, x, backward=False):
self.ctx, self.name, self.x, self.output, self.backward = ctx, f"back_{name}" if backward else name, x, None, backward
@@ -30,24 +31,30 @@ class ProfileOp:
return self
def __exit__(self, *junk):
if GRAPH:
+ def nm(x):
+ global global_num_max
+ if getattr(x, 'global_num', None) is None:
+ setattr(x, 'global_num', global_num_max)
+ global_num_max += 1
+ return x.global_num
# connect inputs to outputs
for x in self.x:
for y in self.output:
- G.add_edge(id(x.data), id(y.data), label=self.name, color="blue" if self.backward else "black")
- G.nodes[id(x.data)]['label'], G.nodes[id(y.data)]['label'] = str(x.shape), str(y.shape)
+ G.add_edge(nm(x.data), nm(y.data), label=self.name, color="blue" if self.backward else "black")
+ G.nodes[nm(x.data)]['label'], G.nodes[nm(y.data)]['label'] = str(x.shape), str(y.shape)
# which saved tensors does this backward depend on?
saved_tensors = filter(lambda x: any(isinstance(x, v) for v in Device.buffers.values()), self.ctx.saved_tensors)
if self.backward:
for x in saved_tensors:
for y in self.output:
- G.add_edge(id(x), id(y.data), label=self.name, color="red")
+ G.add_edge(nm(x), nm(y.data), label=self.name, color="red")
# did this forward create any intermediate tensors?
if not self.backward:
- x_data = [x.data for x in self.x]
+ x_data = [nm(x.data) for x in self.x]
for y in saved_tensors:
- if y not in x_data: # if intermediate tensors are inputs they don't count
- for x in x_data:
- G.add_edge(id(x), id(y), label=self.name, color="purple")
+ if nm(y) not in x_data: # if intermediate tensors are inputs they don't count
+ for x in self.x:
+ G.add_edge(nm(x.data), nm(y), label=self.name, color="purple")
if DEBUG:
self.output[0].data.toCPU()
et = (time.time()-self.st)*1000.

fix bug in graph with use of id