diff --git a/README.md b/README.md index 2ea8884..c77a4e3 100644 --- a/README.md +++ b/README.md @@ -44,11 +44,11 @@ Here are models that can be downloaded: | Model | Parameters | Size | Download | |-------------------|------------|-------|------------------------------| -| SAM 91M | 91M | 100MB | `osam run sam:91m` | -| SAM 308M | 308M | 320MB | `osam run sam:308m` | -| SAM 636M | 636M | 630MB | `osam run sam` | -| EfficientSAM 10M | 10M | 40MB | `osam run efficientsam:10m` | -| EfficientSAM 25M | 25M | 100MB | `osam run efficientsam` | +| SAM 100M | 100M | 100MB | `osam run sam:100m` | +| SAM 300M | 300M | 300MB | `osam run sam:300m` | +| SAM 600M | 600M | 600MB | `osam run sam` | +| EfficientSAM 10M | 10M | 40MB | `osam run efficientsam:10m` | +| EfficientSAM 30M | 30M | 100MB | `osam run efficientsam` | PS. `sam`, `efficientsam` is equivalent to `sam:latest`, `efficientsam:latest`. diff --git a/scripts/count_parameter_size.py b/scripts/count_parameter_size.py new file mode 100755 index 0000000..a5d2ec7 --- /dev/null +++ b/scripts/count_parameter_size.py @@ -0,0 +1,35 @@ +#!/usr/bin/env python3 + +import argparse + +from loguru import logger +import onnx + + +def count_parameters(onnx_file): + model = onnx.load(onnx_file) + total_params = 0 + + for initializer in model.graph.initializer: + param_count = 1 + for dim in initializer.dims: + param_count *= dim + total_params += param_count + + return total_params + + +def main(): + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("onnx_file") + args = parser.parse_args() + + logger.debug(f"ONNX file: {args.onnx_file!r}") + num_params = count_parameters(onnx_file=args.onnx_file) + logger.debug(f"Number of parameters: {num_params} = {num_params / 1e6:.2f}M = {num_params / 1e9:.2f}B") + + +if __name__ == "__main__": + main()