WIP

waifairer · Oct 20, 2023 · d50a814 · d50a814
1 parent 48e475d
commit d50a814
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/examples/mask_rcnn/train.py b/examples/mask_rcnn/train.py
@@ -163,6 +163,12 @@ def simple():
     img_metadata = coco.loadImgs(random_img_id)[0]
     img_filename = os.path.join(BASEDIR, 'train2017', img_metadata['file_name'])
     print("training", img_filename)
+    print_gpu_memory("before_load")
+    t = Tensor(build_transforms()(Image.open(img_filename).convert("RGB")).numpy(), requires_grad=True).realize()
+    print_gpu_memory("after_load")
+    del t
+    gc.collect()
+    print_gpu_memory("after_del")
     img = [Tensor(build_transforms()(Image.open(img_filename).convert("RGB")).numpy(), requires_grad=True)]
     images = to_image_list(img)
     print_gpu_memory("before images")