AllenCell / AllenCell/allencell-segmenter-ml

Batch Size: A reasonable default and potential safeguards

Open
#451 2 comments 0 reactions 1 assignee Assigned to @benjijamorris View on GitHub
Dominant language
Python
Stars
5
Forks
3
PR merge metrics
No merged PRs in 30d

Description

## Why Batch Size is an issue

We discussed that the current default `batch_size = 1` for the plugin configuration is not reasonable for training an effective deep learning model. The reason for this default was the assumption of limited resources on the end-user's side, i.e. someone using the plugin will most likely not have sufficient resources to train a medium to large model (`1 GPU w/16GB, 16 cores, 64GB of RAM`).

#### Theoretical Context

Ideally we want the model to see several samples per update, but if the batch size is one, the model updates after every image---leading to noisy gradient estimates and potentially unstable training. In contrast, a batch size of 16 means the model sees 16 samples before a single update, resulting in more stable and reliable gradients, which generally improves convergence and training stability.

**Limitations:** Larger batch sizes enables faster training times per epoch. However, this is constrained by the memory capacity of the hardware, not everyone has an A100 Machine or even a 4090. Huge batch sizes can lead to out-of-memory errors both on the GPU and CPU/RAM side.

> We want to minimize errors so that people do not face an error and throw the tool away - Matheus 2024

### Proposed Solution

Keep the default batch size at 1, but expose it as a parameter in the plugin that user's can set. We should set an upper bound that is based on the batch size, GPU RAM size, number of CPU cores and RAM. We can also link to some documentation that helps them better understand what a batch means and how to set this given their hardware.

There's also additional overhead like the model size, pre-post processing, and system processes. Maybe I am overthinking this, but a simple buffer or safety factor term that accounts for the overhead. A simple equation that takes the `min(GPU_Memory/Memory_per_Image, Available_RAM/Memory_per_Image)`, and that's the maximum batch size for that system. Idk just throwing something out there.

### Timeline

This is probably for version 1.2 or later. This was the extent of our discussion and I hope this captures what we discussed from an ML perspective. Definitely feel free to add comments so we make optimal design choices. 🚀

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.