3.1 Siamese Neural Network¶
Note
The number of epochs and the real dataset size actually used for training and validation is reduced to speed up the CI pipeline. To improve th performance of the model, feel free to increase these:
CONFIG_DATALOADER["train_pairs_per_epoch"]andCONFIG_DATALOADER["val_pairs"]: number of training and validation pairs actually used per epoch.CONFIG_TRAINER["max_epochs"]: number of epochs
This notebook demonstrates how to train one Siamese Neural Network for image similarity computation. For this, the Oxford Flowers dataset is used, but customized for the Siamese network training.
Import libraries¶
from typing import Any
from collections import defaultdict
import multiprocessing
import random
import datetime
import os
import torch
from torch.utils.data import DataLoader, Dataset, RandomSampler, SequentialSampler
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms
from pyvisim.neural_networks import ContrastiveSiameseNetwork, BCESiameseNetwork
from pyvisim.datasets import OxfordFlowerDataset
/home/runner/work/Python-Visual-Similarity/Python-Visual-Similarity/.venv/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
from .autonotebook import tqdm as notebook_tqdm
Declare the dataset¶
The Siamese Network needs pairs of images during training. Each item in the dataset is a pair of images (img_a, img_b) and a label indicating whether they are similar (interger 1) or dissimilar (integer 0).
The pairs are also built upfront to avoid repeated sampling during training. You can change the seed to get different pairs.
positive_fraction controls the ratio of positive to negative pairs generated during training. For example, if positive_fraction=0.5, then half of the pairs will be positive and half will be negative. The default ratio of 0.5 assumes a balanced dataset. For skewed datasets, adjust this parameter.
class OxfordSiamesePairs(Dataset[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]):
def __init__(
self,
purpose: str,
transform: transforms.Compose,
positive_fraction: float = 0.5,
) -> None:
self._base = OxfordFlowerDataset(purpose=purpose)
self.transform = transform
self.positive_fraction = positive_fraction
# Build a label -> list[index] map for fast pair mining. Labels are
# read directly; indexing the base dataset would decode every image.
self._label_to_indices: dict[int, list[int]] = defaultdict(list)
for idx, label in enumerate(self._base.labels):
self._label_to_indices[label].append(idx)
self._labels_list = list(self._label_to_indices.keys())
if len(self._labels_list) < 2:
raise ValueError("Need >= 2 classes for contrastive training.")
# For reproducibility
rng = random.Random(SEED)
# Build the pairs list upfront to avoid repeated sampling during training,
# which has O(n) complexity.
self._pairs: list[tuple[int, int]] = [
self._sample_partner(idx, rng) for idx in range(len(self._base))
]
def __len__(self) -> int:
return len(self._base)
def _get_tensor(self, idx: int) -> torch.Tensor:
image, _, _ = self._base[idx]
tensor: torch.Tensor = self.transform(image)
return tensor
def _sample_partner(self, idx: int, rng) -> tuple[int, int]:
label_a = self._base.labels[idx]
if rng.random() < self.positive_fraction:
# Positive: another image of the same class.
same_indices = [i for i in self._label_to_indices[label_a] if i != idx]
idx_b = rng.choice(same_indices) if same_indices else idx
return idx_b, 1
# Negative: an image from a different class.
neg_label = rng.choice([lbl for lbl in self._labels_list if lbl != label_a])
idx_b = rng.choice(self._label_to_indices[neg_label])
return idx_b, 0
def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
img_a = self._get_tensor(idx)
idx_b, pair_label = self._pairs[idx]
img_b = self._get_tensor(idx_b)
return img_a, img_b, torch.tensor(pair_label, dtype=torch.float32)
Training seeds¶
Training on every pair of the training split for many epochs takes hours without a GPU. train_pairs_per_epoch draws that many random pairs in every epoch and val_pairs validates on the first pairs of the validation split, so that this notebook also finishes on a CPU. The splits hold one pair per image, so raise both up to the size of the splits to train on every pair.
SEED = 42
random.seed(SEED)
torch.manual_seed(SEED)
device = "cuda" if torch.cuda.is_available() else "cpu"
CONFIG_DATALOADER: dict[str, Any] = {
"batch_size": 64,
"num_workers": 4,
"train_pairs_per_epoch": 640,
"val_pairs": 160
}
def seed_worker(worker_id: int) -> None:
random.seed(torch.initial_seed() % 2**32)
Transforms¶
The transforms below match the ones used in the pre-trained backbone (ResNet-18, in the default configuration, which is also the configuration below).
_NORMALIZE = transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
)
TRAIN_TF = transforms.Compose(
[
transforms.ToPILImage(),
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.RandomApply(
[
transforms.RandomChoice([
transforms.RandomHorizontalFlip(p=1.0),
transforms.RandomRotation(15),
]),
],
p=0.5,
),
transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.1),
transforms.ToTensor(),
_NORMALIZE,
]
)
VAL_TF = transforms.Compose(
[
transforms.ToPILImage(),
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
_NORMALIZE,
]
)
Multiprocessing configs¶
This was necessary on Linux and Python 3.14+. You can comment it out if it caused you trouble.
# Python 3.14 made "forkserver" the default start method on Linux. A forkserver
# worker re-imports the payload in a fresh interpreter, so it cannot resolve
# classes defined in a notebook (they live in ``__main__``). Forking inherits the
# interpreter state instead, which keeps ``OxfordSiamesePairs`` usable.
MP_CONTEXT = "fork" if "fork" in multiprocessing.get_all_start_methods() else None
def make_loader(
dataset: Dataset[tuple[torch.Tensor, torch.Tensor, torch.Tensor]],
num_samples: int,
shuffle: bool,
) -> DataLoader[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
num_workers = CONFIG_DATALOADER["num_workers"] if MP_CONTEXT is not None else 0
# A shuffled loader draws new random items in every epoch, the other one
# always reads the first items of the dataset.
sampler = (
RandomSampler(dataset, num_samples=num_samples)
if shuffle
else SequentialSampler(range(min(num_samples, len(dataset))))
)
return DataLoader(
dataset,
batch_size=CONFIG_DATALOADER["batch_size"],
sampler=sampler,
num_workers=num_workers,
pin_memory=True,
worker_init_fn=seed_worker,
multiprocessing_context=MP_CONTEXT if num_workers > 0 else None,
persistent_workers=num_workers > 0,
)
Part 1: The contrastive Siamese network¶
Introduced by LeCun et al. in 2005, this was the classical metric-learning approach for image similarity. The network is trained to minimize the distance between embeddings of similar images and maximize the distance between embeddings of dissimilar images.
First, it encodes the input images into embeddings using a backbone network (ResNet-18 in this case). Then, it computes the distance between the embeddings. The ContrastiveLoss is used as criterion.
ContrastiveLoss is defined as follows:
\[
L = \frac{1}{2N} \sum_{n=1}^{N} \Big[ y_n d_n^2 + (1 - y_n) \max(0, m - d_n)^2 \Big]
\]
Whereas:
N: number of pairs in the batch
y_n: label of the pair (1 for similar, 0 for dissimilar)
d_n: distance between the embeddings of the pair
m: margin: controls how far apart the embeddings of dissimilar pairs should be
from pyvisim.neural_networks.losses import ContrastiveLoss
Train parameters¶
CONFIG_CONTRASTIVE: dict[str, Any] = {
"embedding_dim": 64,
"margin": 1.0,
"lr": 1e-4,
"weight_decay": 1e-4,
"num_epochs": 5,
"checkpoint_dir": "checkpoints",
"log_every_n": 50,
"gamma": 0.8,
}
Some train helpers¶
def _run_epoch_contrastive(
model: ContrastiveSiameseNetwork,
loader: DataLoader[tuple[torch.Tensor, torch.Tensor, torch.Tensor]],
criterion: torch.nn.Module,
optimizer: torch.optim.Optimizer,
is_train: bool,
epoch: int,
writer: SummaryWriter
) -> float:
model.train() if is_train else model.eval()
total_loss = 0.0
ctx = torch.enable_grad() if is_train else torch.no_grad()
with ctx:
for step, (img_a, img_b, labels) in enumerate(loader):
img_a = img_a.to(device)
img_b = img_b.to(device)
labels = labels.to(device)
emb_a = model(img_a)
emb_b = model(img_b)
loss = criterion(emb_a, emb_b, labels)
if is_train:
optimizer.zero_grad()
loss.backward()
# Gradient clipping to prevent exploding gradients.
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
total_loss += loss.item()
if is_train and (step + 1) % CONFIG_CONTRASTIVE["log_every_n"] == 0:
print(
f" Epoch {epoch:03d} | step {step + 1:04d} "
f"| loss {total_loss / (step + 1):.4f}"
)
loss_str = "Loss/Train" if is_train else "Loss/Validation"
global_step = (epoch - 1) * len(loader) + step
writer.add_scalar(loss_str, loss.item(), global_step=global_step)
return total_loss / len(loader)
def train_contrastive(model: ContrastiveSiameseNetwork, criterion: torch.nn.Module) -> None:
session_id = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
writer = SummaryWriter(log_dir=f"runs/siamese/session_{session_id}_contrastiveloss")
try:
os.makedirs(CONFIG_CONTRASTIVE["checkpoint_dir"], exist_ok=True)
train_ds = OxfordSiamesePairs("train", transform=TRAIN_TF)
val_ds = OxfordSiamesePairs("validation", transform=VAL_TF)
train_loader = make_loader(train_ds, CONFIG_DATALOADER["train_pairs_per_epoch"], shuffle=True)
val_loader = make_loader(val_ds, CONFIG_DATALOADER["val_pairs"], shuffle=False)
print(
f"Train pairs per epoch: {len(train_loader.sampler):,} of {len(train_ds):,} "
f"| Val pairs: {len(val_loader.sampler):,}"
)
print(f"Flower classes: {len(train_ds._labels_list)}")
optimizer = torch.optim.AdamW(
model.parameters(), lr=CONFIG_CONTRASTIVE["lr"], weight_decay=CONFIG_CONTRASTIVE["weight_decay"]
)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=4, gamma=CONFIG_CONTRASTIVE.get('gamma', 0.8))
best_val = float("inf")
for epoch in range(1, CONFIG_CONTRASTIVE["num_epochs"] + 1):
train_loss = _run_epoch_contrastive(
model, train_loader, criterion, optimizer, is_train=True, epoch=epoch, writer=writer
)
val_loss = _run_epoch_contrastive(
model, val_loader, criterion, optimizer, is_train=False, epoch=epoch, writer=writer
)
scheduler.step()
print(
f"Epoch {epoch:03d} | train {train_loss:.4f} | val {val_loss:.4f} "
f"| lr {scheduler.get_last_lr()[0]:.2e}"
)
model.save_to_disk("contrastive_siamese.safetensors")
if val_loss < best_val:
best_val = val_loss
model.save_to_disk("contrastive_siamese_best.safetensors")
print(f"New best val loss: {best_val:.4f}")
print(f"\nTraining complete. Best val loss: {best_val:.4f}")
finally:
writer.close()
criterion = ContrastiveLoss(margin=CONFIG_CONTRASTIVE["margin"])
model_contrastive = ContrastiveSiameseNetwork(
backbone="resnet18",
embedding_dim=CONFIG_CONTRASTIVE["embedding_dim"],
device=device,
pretrained_backbone=True,
)
train_contrastive(model_contrastive, criterion)
Downloading: "https://download.pytorch.org/models/resnet18-f37072fd.pth" to /home/runner/work/Python-Visual-Similarity/Python-Visual-Similarity/.cache/torch/hub/checkpoints/resnet18-f37072fd.pth
0%| | 0.00/44.7M [00:00<?, ?B/s]
14%|█▍ | 6.38M/44.7M [00:00<00:00, 66.8MB/s]
78%|███████▊ | 35.0M/44.7M [00:00<00:00, 204MB/s]
100%|██████████| 44.7M/44.7M [00:00<00:00, 178MB/s]
Train pairs per epoch: 640 of 6,149 | Val pairs: 160
Flower classes: 102
/home/runner/work/Python-Visual-Similarity/Python-Visual-Similarity/.venv/lib/python3.10/site-packages/torch/utils/data/dataloader.py:1102: UserWarning: 'pin_memory' argument is set as true but no accelerator is found, then device pinned memory won't be used.
super().__init__(loader)
Epoch 001 | train 0.1315 | val 0.0973 | lr 1.00e-04
New best val loss: 0.0973
Epoch 002 | train 0.0942 | val 0.0820 | lr 1.00e-04
New best val loss: 0.0820
Epoch 003 | train 0.0802 | val 0.0699 | lr 1.00e-04
New best val loss: 0.0699
Epoch 004 | train 0.0726 | val 0.0632 | lr 8.00e-05
New best val loss: 0.0632
Epoch 005 | train 0.0634 | val 0.0610 | lr 8.00e-05
New best val loss: 0.0610
Training complete. Best val loss: 0.0610
Load models after training¶
model_contrastive = ContrastiveSiameseNetwork.load_from_disk("contrastive_siamese_best.safetensors")
model_contrastive.eval()
ContrastiveSiameseNetwork(similarity_func=cosine)
Visualize results¶
Now, we will pick three random images from the validation set, two of which are similar and one is dissimilar.
import matplotlib.pyplot as plt
import numpy as np
def visualize_image_pair(image_a: np.ndarray, image_b: np.ndarray) -> None:
fig, axes = plt.subplots(1, 2, figsize=(8, 4))
axes[0].imshow(image_a)
axes[0].set_title("Image A")
axes[0].axis("off")
axes[1].imshow(image_b)
axes[1].set_title("Image B")
axes[1].axis("off")
plt.tight_layout()
plt.show()
test_dataset = OxfordFlowerDataset(purpose="test")
img_a = test_dataset[4][0]
img_b = test_dataset[5][0]
img_c = test_dataset[100][0]
print("Image shapes:\n")
print(f"Image A: {img_a.shape}")
print(f"Image B: {img_b.shape}")
visualize_image_pair(img_a, img_b)
visualize_image_pair(img_a, img_c)
Image shapes:
Image A: (500, 753, 3)
Image B: (500, 748, 3)
Compute similarity¶
As observed, the pair of the same class yields a similarity score closer to 1, while the pair of different classes yields a similarity score should be much lower.
score_same_class = model_contrastive.similarity_score(img_a, img_b).item()
score_different_class = model_contrastive.similarity_score(img_a, img_c).item()
print(f"Similarity score (same class): {score_same_class:.4f}")
print(f"Similarity score (different class): {score_different_class:.4f}")
Similarity score (same class): 0.9336
Similarity score (different class): 0.4944
Part 2: BCELoss Siamese network¶
This is a more recent approach to train a Siamese network. Instead of learning the metric, the network learns to predict the similarity score directly using L1 distance between the embeddings and a sigmoid activation function. The network is trained using BCELoss, as familiar in binary classification tasks.
NOTE¶
This network, unlike the contrastive network, does not learn to produce embeddings. Instead, they learn to model the similarity score directly.
Import the loss function¶
from torch.nn import BCEWithLogitsLoss
Train parameters¶
CONFIG_BCE: dict[str, Any] = {
"embedding_dim": 128,
"margin": 1.0,
"lr": 1e-4,
"weight_decay": 1e-4,
"num_epochs": 5,
"checkpoint_dir": "checkpoints",
"log_every_n": 50,
"gamma": 0.8,
}
Train helpers¶
def _run_epoch_bce(
model: BCESiameseNetwork,
loader: DataLoader[tuple[torch.Tensor, torch.Tensor, torch.Tensor]],
criterion: torch.nn.BCEWithLogitsLoss,
optimizer: torch.optim.Optimizer,
is_train: bool,
epoch: int,
) -> float:
model.train() if is_train else model.eval()
total_loss = 0.0
ctx = torch.enable_grad() if is_train else torch.no_grad()
with ctx:
for step, (img_a, img_b, labels) in enumerate(loader):
img_a = img_a.to(device)
img_b = img_b.to(device)
labels = labels.to(device)
logits = model(img_a, img_b)
loss = criterion(logits, labels)
if is_train:
optimizer.zero_grad()
loss.backward()
# Gradient clipping to prevent exploding gradients.
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
total_loss += loss.item()
if is_train and (step + 1) % CONFIG_BCE["log_every_n"] == 0:
print(
f" Epoch {epoch:03d} | step {step + 1:04d} "
f"| loss {total_loss / (step + 1):.4f}"
)
return total_loss / len(loader)
def train_bce(model: BCESiameseNetwork, criterion: torch.nn.BCEWithLogitsLoss) -> None:
session_id = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
writer = SummaryWriter(log_dir=f"runs/siamese/session_{session_id}_bceloss")
try:
os.makedirs(CONFIG_BCE["checkpoint_dir"], exist_ok=True)
train_ds = OxfordSiamesePairs("train", transform=TRAIN_TF)
val_ds = OxfordSiamesePairs("validation", transform=VAL_TF)
train_loader = make_loader(train_ds, CONFIG_DATALOADER["train_pairs_per_epoch"], shuffle=True)
val_loader = make_loader(val_ds, CONFIG_DATALOADER["val_pairs"], shuffle=False)
print(
f"Train pairs per epoch: {len(train_loader.sampler):,} of {len(train_ds):,} "
f"| Val pairs: {len(val_loader.sampler):,}"
)
print(f"Flower classes: {len(train_ds._labels_list)}")
model = BCESiameseNetwork(
backbone="resnet18",
embedding_dim=CONFIG_BCE["embedding_dim"],
device=device,
pretrained_backbone=True,
)
criterion = torch.nn.BCEWithLogitsLoss()
optimizer = torch.optim.AdamW(
model.parameters(), lr=CONFIG_BCE["lr"], weight_decay=CONFIG_BCE["weight_decay"]
)
scheduler = torch.optim.lr_scheduler.StepLR(
optimizer,
step_size=4,
gamma=CONFIG_BCE.get('gamma', 0.8)
)
best_val = float("inf")
for epoch in range(1, CONFIG_BCE["num_epochs"] + 1):
train_loss = _run_epoch_bce(
model, train_loader, criterion, optimizer, is_train=True, epoch=epoch
)
val_loss = _run_epoch_bce(
model, val_loader, criterion, optimizer, is_train=False, epoch=epoch
)
scheduler.step()
print(
f"Epoch {epoch:03d} | train {train_loss:.4f} | val {val_loss:.4f} "
f"| lr {scheduler.get_last_lr()[0]:.2e}"
)
writer.add_scalar("Loss/Train", train_loss, epoch)
writer.add_scalar("Loss/Validation", val_loss, epoch)
writer.add_scalar("Learning Rate", scheduler.get_last_lr()[0], epoch)
ckpt = {
"epoch": epoch,
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
"scheduler": scheduler.state_dict(),
"val_loss": val_loss,
"cfg": CONFIG_BCE,
}
model.save_to_disk("bce_siamese.safetensors")
if val_loss < best_val:
best_val = val_loss
model.save_to_disk("bce_siamese_best.safetensors")
print(f"New best val loss: {best_val:.4f}")
print(f"\nTraining complete. Best val loss: {best_val:.4f}")
finally:
writer.close()
criterion = BCEWithLogitsLoss()
model_bce = BCESiameseNetwork(
backbone="resnet18",
embedding_dim=CONFIG_BCE["embedding_dim"],
device=device,
pretrained_backbone=True,
)
train_bce(model_bce, criterion)
Train pairs per epoch: 640 of 6,149 | Val pairs: 160
Flower classes: 102
/home/runner/work/Python-Visual-Similarity/Python-Visual-Similarity/.venv/lib/python3.10/site-packages/torch/utils/data/dataloader.py:1102: UserWarning: 'pin_memory' argument is set as true but no accelerator is found, then device pinned memory won't be used.
super().__init__(loader)
Epoch 001 | train 0.6897 | val 0.6853 | lr 1.00e-04
New best val loss: 0.6853
Epoch 002 | train 0.6743 | val 0.6698 | lr 1.00e-04
New best val loss: 0.6698
Epoch 003 | train 0.6519 | val 0.6459 | lr 1.00e-04
New best val loss: 0.6459
Epoch 004 | train 0.6325 | val 0.6235 | lr 8.00e-05
New best val loss: 0.6235
Epoch 005 | train 0.6232 | val 0.6133 | lr 8.00e-05
New best val loss: 0.6133
Training complete. Best val loss: 0.6133
Load models after training¶
model_bce = BCESiameseNetwork.load_from_disk("bce_siamese_best.safetensors")
model_bce.eval()
BCESiameseNetwork(similarity_func=cosine)
Visualize results¶
Now, we will pick three random images from the validation set, two of which are similar and one is dissimilar.
print("Image shapes:\n")
print(f"Image A: {img_a.shape}")
print(f"Image B: {img_b.shape}")
visualize_image_pair(img_a, img_b)
visualize_image_pair(img_a, img_c)
Image shapes:
Image A: (500, 753, 3)
Image B: (500, 748, 3)
Compute similarity¶
Since a sigmoid function is applied on the absolute difference between embeddings, the score for the similar pair should be closer to 1 and the score for the dissimilar pair should be closer to 0.
score_same_class = model_bce.similarity_score(img_a, img_b).item()
score_different_class = model_bce.similarity_score(img_a, img_c).item()
print(f"Similarity score (same class): {score_same_class:.4f}")
print(f"Similarity score (different class): {score_different_class:.4f}")
Similarity score (same class): 0.4121
Similarity score (different class): 0.2298
References¶
[1] Koch, G., Zemel, R., & Salakhutdinov, R. (2015). Siamese Neural Networks for One-shot Image Recognition. ICML Deep Learning Workshop. https://www.cs.cmu.edu/~rsalakhu/papers/oneshot1.pdf
[2] Hadsell, R., Chopra, S., & LeCun, Y. (2006). Dimensionality Reduction by Learning an Invariant Mapping. In Proceedings of the 2006 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR), Vol. 2, 1735-1742. https://doi.org/10.1109/CVPR.2006.100