# Model Guide


# Overview

cellclass includes a curated set of convolutional architectures suited to small fluorescence microscopy crops (typically 50 × 50 px, 1-3 channels). The models fall into two groups: **lightweight baselines** that train quickly and are easy to deploy, and **pretrained transfer-learning models** that leverage ImageNet features when labelled data is scarce.

------------------------------------------------------------------------


# Recommended models for small cell images


## Lightweight (no pretrained weights)

These models train from scratch and are well suited to grayscale or 2-channel images where ImageNet features are less relevant.

| Model | Key strengths | Recommended batch size | Notes |
|----|----|----|----|
| [shufflenet2x1_0](../reference/Model.md#cellclass.Model.shufflenet2x1_0) | Fast, memory-efficient, excellent accuracy/speed trade-off | 128 | **First choice** for most cell classification tasks |
| [shufflenet2x1_5](../reference/Model.md#cellclass.Model.shufflenet2x1_5) | Slightly wider than `2x1_0`, marginal accuracy gain | 64 | Use when `2x1_0` plateaus |
| [shufflenet2x2_0](../reference/Model.md#cellclass.Model.shufflenet2x2_0) | Widest ShuffleNet, highest capacity | 32 | Use for complex multiclass problems |
| [squeezenet1_0](../reference/Model.md#cellclass.Model.squeezenet1_0) | Smallest model, fastest inference | 128 | Good sanity-check / baseline |
| [densenet121](../reference/Model.md#cellclass.Model.densenet121) | Strong regularisation via dense connections | 64 | Reliable performer across datasets |
| [densenet161](../reference/Model.md#cellclass.Model.densenet161) | More capacity than 121 | 32 | Diminishing returns on 50 px crops |


## Pretrained (ImageNet transfer learning)

Use `--freeze-weights` to freeze the backbone and only train the classification head. This is beneficial when labelled data is limited (\< ~500 crops per class).

| Model | Input size | Key strengths | Notes |
|----|----|----|----|
| [efficientnetb3s](../reference/Model.md#cellclass.Model.efficientnetb3s) | 300 × 300 (resized automatically) | Best accuracy with pretrained features | **Recommended pretrained model** |
| [efficientnetb4s](../reference/Model.md#cellclass.Model.efficientnetb4s) | 380 × 380 | Highest accuracy | Slower; use when `b3s` is insufficient |

> **Note on image resizing**: EfficientNet models internally resize inputs. For 50 × 50 px crops this upscaling can hurt performance compared to architectures that accept any input size. If you have ≥ 1000 crops per class, [densenet121](../reference/Model.md#cellclass.Model.densenet121) often matches or exceeds EfficientNet.


## Attention models

The following models add spatial or channel attention modules to GoogLeNet- or AlexNet-based backbones. They are experimental and have shown mixed results on small crops.

| Model | Description |
|----|----|
| [googlenet](../reference/Model.md#cellclass.Model.googlenet) / [googlenet_att](../reference/Model.md#cellclass.Model.googlenet_att) | GoogLeNet with optional spatial attention |
| [alexnet](../reference/Model.md#cellclass.Model.alexnet) / [alexnet_att](../reference/Model.md#cellclass.Model.alexnet_att) | AlexNet with optional spatial attention |
| [simplegooglenet](../reference/Model.md#cellclass.Model.simplegooglenet) | Lightweight GoogLeNet using a subset of inception blocks |
| [simplegooglenet_att](../reference/Model.md#cellclass.Model.simplegooglenet_att) | \+ spatial attention |
| [simplegooglenet_ch_att](../reference/Model.md#cellclass.Model.simplegooglenet_ch_att) | \+ channel attention |
| [simplegooglenet_ch_sp_att](../reference/Model.md#cellclass.Model.simplegooglenet_ch_sp_att) | \+ both channel and spatial attention |

------------------------------------------------------------------------


# Choosing a model: decision tree

    Do you have ≥ 500 labelled crops per class?
    ├── Yes → Start with shufflenet2x1_0 (fast sweep) + densenet121
    │         Compare on W&B; pick best.
    └── No  → Use efficientnetb3s --freeze-weights
              Transfer learning compensates for small datasets.

    Is inference speed critical (production deployment)?
    ├── Yes → shufflenet2x1_0 or squeezenet1_0
    └── No  → Any of the above

    Do you have > 3 classes or very similar morphologies?
    └── Yes → Try densenet161 or shufflenet2x2_0 for higher capacity

------------------------------------------------------------------------


# Hyperparameter guidance


## Learning rate

Always sweep at least two values: `1e-3` and `1e-4`. The plateau scheduler adapts automatically, so a higher starting LR rarely hurts.


## Batch size

Larger batches train faster but may reduce generalisation on small datasets. Use the recommended defaults from the table above; halve the batch size if you run out of memory.


## Data augmentation (`--flip`)

| Value | Augmentation applied                                 |
|-------|------------------------------------------------------|
| `0`   | None                                                 |
| `1`   | Horizontal flip only                                 |
| `2`   | Vertical flip only                                   |
| `3`   | Both horizontal and vertical flips (**recommended**) |

Cells in fluorescence images have no canonical orientation, so flipping both axes is almost always beneficial.


## Class imbalance

By default, cellclass weights the loss function inversely proportional to class frequency. Disable with `--no-loss-weights` if your classes are well balanced.

------------------------------------------------------------------------


# Performance expectations (50 × 50 px, 2-class)

Based on the mitosis-RPE and ppase-screen datasets:

| Model | Typical test accuracy | Training time (per epoch, GPU) |
|----|----|----|
| [squeezenet1_0](../reference/Model.md#cellclass.Model.squeezenet1_0) | 88-92 % | ~5 s |
| [shufflenet2x1_0](../reference/Model.md#cellclass.Model.shufflenet2x1_0) | 92-96 % | ~8 s |
| [densenet121](../reference/Model.md#cellclass.Model.densenet121) | 93-96 % | ~15 s |
| [efficientnetb3s](../reference/Model.md#cellclass.Model.efficientnetb3s) (frozen) | 90-95 % | ~20 s |

Accuracy is highly dataset-dependent. Always evaluate on your own held-out test split.