{
"cells": [
{
"cell_type": "markdown",
"id": "b50523e6",
"metadata": {},
"source": [
"# Train a model\n",
"\n",
"Train a model to find nuclei and identify their cell types. First run\n",
"[data preparation](prepare-data.ipynb) to create the images and labels used here.\n",
"\n",
"This example runs five epochs. Inspect the validation predictions to decide\n",
"whether to train longer.\n"
]
},
{
"cell_type": "markdown",
"id": "ce252923",
"metadata": {},
"source": [
"## Setup"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "cd8abd07",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Device: cuda\n"
]
}
],
"source": [
"import json\n",
"import random\n",
"from pathlib import Path\n",
"\n",
"import albumentations as A\n",
"import matplotlib.pyplot as plt\n",
"import numpy as np\n",
"import pandas as pd\n",
"import torch\n",
"from matplotlib.patches import Patch\n",
"from safetensors.torch import load_model, save_model\n",
"from skimage.color import label2rgb\n",
"from torch.optim.lr_scheduler import CosineAnnealingLR\n",
"from torch.utils.data import DataLoader\n",
"\n",
"from cellseg_models_pytorch.inference.post_processor import PostProcessor\n",
"from cellseg_models_pytorch.losses import (\n",
" CELoss,\n",
" DiceLoss,\n",
" JointLoss,\n",
" MSE,\n",
" MultiTaskLoss,\n",
")\n",
"from cellseg_models_pytorch.metrics import panoptic_quality\n",
"from cellseg_models_pytorch.models.cellpose import CellPose\n",
"from cellseg_models_pytorch.transforms.albu_transforms import (\n",
" AlbuStrongAugment,\n",
" ApplyEach,\n",
" CellposeTransform,\n",
" MinMaxNormalization,\n",
")\n",
"from cellseg_models_pytorch.transforms.functional.normalization import minmax_normalize\n",
"from cellseg_models_pytorch.torch_datasets import TrainDatasetFolder\n",
"from cellseg_models_pytorch.utils import FileHandler\n",
"\n",
"seed = 42\n",
"\n",
"# Seed the libraries that shuffle data and choose random augmentations.\n",
"random.seed(seed)\n",
"np.random.seed(seed)\n",
"torch.manual_seed(seed)\n",
"torch.set_num_threads(2)\n",
"\n",
"# Use the GPU when one is available.\n",
"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
"\n",
"data_root = Path(\"../../data/panoptils-nuclei\")\n",
"class_names = [\n",
" \"background\",\n",
" \"neoplastic\",\n",
" \"stromal\",\n",
" \"inflammatory\",\n",
" \"epithelial\",\n",
" \"other\",\n",
" \"unknown\",\n",
"]\n",
"\n",
"print(f\"Device: {device}\")\n"
]
},
{
"cell_type": "markdown",
"id": "5ff80317",
"metadata": {},
"source": [
"## Load images and labels\n",
"\n",
"Pair each image with its nucleus and cell-type labels. Scale image pixels to 0–1.\n",
"For training, flip images and labels together and vary the image colours.\n",
"Validation images are scaled without random changes.\n"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "51d2714a",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Patches: 17536 training, 4048 validation\n"
]
}
],
"source": [
"# Apply the same spatial changes to images and masks.\n",
"train_transforms = A.Compose(\n",
" [\n",
" A.HorizontalFlip(p=0.33),\n",
" A.VerticalFlip(p=0.33),\n",
" AlbuStrongAugment(seed=seed),\n",
" MinMaxNormalization(copy=True),\n",
" ],\n",
" seed=seed,\n",
")\n",
"\n",
"# Keep validation deterministic, with the same image scaling.\n",
"validation_transforms = A.Compose([MinMaxNormalization(copy=True)])\n",
"\n",
"# Build flow targets from the augmented nucleus masks.\n",
"instance_transforms = ApplyEach([CellposeTransform(deduplicate=False)], as_list=True)\n",
"\n",
"train_dataset = TrainDatasetFolder(\n",
" str(data_root / \"train/images\"),\n",
" str(data_root / \"train/labels\"),\n",
" mask_keys=(\"inst\", \"type\"),\n",
" transforms=train_transforms,\n",
" inst_transforms=instance_transforms,\n",
")\n",
"validation_dataset = TrainDatasetFolder(\n",
" str(data_root / \"validation/images\"),\n",
" str(data_root / \"validation/labels\"),\n",
" mask_keys=(\"inst\", \"type\"),\n",
" transforms=validation_transforms,\n",
" inst_transforms=instance_transforms,\n",
")\n",
"\n",
"# Shuffle training patches; keep validation patches in a fixed order.\n",
"train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, num_workers=4)\n",
"validation_loader = DataLoader(\n",
" validation_dataset,\n",
" batch_size=8,\n",
" shuffle=False,\n",
" num_workers=4,\n",
")\n",
"\n",
"print(f\"Patches: {len(train_dataset)} training, {len(validation_dataset)} validation\")\n"
]
},
{
"cell_type": "markdown",
"id": "dcbfe6e4",
"metadata": {},
"source": [
"## Training recipe\n",
"\n",
"Follow the [Histolytics training recipe](https://hautaniemilab.github.io/histolytics/user_guide/seg/finetuning/)\n",
"with two outputs: cell types and flows. Flows point nucleus pixels toward their\n",
"centre and help separate touching nuclei.\n",
"\n",
"Use cross-entropy and Dice loss for cell types, and mean squared error for flows.\n",
"Enable `apply_sd=True` to add a penalty on large class scores (spectral decoupling).\n",
"Train with AdamW and a cosine learning-rate schedule.\n"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "abb51087",
"metadata": {},
"outputs": [],
"source": [
"criterion = MultiTaskLoss(\n",
" {\n",
" \"cellpose\": JointLoss([MSE()]),\n",
" \"type_mask\": JointLoss([CELoss(apply_sd=True), DiceLoss()]),\n",
" },\n",
" loss_weights={\"cellpose\": 1.0, \"type_mask\": 1.0},\n",
").to(device)\n",
"\n",
"learning_rate = 1e-4\n",
"num_epochs = 5\n"
]
},
{
"cell_type": "markdown",
"id": "99ae01f3",
"metadata": {},
"source": [
"## Initialize the network\n",
"\n",
"Use an EfficientNet-B5 encoder pretrained on ImageNet and randomly initialized\n",
"segmentation layers.\n",
"The encoder weights download on first use.\n",
"\n",
"Set `checkpoint_dir` to a new folder for the saved model files.\n"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "4094ee96",
"metadata": {},
"outputs": [],
"source": [
"model = CellPose(\n",
" 7,\n",
" enc_name=\"efficientnet_b5\",\n",
" enc_pretrain=True,\n",
" device=device,\n",
").model\n",
"\n",
"optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n",
"scheduler = CosineAnnealingLR(optimizer, T_max=num_epochs, eta_min=0)\n",
"\n",
"# A new directory protects checkpoints from earlier runs.\n",
"checkpoint_dir = Path(\"training-checkpoints\")\n",
"checkpoint_dir.mkdir(exist_ok=False)\n"
]
},
{
"cell_type": "markdown",
"id": "optional-freeze-encoder",
"metadata": {},
"source": [
"## Optional: freeze the encoder\n",
"\n",
"To train only the segmentation layers, freeze the encoder before training:\n",
"\n",
"```python\n",
"encoder = model[model.enc_name]\n",
"for parameter in encoder.parameters():\n",
" parameter.requires_grad_(False)\n",
"```\n",
"\n",
"Add this immediately after `model.train()` in the epoch loop to keep the encoder’s\n",
"BatchNorm statistics fixed too:\n",
"\n",
"```python\n",
"encoder.eval()\n",
"```\n",
"\n",
"Skip these steps to train the whole model. Reapply them if you resume with a frozen encoder.\n"
]
},
{
"cell_type": "markdown",
"id": "717dbf3f",
"metadata": {},
"source": [
"## Define the training steps\n",
"\n",
"Use the same flow and cell-type targets for training and validation.\n"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "ab2810d2",
"metadata": {},
"outputs": [],
"source": [
"def prepare_batch(batch: dict[str, torch.Tensor]):\n",
" # Class labels use integers; flow targets use floating-point values.\n",
" images = batch[\"image\"].to(device)\n",
" targets = {\n",
" \"type_mask\": batch[\"type\"].long().to(device),\n",
" \"cellpose\": batch[\"cellpose_inst\"].float().to(device),\n",
" }\n",
" return images, targets\n",
"\n",
"\n",
"def train_epoch(model, loader, optimizer):\n",
" total = 0.0\n",
" for batch in loader:\n",
" images, targets = prepare_batch(batch)\n",
"\n",
" # Clear old gradients before predicting this batch.\n",
" optimizer.zero_grad(set_to_none=True)\n",
" output = model(images)[\"nuc\"]\n",
" loss = criterion(\n",
" {\"type_mask\": output.type_map, \"cellpose\": output.aux_map}, targets\n",
" )\n",
"\n",
" # Backpropagate the loss, then update the model weights.\n",
" loss.backward()\n",
" optimizer.step()\n",
"\n",
" # Weight each batch by its size so the last, smaller batch counts correctly.\n",
" total += float(loss.detach()) * images.shape[0]\n",
"\n",
" return total / len(loader.dataset)\n",
"\n",
"\n",
"def validate_epoch(model, loader):\n",
" total = 0.0\n",
" with torch.inference_mode():\n",
" for batch in loader:\n",
" images, targets = prepare_batch(batch)\n",
"\n",
" output = model(images)[\"nuc\"]\n",
" loss = criterion(\n",
" {\"type_mask\": output.type_map, \"cellpose\": output.aux_map}, targets\n",
" )\n",
"\n",
" total += float(loss) * images.shape[0]\n",
"\n",
" return total / len(loader.dataset)\n"
]
},
{
"cell_type": "markdown",
"id": "training-save-checkpoint",
"metadata": {},
"source": [
"## Save training progress\n",
"\n",
"Save model weights in `model.safetensors` and the optimizer, schedule and random\n",
"states in `training-state.pt`. Each epoch gets its own folder.\n"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "training-checkpoint-helper",
"metadata": {},
"outputs": [],
"source": [
"def save_checkpoint(epoch):\n",
" numpy_rng = np.random.get_state()\n",
"\n",
" # Creating a new epoch folder protects both files from being overwritten.\n",
" epoch_dir = checkpoint_dir / f\"epoch_{epoch:02d}\"\n",
" epoch_dir.mkdir()\n",
" save_model(\n",
" model,\n",
" epoch_dir / \"model.safetensors\",\n",
" metadata={\n",
" \"class_names\": json.dumps(class_names),\n",
" \"normalization\": \"global min-max, RGB float32\",\n",
" },\n",
" )\n",
"\n",
" path = epoch_dir / \"training-state.pt\"\n",
" with path.open(\"xb\") as destination:\n",
" torch.save(\n",
" {\n",
" \"optimizer\": optimizer.state_dict(),\n",
" \"scheduler\": scheduler.state_dict(),\n",
" \"epoch\": epoch,\n",
" \"encoder\": model.enc_name,\n",
" \"class_names\": class_names,\n",
" \"torch_rng\": torch.get_rng_state(),\n",
" \"cuda_rng\": torch.cuda.get_rng_state(device)\n",
" if device.type == \"cuda\"\n",
" else None,\n",
" \"python_rng\": random.getstate(),\n",
" \"numpy_rng\": (numpy_rng[0], numpy_rng[1].tolist(), *numpy_rng[2:]),\n",
" \"data_manifest\": json.loads((data_root / \"manifest.json\").read_text()),\n",
" \"normalization\": \"global min-max, RGB float32\",\n",
" \"augmentation_seed\": seed,\n",
" \"num_workers\": train_loader.num_workers,\n",
" \"num_epochs\": num_epochs,\n",
" \"training_loss\": training_loss,\n",
" \"validation_loss\": validation_loss,\n",
" },\n",
" destination,\n",
" )\n",
"\n",
" return path\n"
]
},
{
"cell_type": "markdown",
"id": "b09d5523",
"metadata": {},
"source": [
"## Train and validate\n"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "564e694d",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 1: training loss 1.172, validation loss 0.820\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 2: training loss 0.964, validation loss 0.767\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 3: training loss 0.909, validation loss 0.752\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 4: training loss 0.868, validation loss 0.735\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 5: training loss 0.841, validation loss 0.828\n"
]
}
],
"source": [
"for epoch in range(num_epochs):\n",
" # Use new random augmentations on each pass through the training set.\n",
" train_transforms.set_random_seed(seed + epoch)\n",
" model.train()\n",
" training_loss = train_epoch(model, train_loader, optimizer)\n",
"\n",
" # Validation changes neither weights nor BatchNorm statistics.\n",
" model.eval()\n",
" validation_loss = validate_epoch(model, validation_loader)\n",
"\n",
" # Advance the learning-rate schedule and save the completed epoch.\n",
" scheduler.step()\n",
" saved_checkpoint = save_checkpoint(epoch + 1)\n",
"\n",
" print(\n",
" f\"Epoch {epoch + 1}: training loss {training_loss:.3f}, \"\n",
" f\"validation loss {validation_loss:.3f}\"\n",
" )\n"
]
},
{
"cell_type": "markdown",
"id": "2157f095",
"metadata": {},
"source": [
"## Load the saved model\n",
"\n",
"Load the last epoch’s safetensors weights for prediction. Restore the separate\n",
"training state to continue training from that epoch.\n"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "9c41c9ed",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Reloaded model, optimizer and scheduler after epoch 5\n"
]
}
],
"source": [
"state = torch.load(saved_checkpoint, map_location=\"cpu\", weights_only=True)\n",
"inference_model = CellPose(\n",
" len(state[\"class_names\"]),\n",
" enc_name=state[\"encoder\"],\n",
" enc_pretrain=False,\n",
" device=device,\n",
")\n",
"\n",
"restored = inference_model.model\n",
"load_model(\n",
" restored, saved_checkpoint.with_name(\"model.safetensors\"), device=str(device)\n",
")\n",
"\n",
"# Restore optimizer and scheduler state to keep the training schedule.\n",
"resumed_optimizer = torch.optim.AdamW(restored.parameters(), lr=learning_rate)\n",
"resumed_scheduler = CosineAnnealingLR(\n",
" resumed_optimizer, T_max=state[\"scheduler\"][\"T_max\"]\n",
")\n",
"resumed_optimizer.load_state_dict(state[\"optimizer\"])\n",
"resumed_scheduler.load_state_dict(state[\"scheduler\"])\n",
"\n",
"\n",
"# Restore random states so a resumed run continues from this epoch.\n",
"random.setstate(state[\"python_rng\"])\n",
"np.random.set_state(state[\"numpy_rng\"])\n",
"torch.set_rng_state(state[\"torch_rng\"])\n",
"\n",
"if device.type == \"cuda\" and state[\"cuda_rng\"] is not None:\n",
" torch.cuda.set_rng_state(state[\"cuda_rng\"], device)\n",
"\n",
"\n",
"# Switch to prediction mode before evaluating the reloaded weights.\n",
"inference_model.set_inference_mode(mixed_precision=False)\n",
"inference_model.post_processor = PostProcessor(\n",
" \"cellpose\", {\"use_gpu\": device.type == \"cuda\"}\n",
")\n",
"\n",
"print(\"Reloaded model, optimizer and scheduler after epoch\", state[\"epoch\"])\n"
]
},
{
"cell_type": "markdown",
"id": "7ac91c28",
"metadata": {},
"source": [
"## Check predictions on validation images\n",
"\n",
"Predict nuclei in 16 validation images at their original 1024-pixel size. These\n",
"images come from patients that were not used for training.\n",
"\n",
"Panoptic quality (PQ) measures how well predicted nuclei match the labels:\n",
"1 is perfect and 0 means no matches. The table shows the average score and nucleus\n",
"counts. Cell-type scores require both the correct type and a matching nucleus;\n",
"they average only images where that type appears in the labels or predictions.\n",
"`NaN` means there were no cells of that type to score. These results describe the\n",
"16 example images; inspect the masks too, especially for rare types.\n"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "a42331d5",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Held-out regions: 16\n"
]
},
{
"data": {
"text/html": [
"\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" PQ | \n",
" target objects | \n",
" predicted objects | \n",
"
\n",
" \n",
" | class | \n",
" | \n",
" | \n",
" | \n",
"
\n",
" \n",
" \n",
" \n",
" | all nuclei | \n",
" 0.643 | \n",
" 6756 | \n",
" 7192 | \n",
"
\n",
" \n",
" | neoplastic | \n",
" 0.509 | \n",
" 3082 | \n",
" 2860 | \n",
"
\n",
" \n",
" | stromal | \n",
" 0.361 | \n",
" 1740 | \n",
" 1986 | \n",
"
\n",
" \n",
" | inflammatory | \n",
" 0.308 | \n",
" 1202 | \n",
" 1115 | \n",
"
\n",
" \n",
" | epithelial | \n",
" NaN | \n",
" 0 | \n",
" 0 | \n",
"
\n",
" \n",
" | other | \n",
" 0.000 | \n",
" 3 | \n",
" 0 | \n",
"
\n",
" \n",
" | unknown | \n",
" 0.219 | \n",
" 657 | \n",
" 1231 | \n",
"
\n",
" \n",
"
\n",
"