fix DEBUG for GPU
MODIFY
examples/train_efficientnet.py
examples/train_efficientnet.py
@@ -45,7 +45,7 @@ if __name__ == "__main__":
print("parameter count", len(parameters))
optimizer = optim.Adam(parameters, lr=0.001)
- BS, steps = int(os.getenv("BS", "64" if TINY else "16")), 2048
+ BS, steps = int(os.getenv("BS", "64" if TINY else "16")), int(os.getenv("STEPS", "2048"))
print("training with batch size %d for %d steps" % (BS, steps))
Tensor.training = True

MODIFY
tinygrad/ops/ops_gpu.py
tinygrad/ops/ops_gpu.py
@@ -14,10 +14,10 @@ def require_init_gpu():
cl_ctx = cl.Context(devices=devices)
# this is an in-order command queue
cl_queue = cl.CommandQueue(cl_ctx)
-require_init_gpu()
class GPUBuffer:
def __init__(self, shape, hostbuf=None):
+ require_init_gpu()
self.shape, self.dtype = tuple(shape), np.float32
self.cl = hostbuf.cl if isinstance(hostbuf, GPUBuffer) else \
cl.Buffer(cl_ctx, cl.mem_flags.READ_WRITE | (cl.mem_flags.COPY_HOST_PTR if hostbuf is not None else 0), 4*np.prod(shape),
@@ -32,6 +32,7 @@ class GPUBuffer:
def toCPU(self):
data = np.empty(self.shape, dtype=np.float32)
+ cl_queue.finish()
cl.enqueue_copy(cl_queue, data, self.cl, is_blocking=True)
return data

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -24,9 +24,10 @@ class ProfileOp:
return self
def __exit__(self, *junk):
if DEBUG:
- # TODO: fix this
- #if cl_queue is not None:
- # cl_queue.finish()
+ if self.output is not None:
+ self.output.data.toCPU()
+ elif self.x[0] is not None:
+ self.x[0].data.toCPU()
et = (time.time()-self.st)*1000.
debug_counts[self.name] += 1
debug_times[self.name] += et

fix DEBUG for GPU