Create create_experiments

create_experiments

@@ -0,0 +1,71 @@
+import argparse
+import csv
+
+def generate_ratios(max_layers, max_width):
+    ratios = []
+    for i in range(1, 9):  # 1/8 steps from 1:1 to 1:3 and vice versa
+        ratio = i / 8
+        if ratio <= 3:  # Ensure the ratio is not greater than 3
+            ratios.append(ratio)
+    return ratios
+
+def generate_experiments(max_layers, max_width):
+    experiments = []
+    ratios = generate_ratios(max_layers, max_width)
+
+    for ratio in ratios:
+        for layers in range(1, max_layers + 1):
+            width = int(max_width * ratio)
+            experiments.append((layers, width))
+
+    return experiments
+
+def estimate_vram(layer_count, width, input_size, output_size):
+    # Calculate the number of parameters
+    # Each layer has (input_size * width) + width parameters (weights + biases)
+    # The last layer has (width * output_size) + output_size parameters
+    param_count = 0
+    for i in range(layer_count):
+        if i == 0:
+            param_count += (input_size * width) + width
+        else:
+            param_count += (width * width) + width
+    param_count += (width * output_size) + output_size
+
+    # Estimate the VRAM usage
+    # Parameters: 4 bytes per parameter (FP32)
+    # Activations: Assume the size of the activations is the same as the input size
+    vram_usage = param_count * 4 + input_size * 4
+
+    return vram_usage
+
+def write_csv(experiments, filename):
+    with open(filename, 'w', newline='') as csvfile:
+        writer = csv.writer(csvfile)
+        writer.writerow(['layer_count', 'width', 'vram_usage'])
+        for experiment in experiments:
+            writer.writerow(experiment)
+
+def main():
+    parser = argparse.ArgumentParser(description='Generate a CSV file with a variety of layer counts and widths.')
+    parser.add_argument('--max_layers', type=int, default=4, help='Maximum number of layers (default: 4)')
+    parser.add_argument('--max_width', type=int, default=2048, help='Maximum width (default: 2048)')
+    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=224*224*3, help='Input size (default: 224*224*3)')
+    parser.add_argument('--output_size', type=int, default=10, help='Output size (default: 10)')
+    args = parser.parse_args()
+
+    experiments = generate_experiments(args.max_layers, args.max_width)
+    experiments_with_vram = []
+
+    for experiment in experiments:
+        layer_count, width = experiment
+        vram_usage = estimate_vram(layer_count, width, args.input_size, args.output_size)
+        experiments_with_vram.append((layer_count, width, vram_usage))
+        print(f'Layer Count: {layer_count}, Width: {width}, Estimated VRAM Usage: {vram_usage} bytes')
+
+    write_csv(experiments_with_vram, args.output_file)
+    print(f'Generated {len(experiments_with_vram)} experiments and saved to {args.output_file}')
+
+if __name__ == '__main__':
+    main()