Add process_batch_size argument to control memory

pytorch-labs · Nov 30, 2023 · ec9a2bf · ec9a2bf
1 parent 35db2a9
commit ec9a2bf
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 1 deletion.
diff --git a/amg_example/amg_example.py b/amg_example/amg_example.py
@@ -41,7 +41,7 @@ def show_anns(anns):
 
 sam = sam_model_fast_registry[model_type](checkpoint=sam_checkpoint)
 sam.to(device=device)
-mask_generator = SamAutomaticMaskGenerator(sam)
+mask_generator = SamAutomaticMaskGenerator(sam, process_batch_size=8)
 
 # Run thrice for warmup
 masks = mask_generator.generate(image)
@@ -69,3 +69,10 @@ def show_anns(anns):
 
 # Save a GPU trace
 profiler_runner(f"amg_example_trace.json.gz", mask_generator.generate, image)
+
+# Write out memory usage
+max_memory_allocated_bytes = torch.cuda.max_memory_allocated()
+_, total_memory = torch.cuda.mem_get_info()
+max_memory_allocated_percentage = int(100 * (max_memory_allocated_bytes / total_memory))
+max_memory_allocated_bytes = max_memory_allocated_bytes >> 20
+print(f"memory(MiB): {max_memory_allocated_bytes} memory(%): {max_memory_allocated_percentage}")
diff --git a/amg_example/amg_example_trace.json.gz b/amg_example/amg_example_trace.json.gz
diff --git a/amg_example/dog_mask_fast.png b/amg_example/dog_mask_fast.png