Update create_experiments.py

create_experiments.py

@@ -1,5 +1,6 @@
 import argparse
 import csv
+import math
 
 def generate_ratios(max_layers, max_width):
     ratios = []
@@ -39,6 +40,15 @@ def estimate_vram(layer_count, width, input_size, output_size):
 
     return vram_usage
 
+def calculate_batch_size(memory_gb=20):
+    memory_bytes = memory_gb * (1024 ** 3)  # Convert GiB to bytes
+    batch_memory_bytes = memory_bytes / 4  # Divide by 4
+    
+    # Find the nearest power of 2
+    batch_size = 2 ** int(math.log2(batch_memory_bytes))
+    
+    return batch_size
+
 def write_csv(experiments, filename):
     with open(filename, 'w', newline='') as csvfile:
         writer = csv.writer(csvfile)
@@ -55,6 +65,7 @@ def main():
     parser.add_argument('--output_file', type=str, default='experiments.csv', help='Output CSV file (default: experiments.csv)')
     parser.add_argument('--input_size', type=int, default=64*64*3, help='Input size (default: 64*64*3)')
     parser.add_argument('--output_size', type=int, default=10, help='Output size (default: 10)')
+    parser.add_argument('--memory_gb', type=int, default=20, help='Total memory in GiB (default: 20)')
     args = parser.parse_args()
 
     experiments = generate_experiments(args.max_layers, args.max_width, args.min_layers, args.min_width)
@@ -69,5 +80,9 @@ def main():
     write_csv(experiments_with_vram, args.output_file)
     print(f'Generated {len(experiments_with_vram)} experiments and saved to {args.output_file}')
 
+    # Calculate and print the batch size
+    batch_size = calculate_batch_size(args.memory_gb)
+    print(f'Recommended batch size: {batch_size}')
+
 if __name__ == '__main__':
     main()